Take the sentence "the cat quietly sat". When the model processes "sat", it would like to know two different things at once. Who is sitting (the subject, "cat")? And what came just before it ("quietly")? One attention head produces one set of weights per query, so it can put its weight on the subject, or on the previous token, or split it between them. It cannot do both cleanly. And if it splits, something worse happens than a loss of sharpness, which we will see in a moment.

The fix in Vaswani et al. (2017) is to run several attention heads side by side, each with its own queries, keys, and values, each free to attend somewhere different, and to combine their outputs at the end. This section explains why that helps, implements it the way every production codebase does (with one reshape and one transpose instead of a loop), counts its parameters and costs exactly, and surveys what trained heads are known to do.

One head blends what it finds

Suppose the values carry token identity as one-hot vectors, and a single head at "sat" splits its weight evenly between the subject "cat" and the previous token "quietly". Its output is

Which of the two is the subject? The output cannot say. A weighted average is symmetric in its inputs, so once two sources are blended into the same coordinates, the information about which source was found by which question is gone. You would get the same vector if "quietly" were the subject and "cat" the previous token.

Two heads fix this. Head 1 attends to the subject and returns $\mathbf{e}_{\text{cat}}$. Head 2 attends to the previous token and returns $\mathbf{e}_{\text{quietly}}$. Their outputs are kept in separate coordinates (concatenated), so the next layer can read "subject is cat" and "previous token is quietly" independently. Each head is a separate question with a separate answer.

There is also a capacity argument. From §4.2, a head's similarity is the bilinear form $\mathbf{W}_Q\mathbf{W}_K^\top$, of rank at most $d_k$. One head with $d_k = d$ has a single full-rank similarity but a single softmax. Splitting the same width into $h$ heads of $d_k = d/h$ gives $h$ separate similarities, each of rank $d/h$, and $h$ separate softmaxes. Neither is strictly more expressive than the other. The bet behind multi-head attention, borne out in practice, is that many cheap, independent lookups are worth more than one rich one. Vaswani et al. found single-head attention 0.9 BLEU worse than their best setting on translation, and quality dropped again with too many heads, so the number of heads is a real trade-off and not "more is better".

The definition

Each head $i \in \{1, \dots, h\}$ has its own projections $\mathbf{W}_Q^{(i)}, \mathbf{W}_K^{(i)} \in \R^{d\times d_k}$ and $\mathbf{W}_V^{(i)} \in \R^{d\times d_v}$, and computes ordinary scaled dot-product attention (§4.3),

with $\mathbf{M}$ the mask of §4.4. The head outputs are concatenated along the feature axis and mixed by an output projection $\mathbf{W}_O \in \R^{hd_v\times d}$,

In plain words, run $h$ small attention heads in parallel, put their answers side by side, and use one linear map to turn the combined answer back into a width-$d$ vector per token.

The standard choice is $d_k = d_v = d/h$. Then the concatenation has width exactly $d$, and, as we will count below, the whole layer costs the same as one head of width $d$. GPT-2 small uses $d = 768$ and $h = 12$, so $d_k = 64$. Llama 3 8B uses $d = 4096$ and $h = 32$, so $d_k = 128$. Head widths of 64 or 128 are by far the most common (Gemma is a notable exception at 256), and models get wider mostly by adding heads.

The output projection, and why heads add up

Split $\mathbf{W}_O$ into $h$ horizontal blocks of $d_v$ rows each, $\mathbf{W}_O^{(1)}, \dots, \mathbf{W}_O^{(h)} \in \R^{d_v\times d}$. Block matrix multiplication gives

In plain words, "concatenate, then project" is the same as "project each head separately, then add".

