31 · Multi-head attention
Run h attention 'heads' in parallel on projected subspaces, then concatenate and project — each head learns a different kind of relationship.
Multi-head attention runs h attentions in parallel on lower-dimensional subspaces, so the model attends to several different relationships at the SAME total compute as one head.
Without this:
A single attention head can only capture one type of relationship at a time; multi-head is what gives transformers their expressive power.
In the last lesson we built self-attention: every position computes Attention(Q,K,V) = softmax(QKᵀ/√d_k)·V, mixing information across the whole sequence in one parallel step. But a single attention head has a limitation. Each head produces ONE softmax distribution per query position — so it can latch onto exactly ONE kind of relationship. Maybe it learns to align verbs with their subjects. But what about the relationship between a pronoun and its antecedent, or a token and the one immediately to its left? One head cannot do all of these at once: its single softmax must compromise.
Multi-head attention is the fix. Instead of running one attention over the full d_model-dimensional vectors, we run h heads in parallel, each over a smaller subspace of size d_k = d_v = d_model / h. Each head gets its OWN learned projections W_i^Q, W_i^K, W_i^V, so each head sees a different "view" of the input and can specialize in a different relationship. We then concatenate the h head outputs and pass them through one final output projection W^O.
Formally, with d_model the model width and h the number of heads:
head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)MultiHead(Q,K,V) = Concat(head_1, …, head_h) · W^O
where each W_i^Q, W_i^K ∈ ℝ^{d_model × d_k}, W_i^V ∈ ℝ^{d_model × d_v}, and W^O ∈ ℝ^{h·d_v × d_model}. Recall the core operation each head runs is exactly the scaled dot-product from lesson 30: Attention(Q,K,V) = softmax(QKᵀ/√d_k)·V. Multi-head simply does h of these on thinner slices and stitches the results back together.
The clever part: because d_k = d_model / h, the total amount of computation across all h heads is the SAME as one full-width attention. We get the expressiveness of many specialized heads for the price of one.
Multi-head shapes: (seq, d_model) → split into (h, seq, d_k) → per-head attention → concat → W_O
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Hover a row (target token) to highlight its attention distribution. Switch heads above.
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
PyTorch read-along: nn.MultiheadAttention with embed_dim, num_heads, and need_weights to retrieve attention maps
Why split attention into multiple heads instead of using one large full-width attention?
- Multi-head attention runs h scaled-dot-product attentions in parallel on d_k = d_model/h subspaces, each with its own W_i^Q, W_i^K, W_i^V, then concatenates and applies W^O.
- The parameter count is 4·d_model² (the four projections), independent of h — adding heads is a reshape, not new weights — and total compute matches a single full-width head.
- Heads empirically specialize (positional, syntactic, induction heads), but more heads is not strictly better: d_k shrinks as h grows, so there is a tunable sweet spot.
Multi-head attention is the core mixing operator in every transformer: BERT, GPT, T5, ViT, Whisper. When you read 'GPT-3 has 96 heads', that's the h in this lesson.
If you remove it: You'd read about 'attention heads' and head specialization in interpretability papers without understanding mechanically how the d_model width is sliced, projected, and recombined.