Scaled Dot-Product Attention with a Causal Mask

The problem

Attention lets every token look at the other tokens and take a weighted mix of their content (see the theory page). Each token has a query (what it looks for), a key (what it offers) and a value (what it passes on). You will write attention itself, and the causal mask that stops a token from looking at the tokens after it.

Attention in three steps

where:

• $Q$, $K$ and $V$ stack the queries, keys and values as rows, and $d$ is the size of a query or key;
• $QK^\top$ holds every query-key dot product: entry $(t, s)$ is how well query $t$ matches key $s$. Divided by $\sqrt{d}$, these are the scores $S$;
• the softmax runs along each row, so the weights of one query add up to 1;
• multiplying by $V$ gives each query the weighted average of the values.

So:

1. scores: query times key, divided by $\sqrt{d}$;
2. weights: softmax of each row of scores;
3. output: weights times values.

The causal mask

When predicting the token at position $t$, the model may only use positions $0, \dots, t$: the later tokens are what it is trying to predict. The mask is a table of True (allowed) and False (forbidden), with True where the key position $s$ is at most the query position $t$: a lower triangle. Before the softmax, every forbidden score is replaced by $-\infty$, so its weight becomes $e^{-\infty} = 0$.

Example. Three tokens, $d = 2$. Queries and keys are the same here: $(1, 0)$, $(0, 1)$, $(1, 1)$. Values: $(1, 0)$, $(0, 10)$, $(5, 5)$.

• causal_mask(3) is [[T, F, F], [T, T, F], [T, T, T]] (T for True, F for False).
• Scores, masked: row 0 is $(0.71, -\infty, -\infty)$, row 1 is $(0, 0.71, -\infty)$, row 2 is $(0.71, 0.71, 1.41)$.
• Weights: row 0 is $(1, 0, 0)$, row 1 is $(0.33, 0.67, 0)$, row 2 is $(0.25, 0.25, 0.50)$. Each row adds up to 1.
• Output: token 0 gets exactly its own value $(1, 0)$; token 1 gets $0.33 \cdot (1, 0) + 0.67 \cdot (0, 10) \approx (0.33, 6.70)$; token 2 gets about $(2.77, 5)$.

Batches and heads

In a real model, q, k and v have extra leading axes, for the texts of the batch and the heads: shape $(B, H, T, d)$. Attention runs independently for every batch and head, on the last two axes. NumPy's @ already works that way: on arrays with more than two axes, it multiplies the last two axes and treats the others as a batch.

What to implement

• causal_mask(T): a boolean $(T, T)$ array, True where attention is allowed ($s \le t$).
• attention(q, k, v, mask=None): q has shape $(\dots, T, d)$, k has shape $(\dots, S, d)$ and v has shape $(\dots, S, d_v)$, with the same leading axes. mask, if given, is a boolean array that broadcasts to $(\dots, T, S)$. Return (out, weights): the $(\dots, T, d_v)$ output and the $(\dots, T, S)$ weights.

Requirements

• Masked scores become $-\infty$ before the softmax, so their weights are exactly 0. Every row of weights adds up to 1.
• Stable softmax: subtract the row maximum.
• No Python loops.
• You can assume every row has at least one allowed position.

Hints

• To transpose only the last two axes of k (whatever the leading axes), use np.swapaxes(k, -1, -2). Plain k.T reverses all the axes.
• np.where(mask, scores, -np.inf) applies the mask; np.tril builds a lower triangle.