Causal Masking, Additive Mask Terms, and Where the Mask Enters the Softmax

A decoder-only language model predicts token $t+1$ from tokens $1 \dots t$. If the attention layer at position $i$ could read position $j > i$, training on a full sequence would leak the answer into the input and the model would learn to copy the next token rather than predict it. The fix is a mask, and the interesting part is not that a mask exists but where in the arithmetic it is applied, because a mask applied in the wrong place produces outputs that are numerically similar and behaviourally wrong.

A mask is an additive term on the logits

The masked score matrix is

$$S = \frac{QK^\top}{\sqrt{d_k}} + M, \qquad M_{ij} = \begin{cases} 0 & j \le i \ -\infty & j > i \end{cases}$$

and the softmax is taken over the row of $S$. Adding $-\infty$ makes $e^{-\infty} = 0$, so the forbidden positions contribute nothing to the numerator and nothing to the denominator. The remaining probabilities renormalize automatically, because the denominator sums only the terms that survived. Nothing else in the layer changes.

The alternative formulation, computing the unmasked softmax, multiplying by a 0/1 mask, then dividing each row by its sum, is algebraically identical and strictly worse. It computes exponentials for entries that will be discarded, which is wasted work proportional to the masked fraction: for a causal mask that is half of all $n^2$ entries. It also has a different failure mode when a row is entirely masked, since the renormalizing division is then $0/0$.

Writing the mask as an additive term also makes composition trivial. Two constraints on the same attention layer, causality and padding, are combined by adding their matrices, or equivalently by a logical OR of their boolean forms before conversion to the additive form. There is no ordering requirement between them and no interaction beyond the union of the forbidden sets.

The value $-\infty$ is not representable in floating point, so implementations substitute either the most negative finite value of the dtype or a large finite negative constant. The choice matters. In 16-bit floating point the largest magnitude is 65,504, so a constant of $-10^9$ rounds to $-\infty$ on conversion and the arithmetic proceeds correctly; a constant of $-60{,}000$ is representable and gives $e^{-60{,}000 + \text{rowmax}}$, which underflows to zero after the max subtraction only if the row’s maximum is not itself near that value. In bfloat16 the exponent range is the same as 32-bit, so $-10^9$ stays finite and the exponential underflows to zero cleanly. Both conventions produce exact zeros for the masked entries; the difference appears when a masked logit participates in the row maximum computation, where a finite substitute can be selected as the max and shift the whole row.

import numpy as np

def masked_softmax(scores, keep):
    # keep: boolean (n_q, n_k), True where attention is allowed
    s = np.where(keep, scores, -np.inf)
    s = s - s.max(axis=-1, keepdims=True)     # shift for stability
    e = np.exp(s)
    return e / e.sum(axis=-1, keepdims=True)

n_q, n_k = 4, 4
causal = np.tril(np.ones((n_q, n_k), dtype=bool))
scores = np.random.default_rng(0).normal(size=(n_q, n_k))
print(masked_softmax(scores, causal))

Two masks in one layer, combined by disjunction

A decoder layer holds two attention sublayers with different masking. Self-attention is causal: position $i$ sees positions $\le i$. Cross-attention reads the encoder output, which is fully computed, so no causal constraint applies there; its only mask is padding, which hides encoder positions that hold no real token.

Padding masks come from batching sequences of unequal length. A batch of three sequences of lengths 40, 61, and 57 is stored as a $3 \times 61$ tensor with the tail positions filled by a placeholder id. Those positions produce embeddings and then keys and values, and without a mask they would be attended to as if they carried content. The mask marks them as invalid on the key axis, so every query row excludes them. The query axis is usually left unmasked, because the padded query rows produce outputs that are simply discarded downstream; masking them as well is correct but costs a branch.

