Multi-head attention

One attention head learns one pattern of routing: one way of deciding which tokens matter. A real model needs many patterns at once, such as one head that tracks the previous word and another that matches a pronoun to its noun. Multi-head attention runs several heads in parallel, each with its own projections, and combines their outputs.

Splitting the width

Start with a model width d_model and h heads. Each head gets a width of d_k = d_model / h, so the total width stays the same. Each head computes its own attention, then the results are concatenated and projected back:

head_i       = Attention(X W_Q^i, X W_K^i, X W_V^i)
MultiHead(X) = Concat(head_1, …, head_h) W_O

W_O has shape d_model × d_model. Adding up the Q, K, V and O projections gives about 4·d_model² parameters per layer, ignoring biases.

Walking through the shapes

Take a Llama-style layer with d_model = 4096 and h = 32 heads:

1. Input X has shape (batch, 128, 4096) for 128 tokens.
2. Each projection gives (batch, 128, 4096). Reshape it to (batch, 32, 128, 128), so each of the 32 heads works in width 128.
3. Each head computes attention, giving (batch, 32, 128, 128).
4. Merge the heads back to (batch, 128, 4096) and apply W_O.

Parameter count for the four projections is 4 × 4096² = 67,108,864 per layer, the same order as a single full-width head. You split the width. You do not multiply the cost by h.

For comparison, GPT-2 small used d_model = 768 with h = 12, so each head had d_k = 64.

Why several heads help

Different heads tend to specialize. Some copy patterns from earlier in the sequence, some follow syntax, and some attend to nearby tokens. Empirically, many small heads do better than one wide head at the same compute.

Common mistakes

• Assuming h heads cost h times as much. You split the width, so the cost stays close to a single full-width head.
• Mixing up query heads and key-value heads. Variants such as MQA and GQA share keys and values across heads. Count the two kinds separately, because the KV count drives cache memory.
• Expecting each head to have a clean label. Heads specialize loosely, and production models rarely have heads you can name cleanly.

Check yourself

1. A model has d_model = 768 and h = 12. What is d_k? 768 / 12 = 64.
2. If you double h and keep d_model fixed, what happens to d_k? It halves.

Key takeaways

• Multi-head attention runs several attention patterns on slices of the model width.
• Each head has width d_k = d_model / h, so the total cost stays about the same.
• Concatenate the heads and project with W_O to get back to d_model.
• The head count is a trade-off between pattern variety and width per head.

Further learning

The last step is the cost of attention itself, and what the variants do about it: cost and long context.