Skip to content

Multi-Head Attention

Overview

A single attention pass can learn only one notion of relevance— the highest-scoring pattern wins. Multi-head attention (MHA) runs h attention passes in parallel, each with its own learned W_Q, W_K, W_V, then concatenates the results and projects them back. With h = 8–32, the model simultaneously tracks syntax, coreference, semantics, and position— instead of forcing one attention pattern to do all of it.

Single head: one set of Q/K/V → one notion of "what matters"
Multi-head: h parallel heads → h complementary views
 → concat → linear projection → richer output
  • Paper: "Attention Is All You Need" (Vaswani et al., 2017)— MHA introduced here
  • Cost: same total FLOPs as a single wide head (heads split d_model)
  • Prerequisite: 02 Scaled Dot Product Attention, [01 Query, Key & Value](/01-modeling/00-fundamentals/01-attention/(01-query-key-value/)
  • Efficiency variant: Mqa & Gqa— share K/V to shrink the KV cache

-

Why One Head Is Not Enough

A single attention function must compromise:

One head, one similarity function
 "The bank rejected my loan application"

 To predict the next word well, the model needs:
 a) "bank" ←→ "loan" (semantic/syntactic link)
 b) "my" ←→ "application" (possessive link)
 c) "rejected" ← — "bank" (subject-verb link)
 d) position rhythm (recent-token locality)

 A single Q/K projection picks ONE similarity notion →
 the other links get no dedicated representation.

MHA's answer: don't choose — maintain h separate Q/K/V spaces,
each free to specialize in one kind of link. The model can use:
 head 3 → subject-verb
 head 7 → nearest-neighbor position
 head 12 → coreference
...simultaneously.

-

How Multi-Head Attention Works

The four sub-computations

Step 1 — PROJECT per head:
 For head i: Qᵢ = X·W_Qⁱ Kᵢ = X·W_Kⁱ Vᵢ = X·W_Vⁱ
 Shapes: X (seq, d_model) × W_Qⁱ (d_model, d_k) → Qᵢ (seq, d_k)

Step 2 — ATTEND per head (independently):
 head_i = softmax(QᵢKᵢᵀ / √d_k) · Vᵢ ← [02 Scaled Dot Product Attention](/01-modeling/00-fundamentals/01-attention/02-scaled-dot-product-attention/)

Step 3 — CONCAT:
 Z = concat(head₁, head₂,..., headₕ) shape: (seq, h·d_k) = (seq, d_model)

Step 4 — PROJECT out:
 Output = Z·W_O with W_O: (d_model, d_model)
 → mixes information across heads

The full equation

MultiHead(Q, K, V) = Concat(head₁,..., headₕ) · W_O

where headᵢ = Attention(Q·W_Qⁱ, K·W_Kⁱ, V·W_Vⁱ)

Visual summary

 X (seq_len, d_model)
 │
 ┌───────┼────────┬────────┬────────┐
 W_Q¹ W_K¹ W_V¹ │ W_Q² W_K²...
 │ │ │ │ │ │
 Q¹ K¹ V¹ │ Q² K² V²
 └──┬───┴───┬──┘ │ └──┬──┴──┬─┘
 attention(Q¹,K¹,V¹) attention(Q²,K²,V²)
 │ │
 └─────────concat───────┘
 │
 W_O (d_model×d_model)
 │
 output (seq, d_model)

-

Head Count and Dimensions

The d_model split trade-off

d_model = h × d_k

Example: LLaMA-2 7B: d_model = 4096, h = 32, d_k = 128

Same total parameters whether you use:
 - 1 head of width 4096 (single head)
 - 8 heads of width 512 (fewer, wider)
 - 32 heads of width 128 (many, narrow) ← typical modern choice

More heads (narrower):
 more specialized patterns, better parallel expressiveness
 more concat/O² overhead, marginal gains past ~32-64 heads

Head counts in real models

Model Heads Head dim d_model
BERT base 12 64 768
GPT-2 small 12 64 768
GPT-3 96 128 12,288
LLaMA-2 7B 32 128 4,096
Mistral 7B 32 128 4,096
GPT-4 ~100+ ~128 ~16,384 (speculated)

-

Why Heads Specialize

Empirical observations (BERT/LLaMA head analysis)

Head 5, layer 3: syntactic dependency
 "The cat sat on the mat"
 └─ "sat" strongly attends to "cat" (subject-verb link)

Head 12, layer 7: coreference
 "John told Mary he liked her"
 └─ "he" attends to "John", "her" attends to "Mary"

Head 3, layer 2: position/rhythm
 attends to the token before or after (local structure)

High-layer heads: broader; some heads attend to sentence boundaries
Low-layer heads: mostly local, syntactic

Heads are never instructed

No head is told "you do coreference".
Specialization emerges purely from the next-token/MLM training loss:
 - heads that reduce the loss get reinforced
 - redundant heads are pushed into different niches by competition
 - the training data's language structure shapes the division of labor

Interesting findings about head redundancy

Research (Voita et al., 2019; Michel et al., 2019):
 - many heads can be PRUNED with little quality loss
 (some BERT runs lose ~1% after removing ~70% of heads)
 - a few "important" heads carry most of the benefit
 - redundancy exists across heads AND across layers
 Implication: MHA is robust — heads overlap functionally,
 which is why pruning/GQA work without catastrophic damage.

-

Multi-Head vs. Single Head: Same Cost, More Power

FLOPs comparison (per layer):

Single wide head:
 Q,K,V: 3 · seq · d_model² ← one big projection
 scores + softmax + weight: 2 · seq² · d_model

Multi-head (h heads × width d_k):
 projections: 3h · seq · d_model · d_k = 3 · seq · d_model² (same!)
 scores: h · 2 · seq² · d_k = 2 · seq² · d_model (same!)

Total compute is identical — only the INTERNAL structure changes:
 - heads parallelize across GPU SMs / tensor cores
 - narrow heads do cheaper small matmuls (more memory-friendly)
 - the concat+W_O adds one small linear layer

The one real extra cost: the KV cache grows with h × d_k per token— which is exactly why MQA/GQA share K/V heads. See Mqa & Gqa.

Key Takeaways

h parallel views: each head attends in its own learned space Concat + W_O: final projection mixes head outputs together Emergent specialization: heads learn syntax, coreference, position— unprompted d_model = h × d_k: same cost as one wide head, far more expressive Redundant heads exist: pruning and GQA exploit this safely Rule of thumb: 12 (small) to 32–96 (large) heads; diminishing returns past that

-