Exercise np-05 — scaled dot-product attention, a causal mask, and multi-head, in NumPy

What you're building

• scaled_dot_product_attention(Q, K, V, mask=None) -> (out, weights) — Q is (..., Tq, d_k), K is (..., Tk, d_k), V is (..., Tk, d_v); any number of leading batch/head dimensions, as long as Q, K and V broadcast together. Scores are Q @ K.swapaxes(-1, -2) / sqrt(d_k). mask, when given, is boolean with the same leading shape as the scores and True where a query may attend to a key; masked scores must become -inf before the softmax, not zeroed after it — the edge case that trap mask-after-softmax.py and mask-with-zero.py both get wrong, in two different ways. weights is softmax(scores, axis=-1); out = weights @ V.
• causal_mask(T) -> np.ndarray — a (T, T) boolean array, True at [i, j] exactly when j <= i (query i may see keys up to and including itself, never a key that comes later).
• multi_head_attention(x, Wq, Wk, Wv, Wo, heads, causal=False) -> np.ndarray — x is (B, T, D); Wq, Wk, Wv, Wo are all (D, D). Project x through each of Wq/Wk/Wv, split the last axis into heads pieces of D // heads, run scaled_dot_product_attention per head (with causal_mask(T) applied to every head identically when causal=True), concatenate the heads back into (B, T, D), and project once more through Wo. The edge case is the reshape: x.reshape(B, T, heads, dh) lines the last axis up into heads groups, but the head axis now sits after the token axis — you must transpose it next to the batch axis before running attention, or every head ends up attending across token boundaries instead of within its own slice of the embedding.

Four of the eight checks are worth reading before you start. Check 2 is a hand-picked 1×2 case — Q=[[1,0]] against two orthogonal keys — where scaled and unscaled attention give different first weights (0.6698 vs 0.7311): a missing / sqrt(d_k) shows up there and nowhere else, because softmax always sums to 1 whether or not the scores were scaled first. Check 3 masks 3 of 6 keys and demands the masked weights be exactly 0.0, to 1e-9 — scores * mask (zero, not -inf) leaves exp(0) in the sum, a small but nonzero leak. Check 4 changes the keys and values of every token after a fixed position t and demands out[t] not move by more than 1e-12 — softmax's denominator sums over every key, so masking the weights after computing that sum still lets a changed future key shift the surviving weights, which is exactly how both masking traps get caught here too. Check 8 builds block-diagonal Wq/Wk/Wv so head 1 can only ever see input dimensions 0–1 and head 2 only 2–3, then zeroes head 2's input slice and demands head 1's output slice not change at all — the check the forgotten transpose fails, because without it every "head" is actually reading a mix of tokens rather than a clean slice of the embedding. Every number here was measured on the reference solution with NumPy 2.1 under CPython 3.13, and none of them depends on the version — they are arithmetic, not floating-point luck.

If you get stuck

• the attention formula — scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(Q.shape[-1]); with a mask, scores = np.where(mask, scores, -np.inf) before subtracting the row max and exponentiating — subtract the max first (scores - scores.max(axis=-1, keepdims=True)) so an all--inf-except-one row doesn't overflow. weights = exp / exp.sum(axis=-1, keepdims=True), out = weights @ V. The leading dimensions of Q, K, V (batch, heads, whatever) just come along for the ride — @ and swapaxes(-1, -2) both work on the last two axes regardless of what precedes them.
• why -inf before, never 0 after — softmax is defined over whatever scores it is given. Feed it -inf for a position and exp(-inf) = 0 drops that term out of the sum before normalising, so the survivors are still a real distribution over only the allowed keys. Zero a score instead of -inf-ing it and exp(0) = 1 stays in the sum — a "masked" key still gets real weight. Compute the softmax first and zero the resulting weights afterwards, and the row no longer sums to 1 at all, because you threw away probability mass softmax already spent on the masked positions instead of redistributing it.
• causal_mask — j = np.arange(T)[None, :]; i = np.arange(T)[:, None]; return j <= i. Row i (query token i) is True for every column up to and including i.
• multi_head_attention, the reshape — D dimensions become heads groups of D // heads: Q.reshape(B, T, heads, dh) is correct as a first step, but its axis order is (batch, token, head, dim) — attention needs (batch, head, token, dim) so that @ contracts over dim per head, per token, not across tokens. .transpose(0, 2, 1, 3) fixes the order; skip it and scaled_dot_product_attention silently treats the head axis as if it were more tokens, mixing every head's Q against every other head's K. After attention, .transpose(0, 2, 1, 3).reshape(B, T, D) puts the heads back together before the final @ Wo.
• the causal mask inside multi-head — build it once, causal_mask(T), then add two leading size-1 axes so it broadcasts against (B, heads, T, T) scores: mask[None, None, :, :]. Every head and every item in the batch shares the exact same triangle — masking is about token position, not about anything a head or a batch item learns.