Self-Attention: Compute the Mixture, Then Prove It Is Causal

The MLP reads a fixed window by assigning a fixed input slot to each token. Self-attention instead computes a content-dependent mixture for every position. The same projections are reused at all positions, and the number of visible source positions can vary.

This lesson builds single-head self-attention from four helpers and one module. You should leave able to calculate a row by hand, identify the normalization axis, and test that future information cannot change an earlier output. We use a single sequence; there is no batch axis in this exercise.

Separate comparison from content

For T input vectors of width D, apply three learned Linear(D,D) projections to every vector:

Q[i] = W_q(x[i])   query: what this position is looking for
K[j] = W_k(x[j])   key: how a source position is matched
V[j] = W_v(x[j])   value: the content that source contributes

Q, K, and V each have shape (T,D). The word “value” in V is a vector role; Value with a capital V is our scalar autograd class. Queries and keys determine weights, but values are what get mixed.

For destination i and source j, score their compatibility:

score[i,j] = sum_k Q[i,k]*K[j,k] / sqrt(D)

This is a (T,T) matrix. Each row belongs to one destination and each column to a source. Scaling uses the query/key feature width, not the sequence length. Under an independent, zero-mean, unit-variance component model, the unscaled dot product's variance grows with D, motivating division by sqrt(D). The assumption explains the scale; it is not a claim that every learned component always has that distribution. See scaled dot-product attention.

Establish the visibility rule before normalizing

At position i, the next-token target is the token at i+1. The model may use positions j <= i, including the current input token, but not later ones. For T=3:

mask = [[0, -inf, -inf],
        [0,    0, -inf],
        [0,    0,    0]]

Add this mask to scores before softmax. Negative infinity excludes future positions for any finite unmasked scores. A large finite negative number such as -1e9 is only an approximation; a still larger future score can overcome it. Each row has at least its diagonal available, so our causal mask never produces an all-masked row.

Scores may be finite or negative infinity, with at least one finite score per row; positive infinity and NaN are not valid inputs. An all-masked softmax row is outside the helper's contract: subtracting a maximum of negative infinity yields undefined arithmetic. Keep that distinction in mind when later adding padding masks.

Normalize each destination's row

For a row s:

m = max(s)
e[j] = exp(s[j]-m)
a[j] = e[j] / sum(e)

Use .data only to choose m, then keep the remaining calculation in Value operations. A masked negative-infinity score exponentiates to zero. Every row independently sums to 1. Normalizing the whole score matrix makes unrelated destination positions compete for probability mass; normalizing columns changes the meaning of the mixture.

The attended output is z[i,k] = sum_j a[i,j]*V[j,k], with shape (T,D). Finally apply W_o = Linear(D,D) to each attended vector. This output projection can mix features; the probability interpretation belongs to a, not to the final output.

Work one row numerically

At destination i=1, suppose the already scaled scores are [0, log(3), 99]. Let source values be:

V[0] = [2,4]
V[1] = [6,0]
V[2] = [100,100]

The causal mask gives [0,log(3),-inf]. Subtracting the row maximum gives [-log(3),0,-inf]; exponentials are [1/3,1,0]. Dividing by 4/3 gives weights [1/4,3/4,0].

Therefore the attended vector is:

z[1] = (1/4)*[2,4] + (3/4)*[6,0] + 0*[100,100]
     = [5,1]

If W_o is the identity with zero bias, this is also the final output. The enormous future value and score have no effect.

For row i=0, only source 0 is visible; its weight must be exactly 1. At this row the softmax weights cannot vary, so its Q/K gradient contribution is zero. A zero gradient there is expected and is not evidence that the entire model is disconnected.

Translate the equation into function boundaries

Function	Consumes	Produces
attention_scores	(T,d_k) Q and K	(T,T) scaled scores
causal_mask	T	(T,T) numeric mask
softmax	One score row	One probability row
apply_attention	(T,T) weights, (T,d_v) values	(T,d_v) mixtures
SelfAttention.forward	(T,D) inputs, causal flag	(T,D) projected outputs

Test helpers independently before assembling the projections. Scalar constants are fine; replacing a computed mixture with Value(mixture.data) is not. That replacement produces plausible forward values while losing the derivatives through both weights and source vectors.

Prove causality with a counterfactual

Construct two sequences with an identical prefix and different suffixes. Run both through the same parameters with causal=True. Every output in the shared prefix must match. This checks behavior through the full projection/mask/softmax path, not merely the shape of a mask matrix.

With causal=False, suffix changes are permitted to affect earlier outputs; use a nondegenerate controlled example if testing that they actually do. Also differentiate an early scalar output and verify that future input scalars receive zero gradients.

Checkpoint. Why is “softmax first, set future weights to zero” wrong?

Check your reasoning

the future scores already changed the denominator. For raw probabilities [0.1,0.2,0.7], zeroing the last entry leaves weights summing to 0.3. Masking before normalization would give [1/3,2/3,0] instead.

The attention core costs O(T*T*D) operations and O(T*T) score storage. Dense projections add O(T*D*D) work. Scalar autograd creates many Python objects as well, so keep diagnostic sequences tiny. Next we will project into several heads, then add normalization and residual paths.