Multi-Head Self-Attention

The problem

Multi-head attention runs $H$ attentions side by side, so that each one can learn a different pattern (one head may follow grammar, another a name ten tokens back). It does not add parameters: it cuts each query, key and value vector of $D$ numbers into $H$ slices of $d = D / H$ numbers, and runs one attention per slice. The hard part is the shape work, which you will write here.

The input x has shape $(B, T, D)$: $B$ texts, $T$ tokens each, $D$ numbers per token. The four weight matrices Wq, Wk, Wv, Wo have shape $(D, D)$.

The five steps

1. Project: $Q = xW_q$, $K = xW_k$, $V = xW_v$, each of shape $(B, T, D)$.
2. Split heads: $(B, T, D) \to (B, T, H, d) \to (B, H, T, d)$. Head $h$ takes the numbers $hd$ to $(h+1)d - 1$ of every token.
3. Causal attention in every head at once, on the last two axes: $(B, H, T, d)$.
4. Merge heads: undo step 2, $(B, H, T, d) \to (B, T, H, d) \to (B, T, D)$.
5. Output projection: multiply by $W_o$, which mixes the heads.

Splitting, with numbers

Example. One text, $T = 2$ tokens, $D = 4$, $H = 2$ heads of $d = 2$. The projected queries are x = [[[0, 1, 2, 3], [4, 5, 6, 7]]]: token 0 is [0, 1, 2, 3], token 1 is [4, 5, 6, 7].

• Head 0 must get the first two numbers of every token: [0, 1] for token 0 and [4, 5] for token 1.
• Head 1 must get the last two: [2, 3] and [6, 7].

Right: x.reshape(1, 2, 2, 2) as $(B, T, H, d)$ cuts each token into its two halves, then .transpose(0, 2, 1, 3) gives $(B, H, T, d)$: head 0 is [[0, 1], [4, 5]], head 1 is [[2, 3], [6, 7]].

Wrong: reading the same reshape(1, 2, 2, 2) directly as $(B, H, T, d)$, with no transpose, makes head 0 [[0, 1], [2, 3]]: the two halves of token 0, as if they were two different tokens. It runs without any error.

reshape never moves numbers: it only reads them in memory order, last axis fastest. So split the last axis first, where the numbers of a head are already next to each other, and then move the head axis into place with transpose.

What to implement

• split_heads(x, n_heads): $(B, T, D) \to (B, H, T, d)$.
• merge_heads(x): $(B, H, T, d) \to (B, T, H \cdot d)$, the exact inverse.
• mha_forward(x, params, n_heads): the five steps, where params is a dict with "Wq", "Wk", "Wv", "Wo".

The attention and causal_mask functions of the previous exercise are given.

Requirements

• Raise ValueError if $D$ is not divisible by n_heads.
• No loops over heads or batch elements.
• merge_heads(split_heads(x, H)) must give back x exactly.