Causal mask: build + apply

Background

GPT-2 is a decoder: at training time it predicts token $t+1$ from tokens $0..t$, so position $t$ must not attend to any position $> t$ — otherwise the model "cheats" by reading the answer straight out of its input. The fix is a causal (lower-triangular) mask applied to the attention scores before softmax. It is one triangular matrix and one where call, but getting the placement right is what makes future positions receive exactly zero probability.

Problem statement

Implement two functions:

1. build_causal_mask(T) — return a (T, T) boolean mask with mask[i, j] = True when $j \le i$ (position $i$ may attend to position $j$) and False when $j > i$. This is the lower triangle including the diagonal.
2. apply_causal_mask(scores) — given a (T, T) score matrix (from $QK^\top/\sqrt{d}$), return a copy with the upper triangle ($j > i$) set to $-\infty$ and the rest unchanged, so a subsequent softmax sends those entries to exactly 0.

Input

• build_causal_mask(T) — T: int, the sequence length.
• apply_causal_mask(scores) — scores: np.ndarray of shape (T, T), the raw attention scores.

Output

• build_causal_mask returns a (T, T) bool array (lower-triangular True).
• apply_causal_mask returns a (T, T) float array: original scores at/below the diagonal, $-\infty$ above it.

Examples

Example 1 — the mask for T = 4

build_causal_mask(4) =
  [[ True, False, False, False],
   [ True,  True, False, False],
   [ True,  True,  True, False],
   [ True,  True,  True,  True]]

Explanation: row $i$ has True in columns $0..i$ (it can see itself and the past) and False afterward — so row $i$ has exactly $i+1$ True entries.

Example 2 — applying it, then softmax

Input:  scores = np.full((4, 4), 1.0)
apply_causal_mask(scores) =
  [[ 1, -inf, -inf, -inf],
   [ 1,    1, -inf, -inf],
   [ 1,    1,    1, -inf],
   [ 1,    1,    1,    1]]
row 0 after softmax = [1, 0, 0, 0]      # future positions are exactly 0

Explanation: the upper triangle is replaced with $-\infty$ while the lower triangle keeps its scores; after a row-wise softmax, the $-\infty$ entries become exactly $0$ and each row's surviving weights sum to 1 (row 0 attends only to position 0).

Constraints

• mask[i, j] = True iff $j \le i$; row $i$ contains exactly $i + 1$ True values.
• Mask before softmax with $-\infty$ (or a large negative number) so masked probabilities are exactly 0 — not $0.001$.
• Do not zero weights after softmax: that leaves each row summing to less than 1 and corrupts the gradient.
• apply_causal_mask leaves the at/below-diagonal scores unchanged.

Notes

• -inf vs -1e9. Both work; -1e9 is friendlier to fp16, while -np.inf is conceptually cleaner (exp(-inf) = 0 exactly). np.tril(np.ones((T, T), bool)) builds the mask in one call.
• Series. apply_causal_mask plugs into build-gpt-06-multi-head-attention; this is the masking half of the attention stack alongside build-gpt-03.