LM3-T4 — Causal Masking

Learning objectives

After this lesson, you should be able to:

• Explain why training a next-token predictor needs a causal mask and what happens without it.
• Apply a causal mask to a score matrix before the softmax and compute the resulting weights by hand.
• Explain why masking with 0 or after the softmax is wrong.
• Write a leakage test and combine the causal mask with a padding mask.

The idea: hide the answer

In the previous lessons’ training setup (Lesson LM2-T5) the target for position i is the token at position i + 1. The input sequence contains that token, only one place along. If attention at position i could read position i + 1, the cheapest way to lower the loss would be to copy it. The model would reach a training loss near zero by learning to look ahead, and at generation time, when there is no future to read, it would produce nonsense.

How the mask works

The score matrix S has one row per query and one column per key. Position i may use only the keys at positions up to i, so the entries above the diagonal must not contribute. We set them to −∞ before the softmax. Then exp(−∞) = 0, the masked weights are exactly 0, and the softmax normalises over the remaining entries, so each row still sums to 1.

Why it is done before, and with −∞

Two tempting alternatives are wrong. Setting the masked scores to 0 does not remove them: exp(0) = 1, so each future token receives a positive weight, often a large one. Zeroing the weights after the softmax removes the future but leaves rows that no longer sum to 1; the weights have lost probability mass, and the output is shrunk by a different amount in each row. Only masking before the softmax gives a proper probability distribution over the allowed tokens.

One pass, many predictions

The mask is what allows a whole sequence to be processed at once. Position 3 sees positions 1 to 3 and predicts token 4; position 4 sees 1 to 4 and predicts token 5; all at the same time in the same matrix product. Without the mask we would need to run the model separately on every prefix. The mask makes training n times more efficient than it would otherwise be.

At generation time

When generating, the future does not exist yet, so a mask is not needed to hide it, but the model was trained with the mask, so the computation must stay the same: each position attends to the past only. This is also why the keys and values of past tokens never change as new tokens arrive, which makes the KV cache possible (Module 10).

The mathematics: the masked softmax

Let S be the (n, n) matrix of scaled scores. Define the mask matrix M with

S

scaled scores Q Kᵀ / √d_k, shape (n, n) (or (B, H, n, n) with batch and heads).

M

the additive mask: zero on and below the diagonal, minus infinity above it.

i, j

query position (row) and key position (column).

The masked weights are

so, for the allowed entries j ≤ i, W_ij = exp(S_ij) / Σ_{j′ ≤ i} exp(S_ij′), and W_ij = 0 for j > i. Row 1 has a single allowed entry, so W₁₁ = 1.

Leakage criterion. A layer is causal if, for every position j, changing the input at position j leaves the outputs at positions i < j unchanged. For the masked layer this holds exactly because the output at i depends only on the rows i′ ≤ i of the keys and values.

Implementation detail. Frameworks often use a large negative number such as −10⁹ instead of −∞ to avoid not-a-number values when a row is entirely masked; in half precision −∞ is safer than −10⁹ if values overflow. In the standard causal mask no row is fully masked, because the diagonal is always allowed.

Padding mask. Batches pad short sequences. If keys at padded positions must not receive attention, their columns are masked too: the allowed set for query i is {j : j ≤ i and j is a real token}. Padded query rows may have an empty allowed set; their outputs are ignored, but they must not produce NaN (Lab LM3-L2 Stage 3).

A worked example: three tokens

Reuse the scores from the previous lesson, S ≈ [[0.707, 0, 0.707], [0, 0.707, 0.707], [0.707, 0.707, 1.414]] and the values V = [[1, 0], [0, 2], [3, 1]].

Masked scores. Replace entries above the diagonal by −∞: [[0.707, −∞, −∞], [0, 0.707, −∞], [0.707, 0.707, 1.414]].

Weights. Row 1: only one allowed entry, so [1, 0, 0]. Row 2: exp(0) = 1 and exp(0.707) = 2.028; the sum is 3.028, so [0.330, 0.670, 0]. Row 3 has all three entries and is unchanged: [0.248, 0.248, 0.503].

