Multi-Head Attention

The idea

Run several attention operations in parallel. Each gets its own query, key and value projections, so each is free to compare words along a different axis and produce a completely different weighting. Then concatenate their outputs and pass the result through one more linear layer to mix them back together.

MultiHead(Q,K,V) = Concat(head₁, …, head_h) W^O

That final projection matters. Without it the heads' outputs would sit in separate slices of the vector, never interacting; W^O is what lets the layer combine what different heads found.

The part that surprises people

Heads are not stacked on top of the model dimension — the model dimension is divided among them. With d_model = 512 and 8 heads, each head works in 64 dimensions, not 512.

So going from 1 head to 8 does not multiply the cost by eight. The projection matrices are the same total size either way, and the parameter count is unchanged. What changes is the shape of the computation: eight narrow attention patterns instead of one wide one.

This is close to a free lunch, and it is why every transformer uses many heads. The cost is that each head is individually less expressive, which is why very high head counts eventually stop helping — at some point d_k gets too small to represent anything useful.

Several attention computations at once

One attention computation produces one set of weights per position — one view of how the tokens relate. Language needs several at a time: which word is the grammatical subject, which pronoun refers to which noun, which bracket closes which.

Multi-head attention runs h independent attention computations in parallel, each with its own learned Q, K and V projections, then concatenates the results and applies one output projection.

MultiHead(X) = Concat(head₁, …, headₕ) W_O

The cost is arranged to be roughly free. With a model dimension of 512 and 8 heads, each head works in 512/8 = 64 dimensions. Eight 64-dimensional attentions cost about the same as one 512-dimensional attention — and give eight independent views instead of one.

Note that the per-head dimension stays around 64–128 across model sizes. Width is added by adding heads, not by making each head wider.

What the heads actually learn

Inspecting trained models finds genuine specialisation, and it is more varied than one might guess:

• Positional heads that mostly attend to the previous or next token.
• Syntactic heads that link verbs to their subjects or objects, corresponding closely to dependency-parse edges.
• Coreference heads that connect pronouns to their referents.
• Delimiter heads that match brackets, quotes and sentence boundaries.
• Rare-token heads that attend to unusual words carrying most of the information.
• Heads that appear to do very little.

That last category is well documented: pruning studies find a substantial fraction of heads can be removed with minimal loss, and in some layers a single head does most of the work. So multi-head attention provides capacity for several relationship types, and the specialisation that emerges is a side effect rather than a design.

Two induction-related heads deserve mention. In language models, "induction heads" that complete a pattern seen earlier in the context appear to underpin much of in-context learning, and their emergence during training coincides with a visible jump in that ability.

The shape gymnastics

In code, all heads are computed in one batched operation rather than in a loop:

B, T, D = x.shape                       # batch, tokens, model dim
H = 8; d = D // H                       # heads, dim per head

q = (x @ Wq).view(B, T, H, d).transpose(1, 2)     # (B, H, T, d)
k = (x @ Wk).view(B, T, H, d).transpose(1, 2)
v = (x @ Wv).view(B, T, H, d).transpose(1, 2)

scores = q @ k.transpose(-2, -1) / d ** 0.5        # (B, H, T, T)
out = (scores.softmax(-1) @ v)                     # (B, H, T, d)
out = out.transpose(1, 2).reshape(B, T, D) @ Wo    # back to (B, T, D)

The view and transpose pair is the whole trick: project once at full width, then reshape so the head dimension becomes a batch dimension. Every head is then computed by the same matrix multiplication.

F.scaled_dot_product_attention does all of this with a FlashAttention kernel underneath, and should be preferred in real code.

Eight small attentions instead of one big one

Splitting the model width across heads costs nothing in parameters and buys the ability to attend to several things at once. Here is the split, the concatenation, and a measurement of what the heads actually do differently.

import numpy as np

rng = np.random.default_rng(0)
words = ["the", "tired", "cat", "quickly", "chased", "the", "small", "mouse"]
T, D_MODEL, H = len(words), 32, 4
D_HEAD = D_MODEL // H
E = rng.normal(0, 1.0, (T, D_MODEL))

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

print("%d tokens, d_model=%d, %d heads -> d_head = %d/%d = %d."
      % (T, D_MODEL, H, D_MODEL, H, D_HEAD))
print()
print("the parameter count is the point. one head at full width against")
print("%d heads at %d each:" % (H, D_HEAD))
one = 3 * D_MODEL * D_MODEL
per_head = 3 * D_MODEL * D_HEAD
proj = D_MODEL * D_MODEL
print("   single head, Q/K/V : 3 x %d x %d           = %6d"
      % (D_MODEL, D_MODEL, one))
print("   %d heads,    Q/K/V : %d x 3 x %d x %d       = %6d   <- the same"
      % (H, H, D_MODEL, D_HEAD, H * per_head))
