Multi-Head Attention

Summary: A full attention block in a transformer is not one attention head but many, each with its own independent $W_Q$, $W_K$, $W_V$, running in parallel. Their outputs all get added to the residual stream, giving the model the capacity to learn many distinct ways that context changes meaning.

Why many heads (run in parallel)?

A single attention head can learn one pattern of “which tokens should influence which, and how?“. But language context works in many ways at once:

• Adjectives update their nouns.
• Pronouns need to find their antecedents.
• “they crashed the” changes what kind of car is expected (structural damage, broken glass).
• The word Harry near wizard vs Harry near Queen, Sussex picks out different people.

Forcing a single head to learn all of these simultaneously would be a nightmare. Running many heads in parallel, each with its own small parameter budget, lets different heads specialise — and summing their contributions gives the token vector access to all of them at once.

The model is given the capacity to learn many distinct ways that context changes meaning.

How it’s wired

For each block and each head $h\in \{1,\dots ,H\}$:

• Independent $W_Q^{(h)}$, $W_K^{(h)}$, $W_V^{(h)}$ (the last factored into $W_V^{(h)\downarrow }$ and $W_V^{(h)\uparrow }$).
• Each head computes its own attention pattern and its own $\Delta {\vec{E}}_j^{(h)}$ for every token position $j$.
• The block’s output for position $j$ is the sum over heads, added to the incoming embedding:

${\vec{E}}_j^{\prime }={\vec{E}}_j+{\sum }_{h=1}^H\Delta {\vec{E}}_j^{(h)}$

Dimensionality example with GPT-3:

Original embedding: ${\vec{E}}_j\in {\mathbb{R}}^{12,288}$ Refinement from head $h$: $\Delta {\vec{E}}_j^{(h)}\in {\mathbb{R}}^{12,288}$ — each head proposes a nudge in the full embedding space Number of heads: $H=96$ Updated embedding: ${\vec{E}}_j^{\prime }\in {\mathbb{R}}^{12,288}$ — same dimensionality, moved to encode context

Summing 96 proposed changes — every head independently produces a 12,288-dim refinement; they all get added together and then added to the original:

$\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0^{\prime }\\ e_1^{\prime }\\ ⋮\\ e_{12287}^{\prime }\end{array}\right]}}=\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ e_{12287}\end{array}\right]}}+\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}\Delta e_0^{(1)}\\ \Delta e_1^{(1)}\\ ⋮\\ \Delta e_{12287}^{(1)}\end{array}\right]}}+\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}\Delta e_0^{(2)}\\ ⋮\\ ⋮\\ \Delta e_{12287}^{(2)}\end{array}\right]}}+\cdots +\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}\Delta e_0^{(96)}\\ ⋮\\ ⋮\\ \Delta e_{12287}^{(96)}\end{array}\right]}}$

Equivalently, as a matrix multiply (the $W_O$ view) — concatenate all heads’ small outputs and multiply by the output matrix:

$\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0^{\prime }\\ e_1^{\prime }\\ ⋮\\ e_{12287}^{\prime }\end{array}\right]}}=\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ e_{12287}\end{array}\right]}}+\underset{W_{O}12,288\times 12,288}{\underbrace{\left[\begin{array}{ccccc}\cdot & \cdots & \cdots & \cdots & \cdot \\ ⋮& & & & ⋮\\ ⋮& & & & ⋮\\ \cdot & \cdots & \cdots & \cdots & \cdot \end{array}\right]}}\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{z}_0^{(1)}\\ ⋮\\ z_{127}^{(1)}\\ z_0^{(2)}\\ ⋮\\ z_{127}^{(96)}\end{array}\right]}}$

The concatenated vector on the right has $96\times 128=12,288$ entries — one 128-dim chunk per head. $W_O$ is square ($12,288\times 12,288$) but is really 96 stacked $V_{\uparrow }$ matrices (each $12,288\times 128$). See terminology gotcha.

GPT-3: $H=96$ heads per block, $L=96$ blocks. So the model runs $96\times 96=9,216$ distinct attention heads in total.

Parameter count (GPT-3)

Per head:

• $W_Q$, $W_K$: each $d_{\text{embed}}\times d_k=12,288\times 128\approx 1.57\text{M}$
• $W_V^{\downarrow }$: $d_k\times d_{\text{embed}}$, same ~1.57M
• $W_V^{\uparrow }$: $d_{\text{embed}}\times d_k$, same ~1.57M
• Total per head: ~6.3M

Per block: $6.3\text{M}\times 96\text{ heads}\approx 600\text{M}$ Across all 96 layers: ~58B parameters devoted to attention. About a third of GPT-3’s 175B.

Low-rank values (why factor $W_V$?)

Without factoring, $W_V$ would be square ($12,288\times 12,288\approx 151\text{M}$) per head — so one head would blow the entire parameter budget of the Q/K matrices combined. Multiply by 96 heads × 96 layers and the arithmetic doesn’t close.

Factoring $W_V$ through the small $d_k$-dimensional space:

1. Drops per-head value params to match $Q$ and $K$ (~3.1M split across the two factors).
2. Across all heads in a block, the up-projection matrices can be stacked and viewed as a single big output projection (often written $W_O$ in papers and textbooks) — same total operation count, different bookkeeping.
3. The overall linear value map has rank $\le H\cdot d_k$, which is still much less than $d_{\text{embed}}$ for early models but grows as $H$ grows.

Terminology gotcha: “value matrix” vs output matrix

Theory vs practice

3b1b’s “low-rank value transformation” framing — each head owning its own $V_{\downarrow }$ and $V_{\uparrow }$ — is conceptually clean, but it is not how papers and real implementations write it. The standard convention is:

• The per-head value matrix is only the “down-projection” $V_{\downarrow }\in {\mathbb{R}}^{128\times 12,288}$ (maps down from embedding space → smaller $d_k$-dim space (${\mathbb{R}}^{12,288}\to {\mathbb{R}}^{128}$ in GPT-3)).
  • When a paper says “the value matrix of head $h$,” this is what it means; often denoted, $W_V\in {\mathbb{R}}^{128\times 12,288}$
• All $H$ “up-projections” $V_{\uparrow }^{(h)}\in {\mathbb{R}}^{12,288\times 128}$ are concatenated into one giant matrix, the output matrix $W_O\in {\mathbb{R}}^{12,288\times 128H}$, which belongs to the multi-head block as a whole, not to any single head.
  • This maps back up to $d_{\text{embed}}$ space.

The two framings compute exactly the same function — ${\sum }_h{V}_{\uparrow }^{(h)}\Delta {\text{attn-output}}^{(h)}$ vs $W_O[\text{concat over }h]$ — but with the per-head/per-block split drawn in different places. If you only read 3b1b and then open the original Attention Is All You Need paper, this renaming is the single most confusing thing. Same operation, different bookkeeping.

What you actually get

• Parallelism within a block. All heads can run simultaneously on a GPU — this is one reason attention is GPU-friendly.
• Specialisation. Different heads learn to attend to different relationships. Interpretability work has found heads that do things like “attend to the previous token”, “attend to the same token earlier in the sequence” (induction heads), “attend to the subject of the current sentence”, etc.
• Redundancy and ablation. Many heads can be pruned after training with little loss — evidence that the model over-provisions capacity during training and then settles into a smaller effective set.

See also

• attention-mechanism — a single head in detail
• transformer-architecture — where multi-head attention blocks sit in the full stack
• src-3b1b-llms-ch3-attention

Sources

• src-3b1b-llms-ch3-attention