Outputs. Row 1 = [1, 0] (just its own value). Row 2 = 0.330 · [1, 0] + 0.670 · [0, 2] = [0.330, 1.340]. Row 3 = [1.759, 1.000] as before.

The output of token 1 has changed from a blend of all three values to its own value only, and nothing about the later tokens can influence it. The last token has the full history, so its row is unchanged.

The mask in PyTorch, and a leakage test

import torch

n, d = 3, 2
Q = K = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
V = torch.tensor([[1., 0.], [0., 2.], [3., 1.]])

r3 = lambda t: [[round(v, 3) for v in row] for row in t.tolist()]

def attend(Q, K, V):
    scores = Q @ K.T / d ** 0.5
    future = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)   # True above the diagonal
    scores = scores.masked_fill(future, float("-inf"))                    # mask BEFORE the softmax
    w = torch.softmax(scores, dim=-1)
    return w @ V, w

out, w = attend(Q, K, V)
print(r3(w))            # [[1.0, 0.0, 0.0], [0.33, 0.67, 0.0], [0.248, 0.248, 0.503]]
print(r3(out))          # [[1.0, 0.0], [0.33, 1.34], [1.759, 1.0]]

# leakage test: change the LAST token, earlier outputs must not move
K2, V2 = K.clone(), V.clone()
K2[2], V2[2] = torch.tensor([5., -5.]), torch.tensor([9., 9.])
out2, _ = attend(Q, K2, V2)
print(torch.allclose(out[:2], out2[:2]))       # True: positions 1 and 2 did not notice
print(torch.allclose(out[2], out2[2]))         # False: position 3 did (it reads itself)

The test is worth keeping in any codebase that implements attention. If someone later replaces the layer with a faster version and breaks the mask, the first print becomes False.

Engineering implications and limitations

Always test for leakage. Training loss looks wonderful when the mask is wrong. A model that quickly reaches a loss far below what the data could allow is a warning. The leakage test above (change a later token, compare earlier outputs) is cheap and decisive.

Register the mask as a buffer. In an nn.Module store the triangular boolean matrix with register_buffer so that it moves to the right device with the model but is not a trainable parameter. Slice it to the actual sequence length: mask[:n, :n].

Use is_causal=True in fused kernels. scaled_dot_product_attention(..., is_causal=True) applies the mask internally and never builds the full matrix, which saves memory. Do not pass a mask and also set is_causal.

Be careful combining masks. The causal mask and the padding mask are combined by logical AND of the allowed entries. With left padding, early rows may be entirely masked; handle them explicitly to avoid NaN.

Mask at the right place in the pipeline. The mask affects attention only. Layers that mix positions in other ways (for example convolutions across the sequence) would need their own causal handling. The usual Transformer block mixes positions only through attention.

Bidirectional models are different. Encoder models such as BERT read both directions and are trained to fill in masked tokens, not to predict the next one. The GPT-style models in this track are decoder-only and causal.

Prefix attention. Some instruction-tuned and vision-language models let a prompt prefix attend to itself in both directions and only apply the causal mask to the generated part. That is a different mask, built from the same machinery.

Common misconceptions

Engineering takeaways

1. The causal mask sets the scores of future positions to −∞ before the softmax, so their weights are exactly 0 and each row still sums to 1.
2. It lets a whole sequence be processed in one pass: every position predicts its next token without seeing it.
3. Masking with 0, or zeroing after the softmax, is wrong: the first leaks, the second breaks normalisation.
4. Test for leakage: changing a later token must not change any earlier output.
5. Combine with a padding mask using logical AND, and guard against rows with nothing left to attend to.

References and further reading

See Chapter 3 of the book (hiding future words with causal attention) and the PyTorch documentation for scaled_dot_product_attention.

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762 (masking in the decoder, Section 3.2.3).
• Build a Large Language Model (From Scratch), Chapter 3: Coding Attention Mechanisms (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: multi-head attention, which runs several attention patterns in parallel.

Related content

Nothing below is required to finish this lesson.

• Before
  LM3-T3: Queries, Keys, Values and Scaled Dot-Product Attention

  Compute attention scores, weights and context vectors from queries, keys and values.

• Lab
  LM3-L2: Implement Causal Masking

  Add a causal mask to attention and prove that changing a future token cannot change an earlier output.