print("   output projection  : %d x %d               = %6d"
      % (D_MODEL, D_MODEL, proj))
print("   %-18s %22s %6d" % ("total, multi-head", "", H * per_head + proj))
print()
print("   the Q/K/V projections cost EXACTLY the same either way -- %d heads"
      % H)
print("   of width %d is the same %d numbers as one head of width %d."
      % (D_HEAD, one, D_MODEL))
print("   multi-head adds only the output projection, and that is the")
print("   entire price of being able to attend to several things at once.")
print()

Wq = rng.normal(0, 0.4, (H, D_MODEL, D_HEAD))
Wk = rng.normal(0, 0.4, (H, D_MODEL, D_HEAD))
Wv = rng.normal(0, 0.4, (H, D_MODEL, D_HEAD))
Wo = rng.normal(0, 0.4, (D_MODEL, D_MODEL))

heads, weights = [], []
for h in range(H):
    Q, K, V = E @ Wq[h], E @ Wk[h], E @ Wv[h]
    A = softmax(Q @ K.T / np.sqrt(D_HEAD))
    weights.append(A)
    heads.append(A @ V)

print("each head runs the same three lines on its own %d-dimensional slice:"
      % D_HEAD)
for h in range(H):
    print("   head %d: Q,K,V are %s -> attention %s -> output %s"
          % (h, (T, D_HEAD), weights[h].shape, heads[h].shape))
print()

cat = np.concatenate(heads, axis=1)
out = cat @ Wo
print("concatenate along the feature axis, then one linear projection:")
print("   %d heads x %d dims = %s, projected back to %s"
      % (H, D_HEAD, cat.shape, out.shape))
print("   the output projection is what lets the heads talk to each other.")
print("   without it they would be %d independent streams that never mix." % H)
print()

print("WHAT THE HEADS ACTUALLY DO. each row is 'chased', showing where that")
print("token looked according to each head:")
i = words.index("chased")
print("%8s %s" % ("", "".join("%9s" % w for w in words)))
for h in range(H):
    print("%8s %s" % ("head %d" % h, "".join("%9.4f" % a for a in weights[h][i])))
print("%8s %s" % ("mean", "".join("%9.4f" % a
                                  for a in np.mean([w[i] for w in weights], 0))))
print()
print("the heads disagree, and that disagreement is the whole reason to have")
print("more than one. measure it -- how similar is each pair of heads?")
print("%8s %s" % ("", "".join("%9s" % ("head %d" % h) for h in range(H))))
for a in range(H):
    row = []
    for b in range(H):
        fa, fb = weights[a].ravel(), weights[b].ravel()
        row.append(fa @ fb / (np.linalg.norm(fa) * np.linalg.norm(fb)))
    print("%8s %s" % ("head %d" % a, "".join("%9.4f" % r for r in row)))
off = [weights[a].ravel() @ weights[b].ravel()
       / (np.linalg.norm(weights[a].ravel()) * np.linalg.norm(weights[b].ravel()))
       for a in range(H) for b in range(H) if a != b]
print("   average similarity between different heads: %.4f" % np.mean(off))
print()
print("   these are random projections, so the differences here are luck.")
print("   in a trained model they specialise -- published analyses find")
print("   heads that track syntactic dependencies, heads that attend to the")
print("   previous token, heads that follow coreference, and a large number")
print("   that do nothing identifiable at all.")
print()

print("that last group is real and worth knowing about:")
entropies = []
for h in range(H):
    p = np.clip(weights[h], 1e-12, 1)
    entropies.append(-(p * np.log(p)).sum(1).mean())
uniform = np.log(T)
for h in range(H):
    print("   head %d: average attention entropy %.4f (uniform would be %.4f)"
          % (h, entropies[h], uniform))
print("   a head whose entropy sits at the uniform value is spreading its")
print("   attention evenly, which is another way of saying it is doing")
print("   nothing. pruning experiments routinely remove 20 to 40 percent of")
print("   heads from a trained transformer with almost no loss in accuracy.")
print()
print("the shapes, once, because the reshaping is where implementations go")
print("wrong:")
print("   x            (batch, seq, d_model)")
print("   after qkv    (batch, seq, 3 * d_model)")
print("   reshape      (batch, seq, heads, d_head)")
print("   transpose    (batch, heads, seq, d_head)   <- heads become a batch")
print("   attention    (batch, heads, seq, seq)")
print("   transpose    (batch, seq, heads, d_head)")
print("   reshape      (batch, seq, d_model)         <- the concatenation")
print("   the transposes matter: reshaping without them interleaves the")
print("   heads' dimensions instead of stacking them, and the bug is silent.")

Try it yourself

