Causal Masking

Attention as written attends to everything. To model $p(x_t \mid x_{<t})$ the operation has to be prevented from reading forward, and the cheapest way to do that is to add $-\infty$ to the logits it must not use.

The mask is additive rather than multiplicative because it has to act before normalisation. Zeroing entries of $A$ afterwards would leave the rows summing to less than one; adding $-\infty$ to the logits removes the terms from the denominator as well.

What it buys

Every row of the masked attention matrix is a valid prediction context. A sequence of length $n$ therefore yields $n$ training examples from one forward pass, at the cost of one — this is the reason decoder-only pre-training is as sample-efficient as it is, and it is a property of the mask, not of the architecture.

Variants worth knowing

The strictly triangular mask is the common case, not the only one. Document packing wants a block-diagonal mask so that concatenated documents cannot attend across the join. Prefix-LM wants a rectangular bidirectional region. Sliding-window attention is a band mask. All of them are the same mechanism with a different boolean array, which is why a kernel that takes an arbitrary mask is more useful than one with is_causal hard-coded — though the hard-coded one is faster, because it can skip whole tiles.

The off-by-one that will bite you

The mask allows $j \le i$, not $j < i$: a position attends to itself. Getting this wrong shifts the model by one token and produces a loss curve that looks plausible, converges, and generates nonsense. Check it by feeding the model a sequence and confirming that the logits at position $t$ are unchanged when position $t+1$ is altered.

import torch
from torch import Tensor


def causal_mask(n: int, device: torch.device) -> Tensor:
    """True where attention is forbidden. Note `diagonal=1`: self is allowed."""
    return torch.ones(n, n, dtype=torch.bool, device=device).triu(diagonal=1)


def document_mask(doc_ids: Tensor) -> Tensor:
    """Causal *and* within-document, for packed training batches. [B, N] -> [B, 1, N, N]."""
    n = doc_ids.size(-1)
    causal = torch.ones(n, n, dtype=torch.bool, device=doc_ids.device).triu(1)
    cross = doc_ids[:, :, None] != doc_ids[:, None, :]
    return (causal | cross)[:, None]


def apply(logits: Tensor, mask: Tensor) -> Tensor:
    # A finite floor, so a fully masked row degrades to uniform rather than NaN.
    floor = torch.finfo(logits.dtype).min / 2
    return logits.masked_fill(mask, floor)

Build the mask once and cache it by length; regenerating a $4096 \times 4096$ boolean tensor per step is a measurable fraction of a small model’s step time. Better still, pass is_causal=True to a fused kernel and never materialise it — the tiles above the diagonal are then skipped rather than computed and discarded, which is most of the point.