The causal mask

A language model must predict each token without seeing it. The causal mask is a triangle of −∞ added to the attention scores so that every token can only look at itself and earlier tokens. It also lets training check a prediction at every position of a text in one pass.

Key idea: Add −∞ above the diagonal of the score matrix: softmax turns it into weight 0, so position i only reads positions 1 to i
Real world: A teacher gives a class a story one line at a time, with a card covering the rest of the page, and asks after each line: “what word comes next?”. Everyone answers for every line, but nobody can read ahead.
In an LLM: GPT-2 reads “The cat sat on the mat because it was tired.” in one pass and makes 11 predictions, one per position. Thanks to the mask, the prediction made at “sat” was computed without seeing “on”, so it is a fair test.

The rule

A language model generates text left to right. When it predicts token 6, tokens 7, 8 and 9 don’t exist yet. So when it learns, it must learn under the same conditions: each position may only use the tokens up to and including itself.

Attention is the only place where tokens read from each other, so that is where the rule is enforced. Before the softmax, every score where the key comes after the query is replaced by $-\infty$ :

M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases} \qquad A = \operatorname{softmax}\!\Big(\frac{QK^\top}{\sqrt{d}} + M\Big)

Since $e^{-\infty} = 0$ , those keys get a weight of exactly zero, and the remaining weights in the row still add up to 1. The allowed cells form a lower triangle. This is the : “causal” because each token can only be influenced by what came before it, like causes before effects.

row 5: “the”

written so far:

scores + mask5 × 5

ThecatsatontheThe0.8−∞−∞−∞−∞cat0.00.2−∞−∞−∞sat0.10.60.5−∞−∞on0.61.60.30.0−∞the0.81.70.50.10.2

→

weights (softmax)5 × 5

ThecatsatontheThe1.000.000.000.000.00cat0.440.560.000.000.00sat0.240.410.350.000.00on0.210.530.150.110.00the0.180.490.140.090.10

−∞ (masked): weight exactly 0

Row “the” (position 5) can see 5 tokens: itself and everything before it. It is the last token, so nothing is masked in its row. Stacked up, the allowed cells form a triangle.

The toy sentence “The cat sat on the”. Press play to compute one row at a time, as if each word had just been written. Switch to bidirectional to see what changes when every word may look at every other.

Switch to bidirectional and every word can read every other. The weights change: in our toy, “cat” now looks ahead to “sat”. That is the attention of BERT and other encoder models, built to understand a finished text. GPT, Llama, Claude and every other text generator are decoder models, and use the causal mask.

example · covering the page

Reading a mystery novel, you can guess the culprit at the end of each chapter using everything you have read so far. If you had flicked to the last page first, your guesses would be perfect, and you would have learned nothing about spotting clues. The card that covers the rest of the page is the causal mask.

Training on every position at once

The mask has a second, less obvious benefit. Because no row can see its own answer, every row can be a separate training example, all in the same forward pass.

Feed in a text of $T$ tokens. Position 1 predicts token 2, position 2 predicts token 3, and so on: $T - 1$ predictions to check against tokens you already have, and one more at the end that predicts what comes next. One big matrix multiplication computes all the rows at once on a GPU.

Loading GPT-2 data…

Tick remove the mask and the problem is obvious: every row can now read the token it is meant to predict. A model trained like that learns to copy the next token, its training loss drops to nearly zero, and it is useless the moment it has to write something new.

example · one pass, ten lessons

GPT-2 reading “The cat sat on the mat because it was tired.” (11 tokens) produces 11 probability distributions in a single pass. At “sat” it gives “ on” 22%; at “on” it gives “ the” 45%; at “the” its top guess is “ floor” and the real “ mat” isn’t in its top 5. In training, all ten of those checkable guesses contribute to the loss at once. Its top guess is right at 5 of the 10.

Generation reuses the past

During generation the model adds one token at a time. Thanks to the mask, the rows for the earlier tokens never change when a new token is appended: they couldn’t see it. So their keys and values can be kept from the previous step instead of recomputed. That saved work is the KV cache, and it only works because attention is causal.

Other masks

The mask is just a matrix of 0s and $-\infty$ s, so it can have other shapes. Each shape is a different rule about who may read whom:

causal (GPT): 36 of 64 pairs

123456781 The2 cat3 sat4 on5 the6 mat7 .8 It

bidirectional (BERT): 64 of 64 pairs

123456781 The2 cat3 sat4 on5 the6 mat7 .8 It

BERT-style encoders let every token see every other token, both directions. That is ideal for understanding a whole text (classifying it, filling in a blank), but it can’t generate: a word would see the words it is supposed to predict.

Grey cells are masked (score −∞, weight 0); purple cells are allowed. The left grid is always the causal mask of a GPT-style model; pick a variant for the right grid.

Bidirectional (no mask): encoders such as BERT, and image models, where there is no “future”.
Padding: when sentences of different lengths share a batch, the short ones are filled with padding tokens, and a padding mask hides them. It is combined with the causal mask in a decoder.
Prefix-LM: the prompt is read in both directions, the answer is generated causally.
Sliding window: each token sees only the last $w$ tokens, so cost grows with $w$ instead of the whole length. See efficient attention.

going deeper

Why −∞ and not 0

Setting a weight to 0 after the softmax would leave the row summing to less than 1. Adding $-\infty$ to the score before the softmax removes the key from both the numerator and the denominator, so the remaining weights are renormalised automatically:

a_{ij} = \frac{e^{s_{ij}}}{\sum_{j' \le i} e^{s_{ij'}}} \ \text{ for } j \le i, \qquad a_{ij} = 0 \ \text{ for } j > i.

In code a large negative number (or the dtype’s minimum) is used instead of a true $-\infty$ , to avoid NaN from rows that are entirely masked.

Training objective

With the causal mask, one forward pass over $x_1, \dots, x_T$ gives $p_\theta(x_{t+1} \mid x_{\le t})$ for every $t$ , and the loss is the average over positions:

\mathcal{L} = -\frac{1}{T-1} \sum_{t=1}^{T-1} \log p_\theta(x_{t+1} \mid x_1, \dots, x_t)

This is the chain rule of probability, evaluated in parallel. It is often called teacher forcing: the model always conditions on the true previous tokens, not its own guesses.

In code

T = x.shape[1]
causal = torch.ones(T, T, dtype=torch.bool).tril()   # True = allowed
pad = attention_mask[:, None, None, :].bool()        # (B, 1, 1, T)
allowed = causal & pad
neg = torch.finfo(scores.dtype).min
scores = scores.masked_fill(~allowed, neg)

In papers

Masked decoder self-attention: Vaswani et al., 2017, section 3.2.3.
Bidirectional encoders: Devlin et al., BERT, 2018.
Prefix-LM and other mask patterns compared: Raffel et al., T5, 2019, section 3.2.
Sliding-window attention: Beltagy et al., Longformer, 2020; Jiang et al., Mistral 7B, 2023.