DL S037 · Multi-Head Attention — Why Eight Heads Beat One

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 37 / 80 · Module M07 · ~2 hours

The story

You built single-head self-attention in S036. It works. It even scales — a d_model = 512 head with d_k = 512 has more capacity than a d_k = 64 head. So why not just make one really big head?

Because empirically, splitting the same parameter budget across h = 8 heads of d_k = 64 each beats one head of d_k = 512. Consistently. On every benchmark. The original transformer paper's ablation table (Table 3) is worth taping to your wall — halving heads (single head) costs ~1 BLEU on English–German, doubling heads (16) gains almost nothing. Eight is a sweet spot that's held up for eight years.

The reason "why" has two parts. The obvious one is capacity through specialisation: different heads can learn different types of dependencies (short-range syntax, long-range coreference, position-relative patterns, etc.) rather than one head having to average across all of them. The subtler one is ensembling in the attention distribution: soft-max is a peaky operator, and with one head you're forced to pick one focus per query, while with eight heads you can attend to eight different things simultaneously and let the output projection blend them.

This session we build multi-head attention, count parameters carefully so you never confuse d_model and d_k, look at real attention patterns from a trained BERT, and hit the two shape gotchas that eat 30 minutes of every first-time implementer's life.

Prerequisites

• Session 036 — scaled dot-product attention. This session is a wrapper around it.
• Session 003 — the reshape + permute sequence in MHA is where broadcast fluency pays off.

1 · The one-line change

Take S036's Attention(Q, K, V). Multi-head attention is:

with W_Q^{(i)}, W_K^{(i)}, W_V^{(i)} ∈ ℝ^{d_{model} \times d_k} per head, and W_O ∈ ℝ^{h·d_v \times d_{model}} mixing the concatenated heads back into d_model.

The trick, which is not stated in the equation but is critical for implementation, is that we set d_k = d_v = d_model / h. So the total width of the concatenated heads is exactly d_model, and W_O is (d_model, d_model). This means MHA has the same parameter count as a single head with d_k = d_model (plus the W_O), but the attention pattern is now h different (T, T) matrices instead of one.

2 · Parameter count — do it once, know it forever

For MHA with model dim d_model and h heads (so d_k = d_v = d_model / h):

• h copies of W_Q^{(i)} each (d_model, d_k) = h · d_model · (d_model/h) = d_model².
• Same for W_K and W_V. Total so far: 3 · d_model².
• W_O: (d_model, d_model) = d_model².
• Total: 4 · d_model².

Note: no h in the formula. Number of heads does not change parameter count. It only changes how you slice the same parameters. This is the most common thing people get wrong on transformer interview questions.

3 · The efficient implementation — one big projection, then reshape

Naively, you'd write h separate nn.Linear(d_model, d_k) layers. In practice, everyone uses one nn.Linear(d_model, d_model) for Q (and one each for K, V), then reshapes the output into (B, h, T, d_k). Mathematically identical, dramatically faster on GPU because it's one big matmul instead of h small ones.

Here's the shape journey. Start with X: (B, T, d_model).

Q = X @ W_Q     # (B, T, d_model)         — one big linear
Q = Q.view(B, T, h, d_k)                  # split last dim into heads
Q = Q.transpose(1, 2)                     # (B, h, T, d_k) — put heads before T

Same reshape+transpose for K and V. Now compute attention per head, in parallel using broadcasting:

scores = Q @ K.transpose(-2, -1) / sqrt(d_k)   # (B, h, T, T)
weights = softmax(scores, dim=-1)              # (B, h, T, T)
out = weights @ V                              # (B, h, T, d_k)

Reverse the reshape:

out = out.transpose(1, 2).contiguous()         # (B, T, h, d_k)
out = out.view(B, T, d_model)                  # concat heads
out = out @ W_O                                # (B, T, d_model)

Study those two "reshape journeys" — forward and reverse — until they're automatic. This is the shape choreography every transformer implementation performs.

4 · The implementation

import torch, torch.nn as nn, torch.nn.functional as F, math
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, h, dropout=0.0):
        super().__init__()
        assert d_model % h == 0, "d_model must be divisible by h"
        self.d_model = d_model
        self.h = h
        self.d_k = d_model // h
 
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)
 
        self.dropout = nn.Dropout(dropout)

