Multi-head attention

Instead of one wide head, use $h$ narrow ones, each with its own projections of width $d_h = d/h$. Each head computes its own $T \times T$ attention pattern: one might follow the previous token, another the subject of the sentence.

Code keeps one $d \times d$ weight per projection. Its column blocks are the heads (Part 2), so a single matmul computes all heads’ queries, and view + transpose cut the result into heads (Part 4).

The heads’ outputs are placed side by side, back to width $d$, and one more $d \times d$ matrix $W_O$ mixes information across heads. In block form that’s $\sum_i \text{head}_i\, W_O^{(i)}$: every head writes its own contribution into the residual stream.

Splitting into heads doesn’t add parameters: $h$ heads of width $d/h$ use the same $d \times d$ matrices as one head of width $d$. That’s about 67 million parameters per layer for standard multi-head attention at $d = 4096$, without biases.

From Part 3: each head compares tokens through a $d \times d$ bilinear form of rank at most $d_h$. Many narrow heads see many different low-dimensional “views” of the tokens, instead of one wide view.

Write $d_h$ and the number of attention weight parameters per layer (ignore biases).

About 2.4 million per layer, 28 million across GPT-2 small’s 12 layers.

q = (x @ W_q).view(B, T, h, d_h).transpose(1, 2)   # (B, h, T, d_h)
k = (x @ W_k).view(B, T, h, d_h).transpose(1, 2)
v = (x @ W_v).view(B, T, h, d_h).transpose(1, 2)
o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
o = o.transpose(1, 2).reshape(B, T, d) @ W_o       # concat heads, mix

• $h$ heads of width $d_h = d/h$, each with its own attention pattern.
• One big projection, cut into heads by view + transpose.
• Concatenate and mix: $\operatorname{concat}(\text{heads})\,W_O$.
• $4d^2$ parameters regardless of $h$; each head has rank $\le d_h$.