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

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](/02-llm-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