The two masks have different shapes and different broadcast patterns. A causal mask is $n_q \times n_k$ with a band structure and depends only on the two index vectors. A padding mask is $1 \times n_k$ in the common case where only keys are padded, broadcast across all query rows. In cross-attention with a padded decoder input, both axes can be padded, giving a $n_q \times n_k$ mask that is the outer product of two padding vectors. Implementations that accept a single boolean mask argument must therefore specify which axes are broadcast, and a mask of shape $n_q \times n_k$ that is actually a per-row vector silently produces wrong results when the broadcasting rules coerce it.

There is a third case worth naming because it appears in prefix language models and in some retrieval setups: a document is attended bidirectionally within itself but cannot see later documents in the packed sequence. That is neither a full causal mask nor no mask, but a block-diagonal band: positions inside the same segment are all visible to each other, positions in a later segment are not. It is expressed as a mask whose entries depend on segment ids rather than on index comparison, and the arithmetic is unchanged.

The all-masked row and the NaN it produces

If every entry of a row is masked, the shifted logits are all $-\infty$, the exponentials are all zero, and the normalizing sum is zero. The division returns NaN. On the backward pass the NaN multiplies through the softmax Jacobian and contaminates the gradients of every parameter in the layer, and because gradient accumulation sums across the batch, one bad row poisons the whole step.

This is reachable in ordinary code. A causal mask on a batch where one sequence has zero real tokens, an empty prompt or a filtered sample, masks the entire row. A padding mask applied to both axes with the same vector masks the diagonal too, and any row whose own position is padded becomes fully masked. A cross-attention layer over an empty retrieval set does the same.

Three responses are common, and they differ in what they cost. The first is to detect the condition and set the row’s output to zero, which requires a reduction over the mask and a select, adding a small amount of work per row. The second is to guarantee the diagonal is never masked, which is possible for causal masks and impossible for padding masks on the query axis. The third is to add a small positive constant to the denominator, which avoids the NaN but produces a uniform distribution over the masked positions, silently attending to padding rather than to nothing. The third is the one to avoid: it converts a loud failure into a quiet one.

Block structure decides where masking costs nothing

The mask matrix is $n^2$ entries, and its representation is a real cost. As a boolean array at one byte per entry, a 4,096-token sequence needs 16 MiB per attention layer per batch element; as a float32 additive mask it needs 64 MiB. That is a per-layer cost that grows quadratically while the cache grows linearly, so implementations either generate the mask on the fly from index comparisons or represent it implicitly.

The causal mask is fully described by the comparison $j \le i$, which costs no memory at all and can be evaluated inside the kernel per element. Tiled implementations go further: partition the score matrix into blocks of $B \times B$. A block strictly below the diagonal is entirely allowed and needs no mask at all; a block strictly above is entirely forbidden and can be skipped without computing its scores. Only the blocks on the diagonal need per-element masking. For a sequence of length $n$ and block size $B$ there are $(n/B)^2$ blocks in total and $n/B$ of them lie on the diagonal, so the fraction needing element-wise masking is $B/n$: at $B = 128$ and $n = 4096$, one block in 32. The skipped upper triangle is half of all blocks, which is where the roughly twofold saving over bidirectional attention comes from. That is why block size is a tuning parameter rather than an implementation detail. The tiling itself is covered in IO-aware attention tiling.

Padding masks have no such structure; they are arbitrary per batch element. Their cost is the gather or the comparison against a length vector, which is $O(n)$ rather than $O(n^2)$, and the practical rule is to keep them out of the inner loop by compacting or by grouping sequences of similar length into the same batch. Variable-length batching schemes that pack multiple sequences into one row without padding avoid the mask entirely at the cost of the segment mask described above, which is again a comparison rather than a stored matrix.

The interaction with position encodings deserves a closing note, since masking and position interact in a way that is easy to get wrong. A causal mask already restricts each output position to a prefix, so the layer is not permutation-equivariant even with no positional encoding; the mask supplies part of the ordering signal. Adding an absolute encoding then supplies absolute position on top of that, and the two together are what let a decoder distinguish “the token three positions back” from “the token at absolute index 3”. The distinction between the symmetries each mechanism provides is developed in permutation equivariance and position encodings, and the rotary variant that replaces absolute addition with a rotation is in rotary position embeddings and relative offsets.