Multi-Head Attention

Each head ends up carrying a different relationship

Look closely at what a sentence's heads weighed heavily and fairly distinct personalities show up. One head might consistently lean on the word right before it, catching the flow of speech. Another might link a subject to a verb far apart from it. Another might pile weight onto the start of the sentence or onto punctuation.

No one assigned these roles by hand. As training went on, the heads simply ended up doing different jobs on their own — having several heads do the same thing would be wasteful, so they naturally split apart.

Not every head ends up with a role a person can point to, though. Plenty are hard to explain, and some can be removed with almost no drop in performance.

One strand alone blurs everything

With a single head, only one set of weights gets computed. But a sentence usually has several relationships holding true at the same time. In "the sibling ate the cake Jiyoon bought yesterday," who bought it, what got eaten, and when all need to be tracked separately.

With only one set of weights, all these relationships get crammed into one table and bleed into each other. The result leans strongly toward nothing and sits near the average — like an entire choir singing one melody.

Split the heads and each can watch its own relationship clearly, free to lean weight wherever it wants without stepping on the others.

Splitting doesn't add more work

This is where a common mix-up happens. Grow the heads to eight and it seems like the compute should grow eightfold too — it doesn't. The width one head used to use gets divided into eight instead. Each head works in that narrower slice, and putting them back together restores the original width.

It's the same as a choir splitting into four parts without gaining more singers — just a different arrangement of the same people. That's how the compute cost stays about the same while expressive power goes up.

Everything gets merged into one block at the end

What each head produces doesn't stay separate. The results are concatenated in order into one block at the original width, and that block goes through one more cleanup pass, where the heads' results mix and their weight gets adjusted.

Stop at just concatenating and the heads would remain strangers to each other. That final cleanup is what turns four parts into one chord.

More heads isn't always better

Keep adding heads and the width each one gets keeps shrinking. Too narrow, and there isn't enough room to hold information, so performance can actually drop. Model size sets a reasonable count, usually somewhere between eight and a few dozen.

Looking at a trained model, a fair number of heads turn out to carry little role at all — and methods for trimming those out to make a model lighter have been studied too.