This identity is worth having in your pocket. It says the heads are independent and additive. Each head reads from the input, computes its own pattern, and writes its own contribution, and the contributions add. Elhage et al. (2021) build their whole analysis of Transformers on this view, in which each head writes into a shared residual stream (§5.1), with the head's OV circuit $\mathbf{W}_V^{(i)}\mathbf{W}_O^{(i)}$ deciding what it writes. It also tells you what $\mathbf{W}_O$ is for. Without it, each head's output would be stuck in its own fixed slice of coordinates. With it, every head can write in any direction of the width-$d$ space.

Reshape, transpose, and no loops

The definition suggests a loop over heads, and the first version most people write has one. Production code has none. The trick is to compute all heads' projections with one matrix product and then reshape.

Stack the per-head projections side by side, $\mathbf{W}_Q = [\mathbf{W}_Q^{(1)} \mid \cdots \mid \mathbf{W}_Q^{(h)}] \in \R^{d\times hd_k}$, and the same for keys and values. Then $\mathbf{Q} = \mathbf{X}\mathbf{W}_Q$ of shape $(B, T, h d_k)$ holds every head's queries, with head $i$ in columns $i d_k$ to $(i+1)d_k - 1$. Two moves put the heads on their own axis.

1. Reshape the last axis into two, $(B, T, hd_k) \to (B, T, h, d_k)$. In row-major memory this splits each token's vector into $h$ consecutive chunks of $d_k$, which is exactly head $i$'s columns. No data moves.
2. Transpose the head and time axes, $(B, T, h, d_k) \to (B, h, T, d_k)$. Now each (batch, head) pair owns a $(T, d_k)$ matrix, and the batched matrix product treats $(B, h)$ as batch dimensions.

Attention then runs on all heads at once. Scores are $(B, h, T, T)$ and outputs $(B, h, T, d_v)$. To merge, undo both moves. Transpose back to $(B, T, h, d_v)$ and reshape to $(B, T, hd_v)$, which is the concatenation, then multiply by $\mathbf{W}_O$. (See §0.1 for why reshape is free and transpose only changes strides.)

The order of those two moves is the single most common bug in multi-head attention, and you met it in §0.1, where Exercise 0.1c had you write split_heads and merge_heads. Reshaping $(B, T, d)$ directly to $(B, h, T, d_k)$ gives the right shape and the wrong contents, since what it calls head 1 holds whole tokens from the second half of the sequence instead of the second slice of every token. Attention computed on that runs without complaint, training proceeds, and the model is worse. The widget below shows the two splits in color, for any $d$ and $h$.

Here is the full layer, vectorized, checked against a loop over heads and against the "project each head, then add" identity.

import numpy as np

def multi_head_attention(X, W_Q, W_K, W_V, W_O, h, causal=True):
    raise NotImplementedError

What it costs, counted exactly

Parameters. With $d_k = d_v = d/h$, the stacked $\mathbf{W}_Q$, $\mathbf{W}_K$, $\mathbf{W}_V$ are each $d\times d$, and $\mathbf{W}_O$ is $d\times d$. So

