Causal Masking

Your program will add one rule to the attention you already have: a token may look at itself and the tokens before it, never at tokens after it. This is what turns bidirectional attention into the attention of a GPT.

Why the future must be hidden

A GPT is trained to predict token t+1 from tokens 0..t. If position t could read token t+1 in attention, the model would just copy the answer, the loss would collapse to zero, and nothing would be learned. At generation time the problem is worse: the future tokens do not exist yet.

The mask is also what makes training cheap. Run the model once over a whole sequence of T tokens, and every position t produces a prediction for t+1. Because the mask guarantees that the prediction at t depended only on 0..t, all T predictions are valid at once, which is T training examples for the price of one forward pass. Feeding the true previous tokens instead of the model's own guesses is called teacher forcing.

The mechanism

Compute the scores S[i][j] = Q[i] . K[j] / sqrt(d) as before. Then, before the softmax, replace every score with j > i (key after query) by minus infinity. Softmax turns exp(-inf) into exactly 0, so those keys receive no weight, and the remaining weights in the row are renormalised among themselves.

raw scores S                 after the mask               weights (softmax per row)
[ 2   5   1 ]                [ 2  -inf -inf ]             [ 1.0000  0.0000  0.0000 ]
[ 0   1   4 ]      ->        [ 0   1   -inf ]      ->     [ 0.2689  0.7311  0.0000 ]
[ 3   3   3 ]                [ 3   3    3   ]             [ 0.3333  0.3333  0.3333 ]

The weight matrix is lower triangular: row 0 can only attend to itself, so its weight on token 0 is exactly 1; the last row sees everything. The raw score of 5 in row 0 (which prefers the future token) is simply never considered.

Why minus infinity and not 0

A score of 0 is not "no attention". It is a perfectly ordinary score, and exp(0) = 1 gives that key a weight comparable to every other key. When the real scores in the row are negative, a masked 0 outscores all of them and the future token wins. Only -inf (or a very large negative number such as -1e9) makes exp return 0.

The mask has to go before the softmax. Zeroing weights afterwards leaves a row that no longer sums to 1.

Check it yourself

A masked implementation has a property you can test: change the key and value of token k to anything you like, and the outputs at all positions before k stay the same. If an earlier output moves, the future is leaking.

Not every attention is masked

• Decoder-only models (GPT, Llama, Claude): causal mask.
• Encoder models (BERT) and the encoder half of T5: no mask, every token sees every other.
• Padding masks hide padding tokens so sequences of different length can share a batch. They are separate from the causal mask and are combined with it by masking a score whenever either rule applies.

Common mistakes

• Masking with 0 instead of -inf.
• Masking the diagonal as well (j >= i). Row 0 then has nothing to attend to, every score is -inf, and the softmax divides zero by zero.
• Masking the wrong triangle (j < i), which hides the past and shows the future.
• Applying the mask after the softmax without renormalising.

Your exercise

Starting from the attention of lesson 5 (given), add the mask. The program prints the full weight matrix first and then the attention outputs, so you can see the triangle. CAUSAL 0 turns the mask off, as in an encoder.