S109 · Multi-Head Attention — Parallel Views

Module M13: NLP & Transformers · Session 109 of 130 · Track: LLM 👁️

What you'll be able to do after this session

• Explain Multi-Head Attention to a friend in one minute, out loud, no notes.
• Recognise it when you see it in real code / systems / papers.
• Complete the hands-on exercise at the end without looking up help.

Prerequisites

If you can't answer these in 30 seconds each, redo the linked session first:

• S108 — Q/K/V Math — Scaled Dot-Product Attention Derived

Pre-read (skim before the session)

• 🎥 Intuition (5–15 min) · Multi-head attention — CodeEmporium — Short, clear explanation of why one head isn't enough.
• 🎥 Deep dive (30–60 min) · Let's build GPT — Andrej Karpathy (multi-head section) — Multi-head implemented from scratch.
• 🎥 Hands-on demo (10–20 min) · Multi-Head Attention in PyTorch — Umar Jamil — Types up MHA and BOOMs through the tensor shapes.
• 📖 Canonical article · The Illustrated Transformer — Jay Alammar (Multi-Head section) — Beautiful visualisations of parallel attention heads.
• 📖 Book chapter / tutorial · The Annotated Transformer — Harvard NLP — MultiHeadedAttention class implemented line by line.
• 📖 Engineering blog · Peter Bloem — Transformers from scratch (Multi-head section) — Rigorous walk-through with param counting.

(a) Intuition · 5 min

The one-paragraph mental model. Why this concept exists, what problem it solves, and the analogy that makes it stick.

In one sentence: Why one attention isn't enough.

One attention head learns one kind of relationship — say, syntactic subject-verb pairs. But sentences have many kinds of relationships simultaneously: subject-verb, adjective-noun, coreference (who "it" refers to), long-range topical links, etc. Multi-head attention just runs several attention operations in parallel, each with its own learned Q/K/V projection matrices, and then concatenates and mixes their outputs. Each head can specialise on a different relationship without the model having to trade one off against another.

In code, it's almost embarrassingly simple: instead of one attention with dimension d_model = 768, you slice 768 into h = 12 heads of d_head = 64 each, run 12 independent attentions in parallel (matrix-multiplied together for efficiency), concatenate the 12 outputs back to 768, and pass through one more linear layer W_O to mix them. Total parameter count and compute stay roughly the same as single-head attention with d_k = d_model, but expressiveness goes up substantially.

Interpretability researchers have watched trained heads and found roles like "points to the previous token", "attends to matching brackets", "resolves pronouns to antecedents", "tracks the position of the verb". Not every head is interpretable — many look like noise — but the diversity is what makes the whole architecture work.

Gotcha to carry forever: In multi-head attention, d_model must be divisible by n_heads. d_model = 512, n_heads = 8 → d_head = 64. Getting the reshape wrong is the single most common transformer implementation bug — always print tensor shapes at each step when debugging.

(b) Visual walkthrough · 15 min

Diagrams, worked examples, step-by-step visuals.

The full multi-head flow, shape by shape.

Worked example — reshape gymnastics. With B=2, T=10, d_model=512, n_heads=8, d_head=64:

1. Input X is [2, 10, 512].
2. Apply W_Q ([512, 512]) → Q is [2, 10, 512].
3. .view(B, T, n_heads, d_head) → [2, 10, 8, 64].
4. .transpose(1, 2) → [2, 8, 10, 64]. Now the heads dim sits before the sequence dim, so per-head attention is a batched matmul over [B, n_heads].
5. Attention: Q @ K^T → [2, 8, 10, 10]. Softmax + V. Output [2, 8, 10, 64].
6. .transpose(1, 2).contiguous().view(B, T, d_model) → back to [2, 10, 512].
7. Final W_O linear → [2, 10, 512].

Parameter count. For a single MHA layer with d_model = 768:

• W_Q, W_K, W_V, W_O each [768, 768] = 589,824 params each → 2,359,296 total.
• That's the vast majority of a transformer block's parameters. The FFN block adds another 768 · 3072 · 2 = 4.7M.

(c) Hands-on · 20 min

Code you type and run. Not pseudo-code — real, runnable, copy-pasteable.

# multi_head_attention.py — a minimal but real MHA module
import torch, torch.nn as nn, torch.nn.functional as F, math
 
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model, self.n_heads = d_model, n_heads
        self.d_head = d_model // n_heads
        # One linear layer producing Q, K, V at once (a common efficiency)
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)
 
    def forward(self, x, is_causal=False):
        B, T, D = x.shape
        qkv = self.qkv(x)                                   # [B, T, 3*D]
        q, k, v = qkv.chunk(3, dim=-1)                      # each [B, T, D]
        # split heads: [B, T, nh, dh] -> [B, nh, T, dh]
        q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        # attention (uses Flash Attention when available on GPU)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
        # merge heads back
        y = y.transpose(1, 2).contiguous().view(B, T, D)
        return self.out(y)
 
