Basis Lab

Answers · How models work

How does masking actually make attention causal?

In shortThe mask is applied to the scores before they become weights, so forbidden positions end up with exactly zero weight and the rest are renormalised automatically.

The mask is applied in score space, before the normalisation that turns scores into weights. Forbidden positions are given a very large negative score, so their exponential is zero and the denominator sums only over allowed positions.

Causal attention does not compute the full set of weights and then hide the future entries. It edits the score matrix first. For a query at position ii, every key at position j>ij>i has a large negative number added to its score sijs_{ij}. Softmax then carries out the exclusion by ordinary arithmetic: exponentiate, add, divide.

The mask is usually stored as a matrix of zeros and very negative numbers, or applied as a lower triangular condition. The arithmetic is the same either way: forbidden scores are not deleted from the tensor, they are made so negative that their contribution to the sum is zero. That is why the same softmax function serves an encoder, which masks nothing, and a decoder, which masks the upper triangle.

s=(2, 1, 0),mask position 3\mathbf{s} = (2,\,1,\,0),\qquad \text{mask position }3

Take three scores before masking, s=(2,1,0)s=(2,1,0). Without a mask, e2=7.389e^2=7.389, e1=2.718e^1=2.718 and e0=1e^0=1, so the denominator is 7.389+2.718+1=11.1077.389+2.718+1=11.107. The weights are 7.389/11.107=0.6657.389/11.107=0.665, 2.718/11.107=0.2452.718/11.107=0.245 and 1/11.107=0.0901/11.107=0.090. Now forbid the third position. Replace its score by negative infinity before exponentiating. Its term becomes e−∞=0e^{-\infty}=0, the denominator becomes 7.389+2.718=10.1077.389+2.718=10.107, and the weights are 7.389/10.107=0.7317.389/10.107=0.731, 2.718/10.107=0.2692.718/10.107=0.269 and 00. The forbidden position is not small. It is exactly zero, and the allowed positions renormalise without a second pass.

FIG 1
positionscoreexp(score)weight before maskweight after mask
127.3890.6650.731
212.7180.2450.269
301.0000.0900
ai=esi∑j≤iesj,aj=0 for j>ia_i = \frac{e^{s_i}}{\sum_{j\le i} e^{s_j}},\qquad a_j=0\ \text{for }j>i

The first row shows the extreme case. Token one has one legal target, itself. Its row contains one finite score and the rest are negative infinity. The denominator is a single term, es11e^{s_{11}}, so a11=es11/es11=1a_{11}=e^{s_{11}}/e^{s_{11}}=1. The score is ignored, not because the code special cases position one, but because normalising over a set of size one forces the answer. That is what makes the triangle in a causal mask exact.

Why before the exponential matters

The natural guess is to compute all the weights, zero the future ones, then divide by the new total. That can be made to agree numerically, but it is a different operation. It needs a second normalisation after softmax. The mask that makes attention causal is a score edit inside the one normalisation. Softmax is unchanged by adding a constant to every score, so the masked score only has to sit far enough below the real scores for its exponential to underflow. The normaliser has no separate masking branch to get wrong.

esi+c∑jesj+c=ecesiec∑jesj=esi∑jesj\frac{e^{s_i+c}}{\sum_j e^{s_j+c}} = \frac{e^c e^{s_i}}{e^c \sum_j e^{s_j}} = \frac{e^{s_i}}{\sum_j e^{s_j}}

That identity is also why implementations subtract the row maximum before exponentiating. If the largest allowed score is 22, subtracting 22 gives exponents 0,−1,−20,-1,-2, whose exponentials are 1,0.368,0.1351,0.368,0.135 and whose sum is 1.5031.503. The weights are 1/1.503=0.6651/1.503=0.665, 0.368/1.503=0.2450.368/1.503=0.245, 0.135/1.503=0.0900.135/1.503=0.090, the same as before. The subtraction changes no weight; it only keeps the exponent from overflowing. A finite very negative mask value works for the same reason. True negative infinity gives an exact zero, but multiplying it by a zero gradient can produce a NaN in the backward pass, so code uses a large finite negative number instead.

What the mask does not do

Masking enforces which positions are available. It does not give the head a way to attend to nothing. Softmax always produces a full distribution over the surviving positions, so a head that finds nothing useful must still put its weight somewhere. Trained models often learn to dump spare weight on the first token, the attention sink. If that token is removed, the place where the spare weight was stored is removed too. The causal constraint is exact; the usefulness of what remains is learned.

Worth keeping

  • A causal mask edits the scores before softmax, so forbidden positions have exponential zero and contribute nothing to the denominator.
  • Softmax is shift invariant, so subtracting the row maximum changes no weight and only prevents overflow.
  • Masking chooses which tokens are legal; it does not let a head abstain, because the weights still sum to one over whatever is legal.

Where this goes next

This question is one step of The Mathematics of Attention, a course that works through its subject from the beginning. The page that teaches this part of it is Softmax Over Scores, and it reads the whole way through whether or not you have an account.

What an account changes is the practice: problems standing under every idea, marked with a reason rather than a tick, in a copy of the course that is yours.

Read that pageThe whole course

Other questions people ask about this

  1. What do multiple attention heads actually give you that one big head doesn't?Several heads let a model track several relationships at once, each in its own lower-dimensional subspace, instead of squeezing every relationship into one weighted average that destroys information.
  2. Why do we divide the attention scores by the square root of the key dimension?The square root of the key dimension is the typical size of an unscaled attention score, so dividing by it stops softmax from becoming a hard pick-one.
  3. Why does attention use softmax instead of just dividing each score by the sum of all the scores?Dividing by the sum can give negative weights and fails when scores cancel; exponentiating first makes weights positive and turns score gaps into ratios.
  4. Why does attention need positional encoding if it already looks at the whole sequence at once?Attention treats tokens as an unordered set, so position has to be added to the vectors; otherwise a sentence and its anagram get the same representations.
  5. Why does the model get stuck repeating the same sentence over and over?The model is not stuck in a memory loop: once it has written a phrase, that phrase is in its context and makes writing it again more likely, so a decoder that always picks the top word cannot leave.
  6. Why does diffusion add Gaussian noise instead of any random corruption?Diffusion adds Gaussian noise because Gaussian errors combine exactly: each step keeps the spread fixed and gives a closed-form noisy picture at any later step, which makes the reverse step learnable.

Read alongside