1. Collapse to one head. Set the Number of Heads slider to 1 and pick "drank" as the Focus Word. A single distribution has to cover subject, object and everything else at once, and Coverage drops.
2. Open them up. Set the Number of Heads slider to 4. The heads separate: one tracks the subject, another the object, another the neighbouring determiner. Distinct Targets rises — the same layer is now representing several relationships simultaneously.
3. Prove the diversity is the point. Tick Make Heads Identical. Every head now shares one projection, so all four rows become the same row, Head Disagreement falls to zero and Distinct Targets drops to 1. Four copies of one opinion are worth exactly one opinion.
4. Watch the width shrink. With Model Dimension at 64, step the Number of Heads slider from 1 to 4 and watch Dimension per Head fall from 64 to 16 while Total Parameters does not move at all. The heads are splitting the budget, not adding to it.
5. Make the heads too narrow. Set the Model Dimension slider to 16 with 4 heads. Each head now has 4 dimensions to work with — barely enough to express a relationship, which is the practical limit on how many heads are worth having.

What heads actually learn

Analyses of trained models have found heads that track syntactic dependencies, heads that attend to the previous token, heads that follow coreference. That is the encouraging half.

The other half is that many heads learn nothing legible. A large fraction attend overwhelmingly to the first token or to punctuation — behaviour usually read as a "no-op", somewhere to park attention when a head has nothing to contribute for this input. Several papers have shown that a majority of heads can be pruned after training with little loss, which suggests the redundancy is real and that heads matter more during training than at inference.

So treat clean per-head interpretations with care. The tidy examples in papers are selected; most heads are not that tidy.

Common mistakes

• Assuming more heads means more capacity. It means the same capacity, divided. Adding heads without raising d_model makes each one narrower.
• Choosing a head count that does not divide d_model. The dimension must split evenly; most implementations simply refuse otherwise.
• Forgetting W^O. Concatenation alone leaves the heads' outputs in disjoint slices. The output projection is what merges them.
• Reading one head's pattern as the model's reasoning. The layer's output is a mixture of all heads, then projected. No single head is the answer.

In one line

Multi-head attention runs several attention operations in parallel, each with its own projections, so a layer can represent several relationships at once instead of compromising between them in a single distribution that must sum to 1. Crucially the model dimension is split across heads rather than multiplied by them, so more heads cost no extra parameters — they buy diversity at the price of making each head narrower, which is why the count cannot be raised indefinitely. The concatenated outputs are mixed by a final projection, and in trained models many heads turn out to be redundant or parked on punctuation, so individual head patterns should be read with caution.

Variants that shrink the KV cache

During generation, the keys and values of previous tokens are cached so each new token costs one step rather than a full recomputation. With many heads that cache becomes the dominant memory cost — often larger than the model weights for long contexts.

Two variants address it directly:

Multi-query attention (MQA) keeps separate query projections per head but shares a single key and value projection across all of them. The cache shrinks by a factor equal to the head count. Quality drops slightly.

Grouped-query attention (GQA) is the compromise now standard in large models: heads are divided into groups, and each group shares one key/value pair. With 32 query heads and 8 key/value groups the cache is a quarter the size, with quality close to full multi-head.

This is a good example of an architectural choice driven entirely by inference economics rather than by modelling quality.

Cost and practical notes

Attention's score matrix is (batch, heads, tokens, tokens), so memory grows with the head count and the square of the sequence length. FlashAttention avoids materialising it at all, which is why it is now the default implementation everywhere.

Choosing the head count in a model you are designing: keep the per-head dimension at 64 or 128 and derive the count from the model width. Very small per-head dimensions (16 or fewer) reduce each head's expressiveness; very large ones lose the benefit of having several views.

For fine-tuning and use, this is not a parameter you choose — it is fixed by the pretrained checkpoint.

Questions people ask

Why not one big attention head? It can only express one weighting pattern per position. Several heads represent several relationships simultaneously.

How many heads is best? 8–32 for typical sizes, with per-head dimension 64–128. Rarely worth tuning.

Do all heads matter? No — pruning studies show many can be removed with little loss. The redundancy appears to help training rather than inference.

What is the output projection for? It mixes the concatenated head outputs back into the model dimension, letting the layer combine what the heads found.

Is multi-head attention more expensive? Not meaningfully — the dimensions are split, so total work is comparable to single-head at full width.

Why does GQA exist? To shrink the KV cache during generation, which is the binding memory constraint for long-context inference.

Recap in one screen

• Several attention computations in parallel, each with its own projections, concatenated and projected out.
• The model dimension is split across heads, so the cost is roughly that of one full-width head.
• Heads specialise — syntax, coreference, delimiters — and many turn out to be redundant.
• All heads are computed as one batched matrix multiplication via reshaping.
• Grouped-query attention shares keys and values across heads to shrink the generation cache.