Core ML

Causal Masking

Force every future position's attention score to negative infinity before the softmax, so a token can only ever attend to itself and whatever came before it.

A lower-triangular mask blocks every position from attending to anything later in the sequence, so the upper-right half of the score grid is forced to zero before softmax

Why Does This Exist?

Self-attention, left alone, lets every token attend to every other token — including tokens later in the sequence. For a model whose only job is to build a representation of a complete, already-written passage, that's exactly right; there's no "later" to worry about because the whole input is available at once.

But training a model to generate text by predicting the next token creates a very specific problem. If token 3's representation is allowed to attend to token 4, and the model is being trained to predict token 4 from token 3's representation, the model can trivially cheat: just copy whatever it already saw at position 4. Training loss looks fantastic, and the model has learned nothing about actually predicting the future — it learned to read the answer key it was handed by accident. At inference time, token 4 doesn't exist yet when you're generating it, so a model trained this way collapses immediately outside of training.

Think of It Like This

An exam where later pages are stapled shut

Imagine an exam booklet where every page after the one you're currently on is stapled closed — you can flip back and reread anything you've already answered, but you physically cannot see ahead to question 12 while you're working on question 5. That's not a memory limitation; it's an enforced rule of the exam, the same for every question in the booklet.

Causal masking is exactly that staple. It doesn't make future tokens hard to attend to — it makes them unavailable, structurally, for every position, so the model that trains under this constraint learns to actually predict rather than to peek.

How It Actually Works

One mask, added before the softmax

The diagram above shows the mechanism directly: build an $n \times n$ mask matrix where position $(i, j)$ is 0 when $j \leq i$ (key position $j$ is at or before query position $i$ ) and negative infinity when $j > i$ (key position $j$ is in the future relative to query $i$ ). Add this mask to the raw attention scores before the softmax:

\text{CausalAttention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V

Adding negative infinity to a score, then exponentiating inside the softmax, sends that term to exactly zero — $e^{-\infty} = 0$ . So every blocked position contributes literally nothing to the weighted output, not a small amount, not an approximation. The diagram's × marks are cells that become exact zeros in the attention-weight matrix, every time, for every input.

Why it's a triangle

Draw out which $(i, j)$ pairs survive — query position $i$ is allowed to see key position $j$ whenever $j \leq i$ — and the surviving region is the lower triangle of the $n \times n$ grid, including the diagonal. The diagonal always survives: a token always attends to itself, since $j = i$ satisfies $j \leq i$ . That's why every accented diagonal cell in the diagram is open while everything above it, later in the sequence, is blocked.

The training-time payoff

Because the mask is applied per-position rather than by physically truncating the input, a causally masked model can compute the loss for every position in a training sequence in one forward pass. Position 5's prediction only ever saw positions 1 through 5, so it's a fair, uncheatable next-token prediction — and so is position 500's, computed in the same batched pass. That parallel-training property, alongside the raw prevention of future-leakage, is a large part of why decoder-only causally masked transformers scaled so well: every position in a sequence supplies a training signal, computed together, with none of them cheating off the others.

Show Me the Code

Building the mask and confirming that blocked positions land at exactly zero after the softmax, not merely small.

1import numpy as np2
3
4def softmax(z: np.ndarray) -> np.ndarray:5    e = np.exp(z - z.max(axis=-1, keepdims=True))6    return e / e.sum(axis=-1, keepdims=True)7
8
9def causal_mask(n: int) -> np.ndarray:10    mask = np.triu(np.ones((n, n)), k=1)      # 1s strictly above the diagonal (the future)11    return np.where(mask == 1, -np.inf, 0.0)   # future -> -inf, everything else -> 012
13
14scores = np.array([[2.0, 1.0, 0.5], [1.0, 3.0, 0.2], [0.1, 0.4, 2.5]])15masked = scores + causal_mask(3)16weights = softmax(masked)17print(np.round(weights, 3))18# -> [[1.    0.    0.   ]19#     [0.119 0.881 0.   ]20#     [0.066 0.089 0.845]]

Row 0 puts all its weight on position 0 — it literally has nothing else available. Row 1 splits weight across positions 0 and 1 only. The upper-right zeros aren't rounding; causal_mask forces them there exactly, before the softmax even runs.

Watch Out For

Masking after the softmax instead of before

Setting attention weights to zero after the softmax, rather than adding negative infinity before it, looks equivalent but isn't. Post-hoc zeroing leaves the softmax's denominator computed as though the blocked positions were still included, so the remaining weights don't actually sum to 1 anymore — the distribution is subtly wrong, and every output is a slightly miscalibrated blend. Masking before the softmax is what guarantees the surviving weights renormalize correctly.

Applying a causal mask to an encoder that doesn't need one

Causal masking is specifically for autoregressive generation — decoder-only models and decoder stacks in encoder-decoder architectures. An encoder that's building a representation of a complete, already-available input (translation source text, an image's patches) has no "future" to hide, and masking it anyway needlessly throws away context the encoder is supposed to have. Confusing which stack needs the mask is a common wiring mistake when assembling a custom transformer from scratch.

The Quick Version

Causal masking adds negative infinity to every attention score where the key position is later than the query position.
After the softmax, those blocked positions land at exactly zero weight — not approximately, exactly.
The surviving pattern is a lower triangle including the diagonal: every position attends to itself and everything before it, nothing after.
This is what prevents a next-token-prediction model from trivially copying the answer during training.
It also lets every position's training loss compute in one parallel forward pass, since no position can see ahead to cheat.