independent of the number of heads. Splitting into heads is free in parameters. For GPT-2 small, which has biases, that is $4 \times 768^2 + 4 \times 768 = 2{,}362{,}368$ per layer, or about 28.3M across its 12 layers, close to a quarter of its 124M parameters. (GPT-2's code stores the first three as one c_attn layer of shape $768 \times 2304$ and the last as c_proj. The totals agree.)

Nothing forces $hd_k = d$, though. Gemma 7B has $d = 3072$ and 16 heads of width $d_k = 256$, so its stacked queries are $4{,}096$ wide, wider than the model, and $\mathbf{W}_O \in \R^{4096\times 3072}$ maps the concatenated heads back down. In general the four matrices hold $4\,d\,hd_k$ weights. Exercise 4.5e counts both cases.

Modern models also often share keys and values across groups of query heads, grouped-query attention. With $h$ query heads and $n_{kv}$ key-value heads, each shared by a group of $h/n_{kv}$ query heads, $\mathbf{W}_K$ and $\mathbf{W}_V$ shrink to $d\times n_{kv} d_k$, and $n_{kv} = h$ is the ordinary layer. Its purpose is to shrink what must be stored at inference time, so it is taught with that cache, in §6.5 and §12.5.

Compute. For one sequence of length $T$, counting a multiply-add as 2 FLOPs, the projections cost $4 \times 2Td^2 = 8Td^2$, the scores cost $h \times 2T^2d_k = 2T^2d$, and the weighted sum of values another $2T^2d$. The number of heads cancels again. §4.7 turns these into real numbers.

Memory. Here the heads do not cancel. A naive implementation materializes one $T\times T$ score matrix per head, so $h T^2$ numbers per sequence and per layer. More heads of the same total width mean more score matrices, which is one reason head widths stay at 64 or 128 instead of shrinking.

GPT-3's attention layers hold about 58 billion of its 175 billion parameters, a third, which is the same fraction as in every standard block ($4d^2$ of $12d^2$, as §5.6 counts).

def attention_param_count(d, h, d_k=None, bias=False):
    return 4 * d * d       # placeholder: right only in one case

What heads learn

Here we leave theorems behind. Everything in this part is an empirical finding about trained models, mostly from studies that look at attention weights or remove heads and measure the damage.

• Heads specialize. Clark et al. (2019) found heads in BERT that attend to the previous or next token, heads that attend broadly, and heads whose patterns track syntax, for example from a verb to its direct object or from a noun to its determiner, with accuracy well above simple baselines. Many heads put most of their weight on the separator token, which looks like the "attend to nothing" behavior of §4.4.
• A few heads do most of the work. Voita et al. (2019) classified the heads of a translation model into positional, syntactic, and rare-word heads, and found that pruning 38 of the 48 encoder heads cost only 0.15 BLEU on English to Russian. Michel et al. (2019) found that many heads can be removed at test time with little loss, and that some layers can be reduced to a single head.
• Some heads form circuits. In language models, previous-token heads (which you built in Exercise 4.2b) combine with induction heads in a later layer to continue repeated patterns, a mechanism found in models from tiny ones to large ones (Olsson et al. 2022). We build one by hand in §4.7, and circuit analysis traces how heads compose into larger algorithms.

Two cautions keep these findings honest. Redundancy at test time does not mean the heads were useless during training. Extra heads may make optimization easier, which is a reason not to start with fewer. And attention weights are not the same as importance. A head can attend strongly to a token whose value vector is tiny, so it contributes almost nothing (Kobayashi et al. 2020), which is why careful studies measure what removing or patching a head does instead of just staring at its weights.

import numpy as np

def head_entropy(A):
    raise NotImplementedError

In PyTorch

A production multi-head attention module is about fifteen lines. This is the shape of nanoGPT's CausalSelfAttention, using the fused kernel from §4.3.

The .contiguous() before the final .view is not decoration. After the transpose the tensor's memory is no longer laid out row by row, and view refuses to reinterpret non-contiguous memory. (reshape would copy silently instead.) The comprehension over (q, k, v) iterates over three tensors, not over heads, so it costs nothing that grows with $h$.

Read more

• Vaswani et al. (2017). Attention Is All You Need. Section 3.2.2 defines multi-head attention, and Table 3 varies the number of heads.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. Heads as independent, additive readers and writers of the residual stream.
• Clark et al. (2019). What Does BERT Look At? A survey of BERT's heads, from positional and separator heads to syntactic ones.
• Voita et al. (2019). Analyzing Multi-Head Self-Attention and Michel et al. (2019), Are Sixteen Heads Really Better than One? Specialized heads, and how many can be pruned.
• Bhojanapalli et al. (2020). Low-Rank Bottleneck in Multi-head Attention Models. Why tying the head size to $d/h$ limits each head, and what to do about it.
• Karpathy. nanoGPT. The reference implementation of causal multi-head self-attention that the PyTorch block above follows.

Watch next

▶

▶

▶