NLP Foundations
Multi-Head Attention in Depth
Why attention runs as several small parallel attentions rather than one large one, what each head actually gets to specialise in, and why most heads turn out to be redundant.
intermediate · 9 min read
A transformer with model dimension 4096 and 32 attention heads does not run one attention operation on 4096-dimensional vectors. It runs 32 separate attention operations, each on 128-dimensional slices, in parallel, then concatenates the results back together. This split is not a minor implementation detail; Vaswani et al., 2017 found empirically that many small attentions beat one large one, and the reason why says something real about what a single attention operation can and cannot represent.
The mechanics of the split
Given model dimension d_model and h heads, each head gets d_k = d_model / h dimensions for its queries and keys, and typically the same for values. Each head has its own W_Q, W_K, W_V (see queries-keys-values), computes attention independently over the full sequence, and produces an output of size d_k per position. The h outputs are concatenated back to d_model and passed through one more learned projection, W_O:
head_i = Attention(x @ W_Q_i, x @ W_K_i, x @ W_V_i)
MultiHead(x) = concat(head_1, ..., head_h) @ W_O
Every head sees the entire sequence; nothing restricts which positions a given head can look at. What differs across heads is only the learned projection, so each head is free to develop a different matching criterion over the same input.
Why one big head is not equivalent
A single attention operation, however large d_k, computes exactly one weighted average per position: one softmax distribution over the sequence, applied uniformly to build one output vector. If a token needs to attend to its immediate neighbour for syntax and simultaneously to a distant coreferent for meaning, a single softmax has to blend both needs into one compromise weighting. Splitting into heads gives the model several independent softmax distributions per position, computed in parallel and then merged only after the fact by W_O. This is a genuine increase in expressive capacity, not just a computational rearrangement: one 512-dimensional head and eight 64-dimensional heads have the same parameter count in the projections, but only the eight-head version can represent eight simultaneously different attention patterns at a single position.
What heads actually learn
Interpretability work on trained transformers finds heads that specialise in identifiable ways: heads that mostly attend to the previous token (positional), heads that track subject-verb agreement across a clause (syntactic), heads that resolve pronouns to antecedents (coreference), and, in models trained with enough scale and repetition, induction heads that implement a simple copy-the-pattern algorithm. No one designs heads to specialise this way; it falls out of training because a head that finds one consistent, useful matching criterion reduces loss more reliably than a head that tries to do everything.
Redundancy is the norm, not the exception
Michel, Levy, and Neubig, 2019 pruned trained transformers head by head and found that most heads could be removed at inference time with little to no loss in accuracy; often a single head per layer carried most of the useful signal, and in some layers nearly all heads could be cut. This is uncomfortable for a clean "each head learns something distinct" story, but it is consistent with how overparameterised networks train generally: the architecture provides capacity for many independent patterns, training uses some of it well, and a lot of the rest is redundant or only marginally useful. Head pruning and head importance scoring are now standard tools for inference-time compression precisely because this redundancy is so consistent across models.
When it falls down
- More heads is not free capacity. Splitting
d_modelacross more heads shrinksd_kper head, and a head with too few dimensions may lack the room to encode a useful matching criterion at all. Head count and head dimension trade off against each other; doubling heads while halvingd_kis not a strict improvement. - Specialisation is not guaranteed. The clean "positional head, syntax head, coreference head" story comes from studying specific trained models on specific tasks. Many heads in practice are diffuse, redundant, or contribute in ways that resist a tidy label.
W_Ore-mixes everything. The final output projection combines all heads' outputs back into one vector, so a head's individual contribution to the model's behaviour cannot be read off in isolation fromW_Oand everything downstream; head-level interpretability claims need to account for this.- Cost scales with heads at inference. Every head needs its own key and value vectors cached for autoregressive decoding, which is exactly the memory pressure that multi-query-grouped-query-attention exists to relieve.
Further reading
- Vaswani et al., 2017, Attention Is All You Need, arXiv:1706.03762 - the original multi-head formulation and the ablation showing multiple heads beat one.
- Michel, Levy, and Neubig, 2019, Are Sixteen Heads Really Better than One?, arXiv:1905.10650 - the head-pruning study showing widespread redundancy.
- Voita et al., 2019, Analyzing Multi-Head Self-Attention, arXiv:1905.09418 - identifies specialised head roles (positional, syntactic) and confirms most heads can be pruned.