Attention masks

A language model predicts the next token, so it must not see the tokens it is trying to predict. A batch also pads short sequences to a common length, and the padding should not affect the result. Both problems are solved with masks: sets of pairs that attention is not allowed to use.

Causal mask: no peeking ahead

Position i may only attend to positions j ≤ i. The blocked pairs form the strict upper triangle of the score matrix. For a sequence of 3 tokens, the mask looks like this, where 1 means blocked:

      j=0  j=1  j=2
i=0    0    1    1
i=1    0    0    1
i=2    0    0    0

Training and generation both rely on this. During training, the model predicts every next token at once. Without the mask, each position would copy the answer from the token after it.

The code below applies the mask, then checks that the weights are zero above the diagonal:

import numpy as np

def causal_attention(q, k, v):
    # q, k, v have shape (n, d) for one head.
    n, d = q.shape
    scores = q @ k.T / np.sqrt(d)
    mask = np.triu(np.ones((n, n), dtype=bool), k=1)   # True above the diagonal
    scores = np.where(mask, -np.inf, scores)            # block the future before softmax
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ v, weights

rng = np.random.default_rng(0)
q, k, v = rng.normal(size=(3, 4, 8))
out, w = causal_attention(q, k, v)
print(np.round(w, 2))
print("upper triangle is zero:", np.all(np.triu(w, k=1) == 0))

Output:

[[1.   0.   0.   0.  ]
 [0.83 0.17 0.   0.  ]
 [0.34 0.41 0.25 0.  ]
 [0.1  0.03 0.68 0.18]]
upper triangle is zero: True

Row 0 can only attend to itself, so its weight is 1. Each later row spreads weight over the positions up to its own.

Padding mask: ignore the filler

Batched sequences have different lengths. Pad the short ones, then forbid attention to the pad positions. Usually you also forbid queries at pad positions from attending anywhere. Combine the padding mask with the causal mask by blocking any pair that either one forbids.

Common mistakes

• Masking after softmax. Softmax has already normalized the scores, so zeroing some entries leaves the rest unfairly renormalized. Always mask scores before softmax.
• Using 0 instead of -∞. exp(0) = 1, so a masked score of 0 still gets weight. Use -∞, or a very large negative number such as -1e9.
• Blocking the diagonal. With diagonal=0 the causal mask hides each token from itself. Use diagonal=1, so position i can still see itself.

Check yourself

1. Which entries of causal_mask(3) are True with diagonal=1? The entries above the diagonal, where j > i: row 0 at columns 1 and 2, and row 1 at column 2.
2. Why set a blocked score to -∞ and not to 0? Because exp(0) = 1 still gives weight, while exp(-∞) = 0 removes it.

Key takeaways

• Masks act on scores before softmax, so blocked pairs get exactly zero weight.
• Causal masks stop the model from reading future tokens, both in training and generation.
• Padding masks stop filler tokens in a batch from affecting the output.
• Use -∞ for blocked scores, and keep the diagonal open for causal masks.

Further learning

Multiple attention patterns run side by side in the next step: multi-head attention.