Multi-Head Attention Primer

A query's output is a weighted average of value vectors, and the weights come from one softmax. That is not a detail of the implementation — it is the whole shape of what a head can produce: one mixture, however the weight gets split among the six keys.

it is still the token to watch. Put a head on it that only ever leans on whichever token came right before, and drag its query anywhere in the sentence:

Notice that this head has no opinion about grammar at all — it does not know it is a pronoun, only that and sits one position back. That is a real, useful signal. It is also not the only one it needs.

Here is a second head, built to lean on whatever the query depends on grammatically instead. Watch both rows as the query moves — at position 4, the two heads point at different tokens entirely:

Because it is the subject of ran, the second head reaches forward to it — the same one-token distance as the first head, but the opposite direction: the first head's rule can only ever look backward, never forward. One row, two genuinely useful, genuinely different answers.

A single head has no third option: it has to blend. Slide the mixture from pure position to pure grammar and read the peak as it crosses the middle:

Watch the peak weight fall as the blend passes through the centre, then climb again toward the other pure answer — the row is never confidently near either token in between. At an even split, it is confidently near neither.

That dip is not a quirk of this sentence. Plotted against the blend itself, the head's best weight on either signal is lowest exactly where it is trying hardest to cover both:

A compromise is not a third answer — it is a worse copy of the two answers it was built from. The fix on the next page is not to make one head smarter. It is to stop asking one head to do two jobs.

A token still arrives as one vector, d_model numbers wide — 8 in the drawings below, 768 in GPT-2 small. Multi-head attention does not add to that width. It cuts it into h equal pieces of d_k = d_model / h each.

Drag to change h and watch the same width get sliced differently — at h = 1 there is no cut at all, one head, the whole vector:

Watch each slice shrink every time a cut is added, while the outer edges never move. Four heads of width 2 are, dimension for dimension, the same budget as one head of width 8 — just a different arrangement.

Here is the mistake this invites: it is tempting to picture each head reading its own quarter of the token. Toggle between the two pictures and watch what each head actually sees at its input, not its output:

Because every head's projection is its own full d_model × d_k matrix, each head reads all eight input numbers and keeps a narrow result — nothing about the token itself is ever chopped up.

What gets cut is W_Q's own output columns — and W_K's, and W_V's, the same way, independently. Change h once more and watch all three split together:

Query, key and value each keep their own d_model × d_model matrix and their own cut — nothing is shared between them but the number h. Section 3 is what each of those narrow slices does once it is on its own.

A head with a d_k = 2 query and key is not a shortcut version of attention — it is attention, full stop, at a narrower width. Nothing about scoring, softmax or the weighted sum changes because the head is one of several.

Drag the query and watch the one that tracks position and the one that tracks grammar answer at full strength, simultaneously, with nothing standing between them and their own row:

Watch that neither row ever softens toward the other. The compromise from Section 1 is gone — not fixed, replaced. Two independent mixtures cost the same as the one blended row did.

Nothing caps this at two. Add the two heads §1 never got to try — one that always pulls toward the same token whatever the query is, and a fourth that commits to neither:

Three of the four rows have a sharp, nameable story. The fourth does not, and it gets no colour here for exactly that reason — Section 4 is what all four turn out to mean in a real, trained model.

One more place independence bites: each head divides by its own √d_k, not the model's. Copy the constant from a single-head implementation and every head gets quietly flatter as d_k shrinks:

Because the correct divisor cancels the width out, the peak weight for the same underlying signal holds flat across every d_k. The wrong one lets that signal wash out — no error, no crash, just a softmax that never gets to be confident.

Voita, Talbot, Moiseev, Sennrich and Titov trained a translator and looked at what its heads actually attend to. Most heads fall into one of three roles — the two already on this page, plus a third: a pull toward whichever token is rarest in the sentence, no matter what the query is.

Put the head that tracks position beside the one that always reaches for the rarest token, and drag the query across the sentence:

Watch the top row's peak walk one step behind the query, every time — and the bottom row never move at all. Same six keys, same control, two structurally different rules.

Plotted for every query position at once, that difference is a shape: the positional head traces a straight diagonal, the rare-token head a flat line pinned on cat:

