Multi-head attention

Key idea

Split d_model into h heads of size d_model ÷ h, run attention in each, concatenate, multiply by W_O

Real world

A newspaper story read by an editorial team at the same time: one checks names, one checks dates, one checks quotes. Each marks up the same text for their own purpose, and the editor combines their notes.

In an LLM

GPT-2 small has 12 heads of 64 numbers in every layer. In layer 4 of GPT-2 reading “The cat sat on the mat because it was tired”, one head reads the previous word, one links “it” to “cat”, and most of the others park their attention on the first token.

One head, one pattern

A single attention head produces one weight matrix: for every token, one set of weights over the earlier tokens. But a word usually needs several kinds of information at once. For “it” in our sentence, the model would like to know what “it” refers to (“cat”), what the previous word was (“because”), and perhaps where the sentence started. A single softmax has to share its weight among all of those, and blurs them.

The fix is simple: run several heads in parallel, each with its own $W_Q$, $W_K$ and $W_V$, and let each one learn its own pattern. That is Multi-head attention. Each Attention head is a full attention computation; a layer just has several.

Split, attend, concatenate, mix

The heads don’t each use the full vector. If the model has width $d_{\text{model}}$ and $h$ heads, each head works with $d_{\text{model}} / h$ numbers. In GPT-2 small that is $768 / 12 = 64$.

1. Project. Compute $Q$, $K$ and $V$ for all heads at once, with full-width matrices: each is $T \times d_{\text{model}}$.
2. Split. Cut the columns into $h$ groups. Head 1 gets the first $d_{\text{model}}/h$ columns, head 2 the next, and so on. No arithmetic: it is a reshape from $(T, d_{\text{model}})$ to $(T, h, d_{\text{head}})$, then a transpose to $(h, T, d_{\text{head}})$ so each head is its own slice.
3. Attend. Every head runs ordinary scaled, masked attention on its slice. All heads are computed together as one batched matrix multiplication.
4. Concatenate. Put the head outputs back side by side: $(T, d_{\text{model}})$ again.
5. Mix. Multiply by an output matrix $W_O$ ($d_{\text{model}} \times d_{\text{model}}$), so each output number can combine what every head found.

Here it is on our toy sentence with $d_{\text{model}} = 4$ and two heads of size 2:

The two heads see the same tokens but have different weights, so they produce different patterns. Head 1 sends the last “the” mostly to “cat”; head 2 sends every word back to “The”. The output of the layer contains both kinds of information.

Real heads in GPT-2

The toy heads were designed by hand. In a trained model each head’s pattern is learned. Here are all 12 heads of a GPT-2 layer, reading the same sentence. (Layers and heads are numbered from 0, so “layer 4” is the fifth of the 12.)

In layer 4, head 11 is a previous-token head: every token gives over 99% of its attention to the token just before it. Head 3 sends “it” to “cat”. Head 7 gives its largest share, about half, to the current token itself. Most of the rest put the bulk of their weight on the first token, a habit called an attention sink, explained on the next page. Move the layer slider and the mix of patterns changes from layer to layer.

How many parameters

Splitting into heads doesn’t add weights. Each head’s $W_Q$, $W_K$ and $W_V$ are $d_{\text{model}} \times d_{\text{head}}$; with $h$ heads that adds up to three $d_{\text{model}} \times d_{\text{model}}$ matrices, the same as one big head. Add $W_O$ and an attention layer has

whatever $h$ is. For GPT-2 small: $4 \times 768^2 = 2{,}359{,}296$, plus $3{,}072$ biases, about 2.36 million per layer and 28.3 million across its 12 layers, just under a quarter of its 124 million parameters.

So the number of heads is a trade-off, not a cost. More heads give more separate patterns, but each compares queries and keys using fewer numbers. Most models settle on heads of 64 or 128.

Why it works

Different relations need different query–key matches. “Which noun does this pronoun mean?” compares meaning; “what was the previous token?” compares positions; “where did this word appear before?” compares token identity. One head’s single dot product can’t be good at all of them at once, because its weights would have to be a compromise. Separate heads can each specialise, and $W_O$ lets the next part of the network use all their findings together.

Not every head is essential. Researchers have pruned many heads from trained models with little loss, which suggests heads are partly redundant. A few, like the induction heads on the next page, turn out to be crucial.

B, T, D = x.shape
q, k, v = (x @ W_qkv).split(D, dim=-1)                 # each (B, T, D)
q = q.view(B, T, h, D // h).transpose(1, 2)            # (B, h, T, d_head)
k = k.view(B, T, h, D // h).transpose(1, 2)
v = v.view(B, T, h, D // h).transpose(1, 2)
o = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # (B, h, T, d_head)
o = o.transpose(1, 2).reshape(B, T, D)                 # concatenate heads
y = o @ W_o