One attention head can only look one way at a time. That single limitation is the whole reason multi-head attention exists, and once it clicks the rest of the mechanism is almost boring.
Start with the atom. Scaled dot-product attention takes queries, keys and values. Every query is scored against every key with a dot product, the scores are divided by the square root of the key dimension, a softmax turns each row into weights that sum to one, and the output is that weighted average of the values:
attention(Q, K, V) = softmax(Q Kᵀ / √dk) · V
The dot product answers "how much should this token listen to that token", the softmax makes it a proper distribution, and the weighted sum blends the values the token chose to attend to. That is it. A whole Transformer is this one operation stacked and wrapped.
Two details people skip past. First, Q, K and V are not the raw embeddings; each token is projected by three separate learned matrices. The query is what a token is looking for, the key is what it advertises, the value is what it hands over if attended to. Giving them different matrices is what lets a token ask for something different from what it offers. Second, the √dk. A dot product of dk numbers has variance roughly dk, so as the head gets wider the raw scores blow up, softmax saturates, almost all weight lands on one key, and the gradient through the rest collapses. Dividing by √dk pulls the scores back to unit variance so softmax stays in its responsive range. It is a tiny constant that quietly decides whether a deep model trains at all.
Now the actual idea. A single head is one lens: one dot-product space, one softmax, so it can emphasise exactly one kind of relationship. But a sentence has many relations happening at once. A word relates to its neighbour, to its syntactic head, to an earlier mention of the same thing, to sentence boundaries. Force all of that through one softmax and the loudest signal wins while everything subtler gets averaged into mush.
Multi-head attention fixes this by running h attention functions in parallel, each in its own slice of the representation. And the beautiful part: it costs almost nothing extra. After projecting to width d you split that width into h contiguous chunks of size dk = d/h. Each chunk is one head's private subspace. In tensor terms you reshape (batch, seq, d) into (batch, seq, h, dk) and move the head axis up front. No new parameters — it is a reshape. The projection matrices are still d×d whether you use one head or twelve, so the parameter count stays at 4d² (the three input projections plus the output one) and the attention FLOPs barely move, because h heads of width dk still sum to d.
Each head then runs its own scaled dot-product attention with its own independent softmax. That independence is the whole point — it is what lets the heads diverge during training. Then you concatenate the h outputs back into one width-d vector and pass it through a final output matrix W_O. Skipping W_O is a common bug: without it the heads sit in fixed disjoint slots and never mix. W_O is what lets information flow across heads.
I built an interactive version on the sentence "the cat sat on the warm mat ." with four heads, and the heatmaps come out visibly different from the same input: one head forms a diagonal band (attend to neighbours), one links same-part-of-speech tokens (the two determiners pair up, the two nouns pair up), one links same-word tokens (the two copies of "the" lock together but cat and mat do not), and one dumps almost all its weight on the final "." — a delimiter sink. Drop to a single head and those four patterns collapse into one blurry map dominated by the strongest signal. You can flip between 1, 2 and 4 heads and watch it happen, and click any row to see the real q·k / √dk → softmax arithmetic:
https://dev48v.infy.uk/dl/day30-multi-head-attention.html
This is not just a cartoon. When people probe trained models they find heads that really do specialise this way: positional heads, verb-to-subject heads, coreference heads, and plenty of delimiter/sink heads. Many are redundant enough to prune after training, which is why head count is a forgiving hyperparameter.
The modern twist is about inference, not quality. Standard multi-head attention stores a separate K and V per head, and during long-context decoding that KV-cache dominates memory. Multi-Query Attention shares one K/V across all query heads; Grouped-Query Attention shares K/V within small groups. You keep many query heads but shrink the cache dramatically, which is why Llama-2 70B, Mistral and most current LLMs use GQA.
In code it is six lines: project, split, attend per head, concat, project, done — or just call nn.MultiheadAttention. But building it once, and watching the heads specialise, is what makes the six lines stop being magic.
Top comments (0)