AI Grimoire

Many heads at once

The subspace argument for multiple heads is clean, standard, and only about half of what is going on. The other half is that most heads are doing nothing in particular.

By the end of part 2 we had one attention operation: a similarity, a softmax, a weighted sum. Part 3 looked at what its output matrix contains. This part is about why nobody uses one of them.

The averaging problem

Consider a single head reading the word bank in “she sat by the bank and counted the notes”. Two things want attention: sat by disambiguates the sense, and notes is the object the eventual prediction depends on.

One softmax row has to produce one distribution. It can split its mass between the two, and then the context vector is a blend — half of one signal and half of the other, added into the same dd dimensions with no way to tell them apart downstream.

c=αby vby+αnotes vnotes,αby+αnotes≤1c = \alpha_{\text{by}}\,v_{\text{by}} + \alpha_{\text{notes}}\,v_{\text{notes}}, \qquad \alpha_{\text{by}} + \alpha_{\text{notes}} \le 1
eq. 1 — one row, one budget

This is not a subtlety I appreciated for a long time. The constraint is not that attention cannot look at two places — it obviously can, and part 1 made a point of that. It is that both retrievals land in the same vector, so they interfere, and the network has no free parameters left with which to separate them.

The standard fix

Run hh attention operations in parallel, each with its own projections, each on a d/hd/h-dimensional slice, and concatenate.

MHA(x)=[head1 ∥ ⋯ ∥ headh]WO,headi=softmax⁡ ⁣(QiKi⊤dk)Vi\mathrm{MHA}(x) = \bigl[\mathrm{head}_1 \,\|\, \cdots \,\|\, \mathrm{head}_h\bigr] W^O, \qquad \mathrm{head}_i = \softmax\!\left(\frac{Q_iK_i^\top}{\sqrt{d_k}}\right)V_i
eq. 2 — h independent operations, one output

Each head has its own α\alpha, so each can attend somewhere different, and each writes into its own slice of the concatenation. The interference is gone because the outputs occupy disjoint dimensions until WOW^O mixes them deliberately.

The parameter count is unchanged. dk=d/hd_k = d/h, so hh heads of width d/hd/h is the same total projection size as one head of width dd — multi-head attention is free relative to single-head, which is the detail that makes the argument persuasive rather than merely appealing.

What they actually learn

Here the tidy story ends, and I think the tidy story is taught for too long.

Michel et al. pruned heads from trained models one at a time. On most layers, most heads could be removed with no measurable loss — many layers ran fine on a single head. Voita et al. found the same and characterised the survivors: positional heads that attend to a fixed offset, syntactic heads that follow dependency relations, and rare-token heads. A small minority doing something identifiable, and a majority that could be deleted.

That is not what “each head learns a different subspace” predicts. The prediction is a division of labour; the observation is a handful of specialists and a lot of redundancy.

The reframing that helped me

Elhage et al.’s circuits framework gave me the version I now use. A head is not a subspace specialist; it is a pair of maps.

WQ⊤WK⏟QK: where to lookWOWV⏟OV: what to write\underbrace{W_Q^\top W_K}_{\text{QK: where to look}} \qquad\qquad \underbrace{W_O W_V}_{\text{OV: what to write}}
eq. 3 — two circuits, and they are independent

QK decides the attention pattern; OV decides what gets copied into the residual stream when it attends. They compose only through the head, so a head is fully described by “which positions does it read, and what does reading them write”.

This is a better frame for two reasons. It explains why heads compose across layers — an induction head is a head whose QK circuit reads the output of a previous head’s OV circuit, which the subspace story has no vocabulary for. And it explains why so many heads look empty: a head with a near-zero OV circuit attends to things and writes nothing, which is exactly what a prunable head is.

The direction it went afterwards

Everything since has reduced the head count on the key and value side while keeping it on the query side. Multi-query attention gives all heads one shared K and V; grouped-query attention shares them among small groups.

The justification is KV-cache size, and it is entirely an inference argument — but it is worth noticing that it works at all. If heads genuinely needed independent subspaces, sharing the keys and values across eight of them should hurt badly. It does not, and the redundancy above is presumably why.

Next

The last part is about what attention still cannot do. Not efficiency — part 4 covered that — but the things the operation structurally cannot express, which turn out to explain several behaviours that look like bugs.

Series

Reference

The settled statements of what this note works through.

References

[1]Vaswani et al. — Attention Is All You Need (2017)arXiv:1706.03762
[2]Michel et al. — Are Sixteen Heads Really Better than One? (2019)arXiv:1905.10650
[3]Voita et al. — Analyzing Multi-Head Self-Attention (2019)arXiv:1905.09531
[4]Elhage et al. — A Mathematical Framework for Transformer Circuits (2021)transformer-circuits.pub

·