NLP Foundations
Causal Masking
How a transformer trained to predict every position in parallel is stopped from cheating by looking ahead, and why the mask is applied before softmax rather than after.
intermediate · 7 min read
Train a language model with teacher forcing (see next-token-prediction-cross-entropy) and every position in the sequence is scored in a single forward pass. That is only valid if position i's prediction cannot see position i+1 or anything after it; otherwise the model could trivially "predict" the next token by copying it straight out of the input it can already see, and training loss would collapse to near zero while the model learned nothing useful. Causal masking is the mechanism that prevents this: it forces every position's self-attention to only look backward.
Where the mask goes
The mask is applied to the raw attention scores, before softmax, not to the attention weights after softmax. For position i attending to position j, the score is set to negative infinity whenever j > i:
scores = Q @ K.T / sqrt(d_k)
scores[i, j] = -inf for all j > i
weights = softmax(scores) # rows now zero out future positions
output = weights @ V
Applying the mask before softmax matters because softmax is what turns scores into a distribution that sums to one; a score of negative infinity becomes exp(-inf) = 0 after exponentiation, so the future position receives exactly zero weight, with no numerical residue and no need to renormalise afterward. Masking the weights directly after softmax and zeroing them out would leave the remaining weights not summing to one, silently changing the effective temperature of the attention distribution.
Why this makes parallel training possible
Because the mask makes each row's attention pattern legal on its own, right down to the last position, the entire sequence can be scored in one matrix multiplication. Every position computes its next-token prediction using only the tokens that would actually be available at that point during real, left-to-right generation, but all of those predictions happen simultaneously across the sequence in a single forward pass. This is what makes transformer pretraining as fast as it is: an n-token sequence produces n training signals per forward pass, not one signal after n sequential steps. Without causal masking this parallelism would be unsound, because a position's "prediction" would have had access to the answer.
Training-time masking versus inference-time reality
At inference, causal masking is not really doing anything extra: when generating token by token, future tokens do not exist yet, so there is nothing to mask. Causal masking's real job is making training (where the whole sequence is present at once) behave as if it were happening left to right, one token at a time, honestly. This is also exactly why the KV cache (see kv-cache) is a safe optimisation at inference: a masked position's output never depended on future keys and values in the first place, so caching past keys and values and only computing the new token's query against them reproduces training-time behaviour exactly, just without redoing work that provably cannot change.
Bidirectional attention is the alternative, not a variant
Encoder-only models like BERT use no causal mask at all; every position attends to every other position, forward and backward, because the training objective (masked-token prediction on scattered positions, not next-token prediction) does not require left-to-right honesty. This is a genuinely different training setup, not a relaxed version of the same one: a bidirectional model cannot be used for autoregressive generation without retraining, because its representations were built with access to the future the whole time.
When it falls down
- The mask is a hard cutoff, not a preference. Causal masking is binary: a future position contributes exactly zero, never a discounted amount. Some architectures relax this deliberately (prefix-LM masking lets a bidirectional prefix precede a causal continuation), trading strict autoregressive honesty for richer context on the fixed part of the input.
- It does not save compute by itself. A naively implemented causal mask still computes the full
n x nscore matrix and then throws away the upper triangle, which is exactly the wasted work that sliding-window-attention and fused kernels like FlashAttention (see flash-attention) are designed to avoid computing in the first place. - Off-by-one errors are a real bug class. Whether position
ican attend to itself (j <= ivsj < i) is easy to get subtly wrong when implementing a mask by hand, and the failure mode, a model that can see one token further ahead than it should, is often invisible until an eval quietly underperforms.
Further reading
- Vaswani et al., 2017, Attention Is All You Need, arXiv:1706.03762 - introduces masked (causal) self-attention in the decoder.
- Radford et al., 2018, Improving Language Understanding by Generative Pre-Training (GPT) - the decoder-only, fully causal architecture that became the LLM default.
- Devlin et al., 2018, BERT: Pre-training of Deep Bidirectional Transformers, arXiv:1810.04805 - the bidirectional (unmasked) contrast case.