Chapter 4 · Section 4.4 · 50 min

Masks, causal and padding

Here is a language model that reaches perfect accuracy on its training data in a single step, and is completely useless. Its one attention head makes every position $t$ look at position $t+1$ and copy the token it finds there. Asked to predict the next token, it reads it. Training loss goes to zero. Then you try to generate text, position $t+1$ does not exist yet, and the model has nothing.

Nobody builds that head on purpose, but gradient descent will find it in an afternoon if you let it. A language model is trained on whole sequences in parallel, with every position predicting its successor at once (§3.1), and attention as we have defined it lets every position read every other. The fix is to forbid some connections. A mask says, for each pair (query $i$ , key $j$ ), whether $i$ may attend to $j$ . You have used one already. In Exercise 3.6a a decoder query gave padded source positions a score of $-\infty$ , and Chapter 3 ended by noting that a decoder must also be kept from reading words it has not written yet. This section turns those two remarks into the masks every practitioner needs (causal, padding, and their relatives), shows how to combine them, what to do when a row is allowed to see nothing at all, and the conventions that trip people up when moving between libraries.

A model that reads the future learns nothing

Let's build the cheating head and watch it work. Tokens are one-hot ids concatenated with one-hot positions, as in Exercise 4.2b, but now the query of position $t$ matches the key of position $t+1$ , a next-token head.

Python · cell 1

Now score it as a next-token predictor, with and without a mask that hides the future.

Python · cell 2

Without a mask the head is a perfect predictor, 100% accuracy, because it copies the answer. With the causal mask it falls to 10%, chance level for 10 tokens, because the answer is out of reach. That 10% is the honest starting point, and everything a real model learns above it is genuine prediction.

The causal mask

The causal mask allows query $i$ to attend to key $j$ exactly when $j \le i$ . As a boolean matrix it is lower triangular, with True on and below the diagonal. As an additive mask $\mathbf{M} \in \R^{T\times T}$ it is

M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases}, \qquad \mathbf{A} = \operatorname{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}} + \mathbf{M}\right).

In plain words, before the softmax, every score that would let a position see its future is replaced by $-\infty$ , and since $e^{-\infty} = 0$ , those weights are exactly zero.

Two details matter. The diagonal is allowed, since a position may always look at itself (the token at position $t$ is part of the input when predicting position $t+1$ ). And the mask goes on the scores, before the softmax, so the remaining weights renormalize among the allowed keys. Zeroing weights after the softmax without renormalizing leaves rows that sum to less than 1. (Zeroing and then renormalizing gives the same result as the $-\infty$ mask, but costs an extra pass.)

Why is this enough to stop all cheating, through every layer of a deep model? By induction. Suppose that after layer $\ell$ , the vector at position $i$ depends only on input tokens $1, \dots, i$ . In layer $\ell+1$ , attention at position $i$ mixes the values at positions $j \le i$ , each of which depends only on tokens up to $j \le i$ . The MLP, the normalization, and the residual connection act on each position separately (§5.1). So after layer $\ell+1$ , position $i$ still depends only on tokens up to $i$ . The input embeddings satisfy the claim, so every layer does. Let's verify the claim numerically, by changing a future token and checking that no past output moves.

Python · cell 3

The causal mask is also what makes training efficient. One forward pass over a sequence of $T$ tokens produces $T$ predictions, each using exactly the context it would have at generation time, so every sequence yields $T$ training examples at once. This is teacher forcing (§3.5), now done in parallel instead of one step at a time. And because the mask matches what a position can see during generation, the keys and values computed for earlier positions never change as new tokens arrive, which is what makes the KV cache possible (§6.5).

Core

Prefix-LM and sliding-window masks

Two more rules from the callout above, as boolean (T, T) masks with True meaning "may attend". prefix_lm_mask(T, P) lets the first P positions (the prompt) attend to each other in both directions, while every later position stays causal and sees the whole prompt. sliding_window_mask(T, w) is causal and lets query $i$ see only the last w positions including itself, keys $j$ with $i - w < j \le i$ . Mistral 7B uses this rule with $w = 4{,}096$ (Jiang et al. 2023).

import numpy as np

def prefix_lm_mask(T, P):
    return np.tril(np.ones((T, T), bool))       # placeholder: plain causal, the prompt cannot see ahead

def sliding_window_mask(T, w):
    return np.tril(np.ones((T, T), bool))       # placeholder: no window

Hidden tests, hints, and AI help, from $1.99 a month.

Padding masks

