The causal mask

A decoder predicts each token from the ones before it, so token $i$ may attend to tokens $j \le i$ only. During training all positions are computed at once, so the future has to be hidden explicitly.

0 on and below the diagonal, $-\infty$ above: the lower-triangular shape from Part 2. It’s added to the scores, so allowed scores pass unchanged.

“the” keeps only its own score; “cat” keeps two; “sat” keeps all three.

$-\infty$ becomes exactly 0. Row “cat”: $e^{0} = 1$ and $e^{2.83} \approx 16.9$, so the weights are $1/17.9 \approx 0.056$ and $0.944$. Every row sums to 1. The shading shows where each query looks.

Which weight in this row is exactly 0? Which of the other two is bigger, and roughly how much bigger?

The masked one (key “sat”) is exactly 0. “cat” beats “the” by a factor of $e^{2.83} \approx 17$: a gap of 2.83 in score.

Zeroing after the softmax leaves a row that no longer sums to 1: “the” would pass on only a third of its value. Adding $-\infty$ before the softmax makes the remaining weights renormalize among themselves.

mask = torch.ones(T, T, dtype=torch.bool).tril()
scores = scores.masked_fill(~mask, float("-inf"))
weights = scores.softmax(dim=-1)
# or all at once:
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

• Mask: 0 on and below the diagonal, $-\infty$ above, added to the scaled scores.
• Softmax each row: future weights become exactly 0, the rest sum to 1.
• Example rows of $A$: $[1, 0, 0]$, $[0.056, 0.944, 0]$, $[0.14, 0.58, 0.28]$.
• Mask before the softmax, not after.