The __init__ allocates the four projections. Now the forward:

    def forward(self, x, mask=None):
        B, T, _ = x.shape
 
        # 1. Project and split into heads
        Q = self.W_Q(x).view(B, T, self.h, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(B, T, self.h, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(B, T, self.h, self.d_k).transpose(1, 2)
        # Q, K, V: (B, h, T, d_k)
 
        # 2. Scaled dot-product attention per head
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)   # (B, h, T, T)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        weights = F.softmax(scores, dim=-1)
        weights = self.dropout(weights)
        out = weights @ V                                        # (B, h, T, d_k)
 
        # 3. Merge heads and project
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_O(out)

Thirty lines. That is multi-head attention. Every transformer library on the planet is a variation on this snippet.

4.1 · Sanity-check test

mha = MultiHeadAttention(d_model=64, h=8)
x = torch.randn(2, 10, 64)   # batch 2, seq len 10
y = mha(x)
print(y.shape)               # torch.Size([2, 10, 64])
assert y.shape == x.shape

Input and output shapes must match. That's the "residual-compatibility" contract — S039 depends on it.

5 · What do the heads actually learn?

If you visualise attention patterns from a trained BERT or GPT, you'll see distinct roles. Voita et al. (2019) and Clark et al. (2019, "What does BERT look at?") catalogued them; the reproducible categories are:

The specialisation isn't perfect and isn't stable across seeds — some heads are messy. But the pattern is real, and it explains why multi-head works: language has many kinds of dependencies, and parallel heads let one model handle several at once instead of forcing every head to be a generalist.

6 · Three implementation bugs to know cold

RuntimeError: view size is not compatible with input tensor's size and stride

7 · Mermaid — the multi-head data flow

8 · Multi-head vs single-head — quick benchmark

Try this at home. Train a 4-layer transformer LM on Tiny Shakespeare (we'll do this in S041) with:

• Config A: d_model = 256, h = 1 (single 256-dim head).
• Config B: d_model = 256, h = 8 (8 heads of dim 32).

Same param count. Same optimiser. Same steps. Config B beats Config A by ~0.05–0.10 val loss consistently. Small but reliable. And attention pattern visualisations show Config B's heads have picked up distinct behaviours (some are diagonal, some skip-1, some CLS-focused) while Config A's one head is a smeary compromise.

9 · Modern-2025 twist — MQA, GQA, and MLA

The MHA you just built has a scaling wart: at inference time you cache K and V for every past token and every head. For Llama 3 70B (d_model = 8192, h = 64, 80 layers), a single 4k-context conversation's KV-cache is ~10 GB in FP16. That memory is what caps how many concurrent users you can serve on one H100. The fix — reduce the number of KV heads while keeping the number of Q heads — turned into a small research industry between 2019 and 2024:

• Multi-Query Attention (MQA) — Shazeer, 2019. Use h query heads but a single shared K and V head. KV-cache shrinks by h×. First deployed at scale in PaLM (2022) and Falcon-40B. Downside: perplexity worsens slightly and training becomes less stable.
• Grouped-Query Attention (GQA) — Ainslie et al., Google, 2023. Compromise: g shared K/V heads where 1 < g < h. Llama 2 70B, Llama 3, Mistral 7B, Mixtral, Gemma all use GQA with g = 8. Nearly MHA quality with MQA-ish memory savings.
• Multi-head Latent Attention (MLA) — DeepSeek-V2/V3, 2024. Compress K and V into a joint low-rank latent (d_c ≈ 4·d_h), decompress inside the attention op. Cuts KV cache 93% vs MHA, matches or beats MHA quality. Only DeepSeek and a few 2025 forks ship it because it interacts subtly with RoPE (they use "decoupled RoPE" — half the head is positional, half is content-latent).

A cheat table to file away (Llama-3-scale numbers, per token, FP16):

We rebuild MHA cleanly this session because it's the pedagogical foundation. In S042 you switch it to GQA (four lines of code) and in the wildcard track we sketch MLA. Every 2025 open model above 7B parameters uses one of MQA/GQA/MLA — none ship pure MHA anymore. It's the single biggest inference-time architectural shift since the transformer itself.

Further reading:

• Shazeer 2019 — MQA original
• Ainslie et al. 2023 — GQA
• DeepSeek-V2 paper §2.1 — MLA derivation with the decoupled-RoPE trick
• Sebastian Raschka: "Understanding MHA, MQA, GQA" (Jan 2024)

Retention scaffold

One-line summary (write it in your own words): _______________________________________________

Spaced review: re-read §3 (shape journey) and §9 (MQA/GQA/MLA table) in 24 hours. Revisit the full session on day 7.

Next session (S038): attention is permutation-invariant — shuffle the tokens, you get shuffled outputs. Language cares about order. We fix it with positional encodings — starting with sinusoids from Vaswani 2017, then previewing RoPE, and closing on 2025's YaRN and LongRoPE for million-token contexts.

Sticky note (keep on your desk): MHA = 4·d_model² params, h heads are just a slicing. Modern models slice Q one way and K/V another (GQA/MLA) to shrink the KV cache.

Previous: ← DL S036 · Next: DL S038 →