Multi-Head Attention

One budget, two relations

The last chapter closed the loop. The sentence arrives as , three learned matrices turn it into , and , and every word walks away with a new representation built out of the whole sentence. Each word receives one row of weights, one distribution, one blend.

Take chased. Its row came out 0.04, 0.18, 0.12, 0.06 and 0.60, and reading that row as a set of relations is enough to spot the problem. A verb needs its object, and mouse takes 0.60. A verb also needs its subject, and cat is left with 0.18. The row is a softmax, so it sums to one whatever happens, which means every point of weight spent on the object is a point not spent on the subject. One relation is served well, the other is not, and both arrive in the same vector.

The obvious response is more room: make large so one query can encode both relations. It does not help, and the reason is not capacity. However wide and become, the dot product still collapses to one number per pair, and what comes out is still one row. Dropping the softmax so the row need not sum to one is worse: without normalisation there are no proportions left to blend in.

Run the lookup more than once

If one lookup produces one distribution, two distributions ask for two lookups. That is the whole idea, and it is as blunt as it sounds. Nothing inside the mechanism changes. It runs more than once over the same sentence, each run with its own , and . Every run scores its own queries against its own keys, takes its own softmax, and hands chased its own row of weights. Two runs, two budgets, and now one of them can spend most of its weight on cat while the other spends most on mouse, with neither paying for the other. Each parallel run is a head.

Nothing new happens inside any head. Writing the -th head out in full:

with running from 1 to , and in the figure. The same goes into every head, and no head reads what another produced, so the runs happen side by side rather than one after another. Each returns an matrix, one row per word, exactly as a single head did.

Two loose ends. Three heads look like three times the parameters and three times the work, and each word now walks away with three output vectors where before it had one.

Splitting the cost

Three heads look expensive only because each one is quietly assumed to be as wide as the single head it replaces. Nothing forces that. The width of a head is a design choice — and were fixed by hand, and training only ever filled matrices of the size it was given. So if heads are going to run, make each one narrower: . Three heads, each a third as wide, and the cost is the one already being paid for one head.

Concretely, heads with in a model wide. Each is rather than , turning the sentence matrix into a matrix of queries. Real models are wider: the original Transformer takes and , which leaves .

The arithmetic works out the same on every axis:

1. Parameters. Each head holds three matrices of size . Across heads that is per role — the count of a single full-width head.
2. Projection. Projecting costs per head, so in total.
3. Scoring and blending. per head, so in total.

Multi-head attention is not several attention blocks bolted together. It is one attention block partitioned.

Putting the heads back together

chased is holding three vectors, one per head, each wide. Write them side by side as a single row: heads of width occupy exactly columns, so concatenating gives an matrix and every word leaves at the width it arrived at.

Concatenation alone is not enough. Each head still owns its own slice — columns 1 to 3 belong to head 1, 4 to 6 to head 2, 7 to 9 to head 3 — and no number in one block has met a number in another. The point of running several heads was to combine perspectives, and side by side is not combined.

So one more learned matrix, , of shape , multiplies the concatenated rows. Every output column is then free to be a mixture of all blocks, which makes the one place in the whole mechanism where one head meets another. It also corrects the cost argument above: this is a fourth matrix a single head never needed, so the count goes from to . The heads are free; re-joining them is what costs.

In code the head index is a leading dimension rather than a loop, so the projections happen for every head at once and the softmax runs on tables instead of one.

def multi_head_attention(X, W_q, W_k, W_v, W_o):
    """X (n, d_model); W_q/W_k/W_v (h, d_model, d_k); W_o (d_model, d_model)."""
    Q, K, V = X @ W_q, X @ W_k, X @ W_v       # (h, n, d_k), no head sees another
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1)          # (h, n, n)
    weights = torch.softmax(scores / d_k**0.5, dim=-1)
    heads = weights @ V                       # (h, n, d_v)
    concat = heads.transpose(0, 1).reshape(heads.shape[1], -1)
    return concat @ W_o, weights              # (n, d_model), (h, n, n)

What the heads actually learn

This chapter has talked as though head 1 handles objects and head 2 handles subjects. Nothing built here makes that happen. Look back for the place where a head is told what to look for: there is none. The heads differ because , and start as different random numbers, and they stay different because nothing pulls them back together. No term in the loss asks them to specialise.

So why do they end up different at all? Not because anything rewards difference, but because nothing removes it. Three sets of matrices start at three random points, and gradient descent moves each towards whatever reduces the loss from where it happens to stand.

The figures illustrate what different heads can look like, not what any particular model's heads do. chased leaves this chapter holding one vector, the width it arrived at, now possibly carrying its object and its subject together.

One thing has not been touched: which positions a head is allowed to read. Every head so far reads the entire sentence, every word scoring against every other. The next chapter takes some of that away.