# ---- exercise -------------------------------------------------------------
torch.manual_seed(0)
mha = MultiHeadSelfAttention(d_model=64, n_heads=4)
x = torch.randn(2, 10, 64)                                   # [B, T, D]
y = mha(x, is_causal=True)
print("input :", x.shape)
print("output:", y.shape)
print("params:", sum(p.numel() for p in mha.parameters()))
 
# Verify shapes for various head counts
for h in (1, 2, 4, 8, 16):
    m = MultiHeadSelfAttention(64, h)
    print(f"n_heads={h:2d}  d_head={64//h:2d}  out shape={m(x).shape}")

Checklist while it runs:

1. Output shape matches input [2, 10, 64] — attention is shape-preserving.
2. params is roughly 4 · 64² = 16,384 (Q/K/V/O linears) plus biases.
3. Try n_heads = 5 (not a divisor of 64) — assertion fires.
4. Compare is_causal=True output rows 0 and 5 — row 5 has access to more info.
5. Time a forward pass with T=2048 on GPU — see Flash Attention's speedup by comparing to a manual Q @ K.T implementation.

Try this modification: Add dropout after attention weights (before multiplying by V) — the standard "attention dropout" used in every real transformer. Wrap F.scaled_dot_product_attention(..., dropout_p=0.1).

(d) Production reality · 10 min

How this shows up in real systems, gotchas, war stories, common bugs.

Multi-head attention is where transformers' compute and memory scale, so most efficiency work targets exactly this layer.

Grouped-Query Attention (GQA). Introduced in Llama-2 70B and adopted by Gemini, Mistral, and Llama-3. Instead of n_heads separate K/V heads (one per Q head), GQA shares each K/V head across a group of Q heads (e.g. 4 Q heads share 1 K/V). This cuts KV-cache size by 4× at almost no quality cost — a huge deal for serving long contexts. Multi-Query Attention (PaLM) is the extreme case: one K/V for all Q heads.

Head pruning. Post-training, many heads turn out to be redundant. Michel et al. (2019) showed you can prune 30–50% of heads in a trained transformer with minimal quality loss, saving inference compute. Production models rarely do this, though, because the savings compound less than expected on modern GPU kernels — the memory-bandwidth cost of loading weights dominates.

Interpretability. Anthropic's transformer circuits research (2021–2023) found specific heads that implement identifiable algorithms: "induction heads" that copy-paste from earlier in the sequence, "IOI heads" that resolve indirect objects. Understanding what individual heads do is now a whole research subfield, and it started from looking at multi-head attention visualisations.

Common bugs.

• Non-contiguous tensor before view. After .transpose(1, 2), calling .view() fails silently in some PyTorch versions or crashes in others. Always .contiguous().view(...) when merging heads.
• Head-dimension mismatch across layers. If you accidentally use different d_head across layers, residual connections break shape.
• Off-by-one in reshape. Confusing [B, T, nh, dh] with [B, nh, T, dh] — the latter is what attention needs, the former is what .view() naturally produces. Always .transpose(1, 2).
• KV-cache mismatch under GQA. If your inference code assumes n_kv_heads == n_heads but the model was trained with GQA, cache alignment silently wrong.

Jay Alammar's illustrated post has the best diagrams of MHA anywhere online; every LLM engineer has that page bookmarked.

(e) Quiz + exercise · 10 min

1. Why use multiple attention heads instead of one very wide attention?
2. Given d_model = 768 and n_heads = 12, what is d_head and why does divisibility matter?
3. Walk through the sequence of tensor shapes from input [B, T, d_model] through MHA back to [B, T, d_model].
4. What does Grouped-Query Attention (GQA) share, and what production benefit does it give?
5. Why must you call .contiguous() before .view() after a .transpose()?

Stretch (connects to S108 — Q/K/V Math): In a single-head attention with d_k = 768, the softmax scaling constant is sqrt(768). In an 8-head attention with d_k = 96, it's sqrt(96). Why is the smaller scale correct even though the model dim is the same?

Explain-out-loud test

If you can't teach these 3 things to a friend without notes, redo the session:

1. What is Multi-Head Attention? (one sentence, no jargon)
2. When would you use it? (one concrete example)
3. When would you NOT use it? (one anti-pattern)

What comes next

• Next session (S110): Positional Encoding — Sinusoidal, Learned, RoPE
• Previous (S108): Q/K/V Math — Scaled Dot-Product Attention Derived
• Hub: The 6-Month Learning Plan

Part of a 130-session evergreen learning series. Session structure: (a) intuition · (b) visual walkthrough · (c) hands-on · (d) production reality · (e) quiz + exercise. Duration: 60 minutes.