To run sequences of different lengths as one batch, we pad them to a common length $T$ with a PAD token (§2.5). This is the mask of Exercise 3.6a, now for every query of a batch at once. Padding is filler, and no real token should attend to it, or the output would depend on how much padding happened to be added. So we mask every key position that holds padding. For a batch with valid lengths $\ell_1, \dots, \ell_B$ , the key-padding mask allows key $j$ of sequence $b$ exactly when $j < \ell_b$ . It has shape $(B, T)$ and is broadcast to $(B, T_q, T)$ , or $(B, h, T_q, T)$ with heads, by inserting axes.

What about padded queries, the rows belonging to pad positions? They compute some output, and we ignore it. Their predictions are excluded from the loss with a loss mask (for example ignore_index in PyTorch's cross-entropy). This is a separate mask with a separate job. The attention mask decides who sees whom. The loss mask decides which positions are graded. Mixing them up is a common bug. Masking prompt tokens out of the loss during fine-tuning, for instance, is a loss mask, and the prompt must stay visible to attention (§14.3).

Masks combine by logical AND, since a connection is allowed only if every rule allows it. With additive masks, the same thing is a sum, because $0 + 0 = 0$ and anything plus $-\infty$ is $-\infty$ .

Python · cell 4

The last rows of sequence 2 are padded queries. They attend to the two real keys, produce an output nobody will use, and cause no harm. With right padding, as here, every query has at least one allowed key, because key 0 is always real and always in the past. That is not true of every layout, and the exceptions are where NaNs come from.

Warm-up

Causal and padding masks for a batch

Implement make_mask(T, lengths=None, causal=True). With lengths=None, return a (T, T) boolean mask (all True if causal is false). With lengths, a list of $B$ valid lengths for right-padded sequences, return a (B, T, T) mask in which query $i$ of sequence $b$ may attend to key $j$ only if $j$ is a real token ( $j <$ lengths[b]) and, if causal, $j \le i$ . True means "may attend".

import numpy as np

def make_mask(T, lengths=None, causal=True):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Attention to nothing

Some rows of a mask can end up entirely False. Here are three ways it happens in practice.

Left padding. Batched generation usually pads on the left, so that all sequences end at the same position and new tokens can be appended in a column. With a causal mask, a pad position at the very start may attend only to keys at or before it, and those are all padding. Its row is empty.
Masks built from several rules, such as padding inside packed sequences or a sliding window that has slid past every real token.
Cross-attention to an empty source, for example an image-conditioned model run on a batch where one example has no image.

With $-\infty$ in every entry of a row, the softmax computes $e^{-\infty - (-\infty)}$ , which is $e^{\text{NaN}}$ , and the row becomes NaN. One NaN row is enough. It contaminates every later layer through the values, and the loss and the gradients become NaN everywhere. Implementations guard against this in one of four ways.

choice	weights in an empty row	other rows	pitfalls
$-\infty$ , plain softmax	NaN	exact	poisons the whole batch
large negative, such as $-10^9$	uniform over the masked keys (if filled in, see below)	exact, since $e^{-10^9}$ underflows to 0	$-10^9$ overflows to $-\infty$ in float16, bringing the NaN back (bfloat16 has float32's range, so it is safe there)
the dtype's most negative value, `finfo(dtype).min`	uniform over the masked keys	exact	adding it to a negative score can overflow to $-\infty$ (in float16, any score below about $-16$ does), so use it with `where` (masked fill), not addition
"safe" softmax, rows with no allowed key return all zeros	zeros, so the output is zero	exact	needs an explicit check, but gives the cleanest semantics

The large-negative options never produce NaN, but they quietly make an empty row attend to the very positions it was told to ignore. With a masked fill every blocked score becomes the same number, so the row is exactly uniform. With an additive mask the row becomes $\operatorname{softmax}(\mathbf{s} - 10^9)$ , which in float64 is almost exactly the softmax of the original scores, while in float32 the spacing between representable numbers near $10^9$ is 64, so the scores all round to the same value and the row is again uniform. Either way the row ignores its mask. That is harmless when the row's output is discarded (a padded query) and wrong when it is not. The safe softmax is the most honest. An empty row means "there is nothing to look at", and a zero output says exactly that. Recent versions of PyTorch return zeros for fully masked rows in scaled_dot_product_attention (we checked the math and CPU kernels of version 2.14), while a bare torch.softmax over a row of $-\infty$ still returns NaN. This behavior has changed across versions and can differ between kernels, so test the one you deploy. Exercise 4.4b asks you to write the safe version.

Python · cell 5

Interview

A numerically safe masked softmax

Implement masked_softmax(S, mask) for scores S of shape (..., n) and a boolean mask broadcastable to it (True means "may attend"), normalizing over the last axis. Requirements, all tested. Masked entries are exactly 0. Rows with at least one allowed entry sum to 1 and match the ordinary softmax over the allowed entries. Rows with no allowed entry are all zeros, not NaN and not uniform. It is stable for scores in the thousands. Do not let NumPy compute -inf - (-inf) anywhere.

import numpy as np

def masked_softmax(S, mask):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

When a head wants to attend to nothing

An empty row is a mask's way of saying "nothing to see". But even with every key allowed, a head often has nothing it wants. Think of a head that tracks the subject of a verb, sitting at a punctuation mark. Softmax gives it no way to abstain, because the weights must sum to 1. Trained models find a workaround. They learn to dump the unwanted weight on a token that carries little information, most often the first token of the sequence. Xiao et al. (2023) named these attention sinks, and showed that the first few tokens receive a large share of attention in many layers of models such as Llama 2, regardless of what those tokens are. Removing the sink tokens from the cache during streaming generation makes quality collapse, which is how they discovered how much the model relies on them.

If abstaining is useful, why not make it explicit? Add a constant to the softmax denominator,

a_{ij} = \frac{e^{s_{ij}}}{e^{b} + \sum_{l} e^{s_{il}}},

so the weights sum to less than 1 and the head can return (nearly) zero by making all its scores small compared with $b$ . With $b = 0$ this is the "softmax plus one" proposed by Miller (2023). OpenAI's gpt-oss models (2025) learn a separate $b$ for every head, which their model card describes as letting a head pay no attention to any token. Equivalently, a sink is an extra key whose score is $b$ for every query and whose value is the zero vector. The weight it absorbs, $e^b/(e^b + \sum_l e^{s_{il}})$ , is weight that no real token receives.

Interview

Attention with a learned sink

Implement sink_attention(Q, K, V, sink, mask=None) for one layer of heads, with Q of shape (h, T_q, d_k), K of shape (h, T_k, d_k), V of shape (h, T_k, d_v), one sink logit per head in sink of shape (h,), and an optional boolean mask broadcastable to (h, T_q, T_k) (True means "may attend"). The weights are $a_{ij} = e^{s_{ij}}/(e^{b} + \sum_l e^{s_{il}})$ with $s_{ij} = \mathbf{q}_i^\top\mathbf{k}_j/\sqrt{d_k}$ and $b$ the head's sink logit, so each row sums to less than 1. Return (O, A). It must be stable for scores and sink logits in the thousands, give masked keys exactly zero weight, and return a zero output, not NaN, for a row whose keys are all masked.

import numpy as np

def sink_attention(Q, K, V, sink, mask=None):
    S = Q @ np.swapaxes(K, -1, -2) / np.sqrt(Q.shape[-1])
    A = np.exp(S) / (np.exp(sink)[:, None, None] + np.exp(S).sum(-1, keepdims=True))   # right formula, unstable
    return A @ V, A

Hidden tests, hints, and AI help, from $1.99 a month.

Packed sequences, a preview

Padding wastes compute, and real documents vary so much in length that padding a batch to its longest member can waste most of every matrix multiply. Krell et al. (2021) found that padding can make up as much as half of all tokens in common NLP datasets. The standard fix is packing, concatenating several documents into one row of length $T$ , separated by end-of-text tokens, so that every row is full. §6.1 builds that pipeline and measures the waste. Here we need only the mask it calls for.

Without one, the first token of document 2 can attend to all of document 1, which it has nothing to do with. Many pretraining runs accept that cross-contamination (GPT-3 relied on the end-of-text token alone, Brown et al. 2020). The exact fix is a block-diagonal mask. Give token $i$ a document id $g_i$ , and allow query $i$ to attend to key $j$ only when

g_i = g_j \quad\text{and}\quad j \le i .

In plain words, every document gets its own causal triangle, and the triangles never touch.

With this mask each document's outputs are identical to what you would get by running it alone, as the cell checks. (With absolute position embeddings the position ids must also restart at each document, a detail §6.1 handles.) The Llama 3 report describes training with exactly this mask (Llama Team 2024), and production kernels never build the $T\times T$ matrix for it. They take the document boundaries and skip the off-diagonal blocks, FlashAttention's variable-length interface of §12.4.

Python · cell 6

Causal masks when queries and keys have different lengths

So far the score matrix was square. During generation with a KV cache (§6.5), the queries are the last $T_q$ positions of a sequence whose keys span all $T_k$ positions so far, with $T_q \le T_k$ . Query $i$ (counting from 0 within the chunk) sits at absolute position $T_k - T_q + i$ , so the causal rule becomes

\text{query } i \text{ may attend to key } j \iff j \le i + (T_k - T_q).

In plain words, the triangle is anchored at the bottom right of the $T_q\times T_k$ score matrix. For a single new token ( $T_q = 1$ ) the row is all True, which is why plain decoding with a cache needs no mask at all. Here is a trap worth knowing now. torch.nn.functional.scaled_dot_product_attention with is_causal=True anchors the triangle at the top left when $T_q \ne T_k$ , so query 0 sees only key 0. (PyTorch provides torch.nn.attention.bias.causal_lower_right for the other alignment, and FlashAttention uses bottom-right alignment.) §6.5 builds the cache and processes prompts in chunks with this mask.

In PyTorch

The same masks in PyTorch. Note how three popular APIs disagree about what True means.

🔥 In PyTorchfor a machine with PyTorch

API	mask argument	`True` means
`F.scaled_dot_product_attention`	`attn_mask` (bool)	may attend
`nn.MultiheadAttention`	`attn_mask`, `key_padding_mask` (bool)	may not attend
Hugging Face models	`attention_mask` (0 or 1, per token)	1 is a real token, 0 is padding
FlexAttention	`mask_mod(b, h, q_idx, kv_idx)`	returns True where attention is allowed

For document masks and other structured patterns, PyTorch's FlexAttention lets you describe the mask as a function and compiles a kernel that skips fully masked blocks. A block-diagonal causal mask looks like this. (It is shown, not run, because it needs torch.compile to be fast.)

🎯 Interview questions on this section

Why does a decoder-only language model need a causal mask during training but not, for one new token, during generation?

Training processes all positions in parallel, with position $t$ predicting token $t+1$ , and without a mask each position could read its own target. The causal mask restricts position $t$ to keys $j \le t$ , exactly the context available at generation time. When generating one token with a KV cache, the single query is at the last position, and every cached key is in its past, so the mask row is all True and no mask is needed.

Prove that a causal mask in every attention layer makes the whole Transformer causal.

By induction on layers. Embeddings are per-position, so position $i$ depends only on token $i$ . If after layer $\ell$ position $i$ depends only on tokens $\le i$ , then layer $\ell+1$ 's attention at $i$ mixes values at $j \le i$ , each depending only on tokens $\le j \le i$ , and the MLP, normalization, and residual connection are per-position. So position $i$ still depends only on tokens $\le i$ . One leak anywhere breaks it, for instance a normalization computed across the sequence dimension.

Should you mask with

-\infty

or with a large negative number like

-10^9

$-\infty$ gives exact zeros but turns fully masked rows into NaN ( $-\infty - (-\infty)$ ), which then poisons everything downstream. A large negative avoids NaN, still gives exact zeros in normal rows because the exponential underflows, but makes fully masked rows attend uniformly to the masked keys. In float16, $-10^9$ itself overflows to $-\infty$ , and adding the dtype minimum to a score below about $-16$ overflows too, so implementations use masked fill with the dtype's minimum, or a safe softmax that returns zeros for empty rows.

When can a row of the attention mask be entirely masked, and what do you do about it?

With left padding plus a causal mask (pad queries at the start see only pad keys), with combinations of window, document, and padding masks, or in cross-attention to an empty source. Options are a safe softmax that outputs zeros, a large finite negative fill so the row is uniform and its output discarded, or adding a sink (an always-visible key with zero value, or a constant in the denominator) so every row has somewhere to put its weight.

What is an attention sink, and why do models create one?

Softmax weights must sum to 1, so a head with nothing relevant to attend to still has to put its weight somewhere. Trained models learn to put it on a low-information token, typically the first one. Xiao et al. (2023) showed that the first tokens receive outsized attention across many heads, and that evicting them from the cache during streaming breaks generation. Explicit fixes add a constant or a learned per-head logit to the softmax denominator, so weights can sum to less than 1.

How do you train on packed sequences without cross-contamination? What else besides the mask must change?

Give each token a document id and allow attention only when ids match and $j \le i$ , a block-diagonal causal mask, so each document's outputs equal those of the document alone. With absolute position embeddings, position ids must also restart at 0 at each document boundary (with rotary embeddings the mask alone suffices, since scores depend only on offsets within a document). Efficient kernels take the document boundaries (cumulative sequence lengths) and skip off-diagonal blocks instead of materializing a $T\times T$ mask. The data side, how to pack and how much it saves, is §6.1.

What is the difference between an attention mask and a loss mask?

The attention mask decides which positions can read which, and changes the activations. The loss mask decides which positions' predictions count toward the loss, and changes only the gradient signal. Padding typically needs both (pad keys hidden, pad predictions ignored). Prompt tokens in supervised fine-tuning need only a loss mask, since the response must still attend to the prompt.

How would you let a head attend to nothing, and what does it change?

Add a sink, an extra logit $b$ in the softmax denominator, $a_{ij} = e^{s_{ij}}/(e^{b} + \sum_l e^{s_{il}})$ , which is the same as an extra key with score $b$ and a zero value. The weights then sum to less than 1, so a head whose scores are all far below $b$ outputs nearly zero instead of being forced to average something. With $b = 0$ fixed this is Miller's "softmax plus one", and gpt-oss learns $b$ per head. A side benefit is that a fully masked row is no longer NaN, since the sink keeps the row maximum finite, and its output is exactly zero. The cost is that the output is no longer a convex combination of the values.

You process a long prompt in chunks of 512 tokens with a KV cache. What mask does each chunk need?

The chunk's 512 queries are the last 512 of $T_k$ positions so far, so query $i$ may attend to key $j$ when $j \le i + T_k - 512$ , a causal triangle anchored at the bottom right, combined with any padding or document mask. Using a top-left aligned causal mask (PyTorch's is_causal=True for non-square inputs) would hide almost the whole cache from every query.

🧭 What to remember

A mask decides, for each (query, key) pair, whether information may flow. It is applied to the scores before the softmax, so allowed weights renormalize and blocked ones are exactly zero.
The causal mask ( $j \le i$ ) stops a language model from reading its own targets, makes the whole stack causal by induction, and lets one forward pass train on every position.
Padding masks hide pad keys. Pad queries are harmless and are dropped by a separate loss mask. Masks combine by AND, or by adding additive masks.
Other rules are other boolean matrices. A prefix LM is bidirectional over the prompt, and a sliding window of $w$ keeps the last $w$ keys, so $L$ layers reach $L(w-1)$ back.
Rows with nothing allowed turn $-\infty$ into NaN. Use a safe softmax, a careful finite fill, or a sink, an extra logit that lets a head attend to nothing. Trained models create sinks of their own on the first token.
Packed documents need a block-diagonal causal mask, which makes each document's outputs exactly those it would have alone (§6.1 builds the pipeline).
With fewer queries than keys, as with a KV cache, the causal triangle is anchored bottom right, and PyTorch's is_causal is not (§6.5).
Next, §4.5 runs many attention heads side by side, each with its own pattern, and counts exactly what that costs.

Vaswani et al. (2017). Attention Is All You Need. Section 3.2.3 introduces masking out illegal connections in the decoder's self-attention.
Phuong & Hutter (2022). Formal Algorithms for Transformers. Masks written as explicit matrices for encoder, decoder, and encoder-decoder models.
Krell et al. (2021). Efficient Sequence Packing without Cross-contamination. How much padding costs, packing algorithms, and the masks and position ids that make packing exact.
Jiang et al. (2023). Mistral 7B. A widely used model with sliding-window attention, and how stacked windows reach far back.
Llama Team (2024). The Llama 3 Herd of Models. Section 3.2 describes the document mask used in pretraining.
Xiao et al. (2023). Efficient Streaming Language Models with Attention Sinks. The discovery of attention sinks and why the first tokens matter.
Miller (2023). Attention Is Off By One and OpenAI (2025), gpt-oss-120b and gpt-oss-20b Model Card. Letting heads attend to nothing, by a constant and by a learned sink.
PyTorch documentation. scaled_dot_product_attention. The boolean mask convention and the alignment of is_causal.

Zero to Frontier is made for larger screens

Masks, causal and padding

A model that reads the future learns nothing

The causal mask

Prefix-LM and sliding-window masks

Padding masks

Causal and padding masks for a batch

Attention to nothing

A numerically safe masked softmax

When a head wants to attend to nothing

Attention with a learned sink

Packed sequences, a preview

Causal masks when queries and keys have different lengths

In PyTorch

Read more

Watch next