The syntactic line is neither — it jumps by whatever the dependency actually is at each position, forward one at it, back one at ate. Step through it on its own and watch the target refuse to settle into a rule:

Because that head is answering "who does this depend on", not "how far back", its target moves however the grammar does. It is the one signal §1's single row could not fake with position alone.

The payoff: heads with a story this clean are the ones a model cannot afford to lose. Prune all but ten of the 48 heads in Voita et al.'s translator and the specialised ones survive:

The heads left standing are overwhelmingly positional, syntactic or rare-token ones — the fourth, uninterpretable kind from §3 is exactly what gets cut first. A head that specialises earns its width; one that does not is close to free to remove.

The obvious move is the real one: lay the four outputs side by side. Place all four and the width comes back to exactly where §2 cut it:

Watch the last empty slot fill and nothing else move — concatenation does not touch a number that is already placed. It only ever adds the next head's two values onto the end.

These are not placeholders. For the token it, each head's output is a real weighted sum over the same six value vectors, using that head's own weights. Scrub to a different token and every one of the eight numbers is real:

Four heads, four honest answers, laid end to end — and, so far, still four separate opinions. Nothing has read head three's number while deciding head one's.

That is what W_O is for. It is one more d_model × d_model matrix, applied to the whole concatenation at once. Turn its mixing up from nothing and watch a single slot start to carry a share of every head's answer:

Because W_O is not block-diagonal, the local head's own slot ends up as a blend of all four — the one place in the whole mechanism where a syntactic finding and a positional one are allowed to influence the same output number.

Four matrices have now done the work: W_Q, W_K, W_V split the token into heads, and W_O puts the heads back together. All four are the same shape:

That shared shape is not a coincidence — it is next section's whole argument. Four matrices at one width, however many heads split that width between them.

W_Q, W_K, W_V and W_O are each d_model × d_model whatever h is — §5 already drew all four the same size. Drag h here and watch the parameter count that formula gives:

Every cut changes d_k and never touches that number. Four matrices at one fixed width is a parameter count with no h in it anywhere — not approximately constant, exactly constant.

The attention matmuls scale with d, and h heads of width d_model / h sum to d_model exactly, so their total is h × d_model / h multiply-adds — h cancels. Drag along h and watch the curve refuse to move off the reference line drawn under it:

One wide head and four narrow ones trace the identical curve at every context length — GPT-2 small's own numbers, not a rounded illustration. The four projection matrices cost about 1.5× that again, and they do not move either: nothing in this section's cost depends on how the width is split.

None of that makes h free. Push it up and d_k keeps shrinking — the same slider, read for what a head can still represent rather than what it costs:

At d_k = 1 a query and a key are single numbers, and two numbers of the same sign always point the "same way" — a head that narrow has only two directions left to tell apart, whatever it is trying to attend to. Nothing errors. The FLOPs and the parameter count you just watched hold perfectly steady. The head simply stops being able to say much, silently, which is why shipped models keep d_k around 64–128 and grow h alongside d_model instead of past it.

q = (x @ Wq).view(n, h, dk).transpose(0, 1)
k = (x @ Wk).view(n, h, dk).transpose(0, 1)
v = (x @ Wv).view(n, h, dk).transpose(0, 1)
s = q @ k.transpose(-2, -1) / dk**0.5
w = s.softmax(dim=-1)
out = (w @ v).transpose(0, 1).reshape(n, h * dk)
out = out @ Wo

view and transpose are the whole split — nothing about x changes, only the shape the three projections' own output is read in. dk**0.5, not d_model**0.5, is §3's trip-up. The final reshape then @ Wo is §5's concatenation and mix. h appears exactly four times, and cost appears in neither line — the same seven lines run, unmodified, from Transformer-base's eight heads to GPT-3's ninety-six — every width a shipped model has picked:

d_k moves far less than h does: the ratio that grew the least across three generations is representational headroom per head, not the model's width. Splitting one wide head into several narrow ones was never about affording more compute — it was about affording more than one answer per token, at no extra price. What those heads still cannot tell each other is which token came first, which is where the next page starts.