A single attention head computes one weighted average per position: one set of compatibility scores, one softmax, one blend of values. That is expressive, but it forces every relationship in the sequence — syntactic agreement, coreference, positional adjacency, topical similarity — through one d-dimensional lens at once. Multi-head attention is the fix, and it is almost embarrassingly simple: instead of one head of width d_model, run h heads of width d_k = d_model / h in parallel, each with its own small projections, then concatenate their outputs and mix them with one final linear layer W_O. The surprise is that this buys the model h distinct representation subspaces for essentially the same parameter count and the same FLOPs as a single full-width head. This article works through the whole mechanism from the shapes up: the split, the per-head projections, the parallel scaled-dot-product attention, the concat-and-project reassembly, exact parameter and compute counts, a fully worked numeric example, and the CPU-and-small-model implications — including where the ‘free lunch’ framing quietly breaks down.
Why one head is not enough
The limitation of a single head is not capacity in the crude sense — it is averaging. Attention produces, for each query position, a convex combination of value vectors weighted by a single softmax distribution. That means one head can express exactly one ‘view’ of which other tokens matter. If the word bank needs to attend to river for its sense and to the for its determiner and to a verb three tokens away for agreement, a lone softmax must somehow blend all of those into one distribution, and the resulting averaged value smears distinct signals together.
Multiple heads let the model attend to different things in different representation subspaces simultaneously. Each head gets its own low-dimensional projection of the token vectors, so it can specialise: one head can learn to track the previous token, another to resolve pronouns, another to match verbs with subjects. Because the heads run in parallel and are combined afterward, the model does not have to choose one relationship at the expense of the others. The original transformer paper motivated exactly this: ‘multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions,’ something a single head, with its one averaging operation, actively prevents.
The dimensional bookkeeping: d_model, h, and d_k
Everything in multi-head attention hangs on one division. Let d_model be the model’s embedding width — the size of the vector that flows along the residual stream. Choose a head count h that divides it evenly, and define the per-head dimension:
d_k = d_v = d_model / hSo with d_model = 512 and h = 8, each head works in a d_k = 64-dimensional space. The key design decision is that the heads partition the width rather than replicate it: eight heads of 64 add back up to 512, not to 8×512. This is what keeps the parameter and compute budget flat as you change h. It also means there is a genuine tension in choosing h: more heads give more independent subspaces (more distinct relationships the layer can track at once), but each head becomes narrower and individually less expressive, since d_k shrinks. Typical configurations keep d_k in the 64–128 range — GPT-style models often use d_k = 64 or 128 and scale h with d_model. The equality h × d_k = d_model is the invariant to hold in your head for the rest of this article.
Per-head projections: W_Q, W_K, W_V
A single head turns the input X: [N, d_model] into queries, keys, and values with three learned matrices. In the multi-head case, head i has its own trio:
W_Q^i : [d_model, d_k]
W_K^i : [d_model, d_k]
W_V^i : [d_model, d_v]
Q_i = X · W_Q^i → [N, d_k]
K_i = X · W_K^i → [N, d_k]
V_i = X · W_V^i → [N, d_v]Each projection reads the full d_model-dimensional token and compresses it down to a d_k-dimensional slice tailored to that head. This is the crucial point about subspaces: the heads do not each receive a hard-partitioned chunk of the input vector; every head sees all of X and learns its own linear projection into a subspace it finds useful. Two heads can look at overlapping information but emphasise different directions. In an implementation you never store h separate matrices — you stack them. The concatenation [W_Q^0 | W_Q^1 | … | W_Q^{h-1}] is one matrix of shape [d_model, h·d_k] = [d_model, d_model], so a single matmul X · W_Q produces every head’s queries at once. The per-head view is the math; the fused matrix is the code.
Scaled dot-product attention, one head at a time
Given a head’s Q_i, K_i, V_i, the attention itself is the standard scaled dot-product operation covered in its own article — here it is only the per-head building block:
head_i = softmax( (Q_i · K_i^T) / √d_k ) · V_i → [N, d_v]The one detail worth flagging is the scale factor: it is √d_k, the per-head dimension, not √d_model. That is correct and important — the dot products Q_i · K_i^T sum over d_k terms, so their variance grows with d_k, and dividing by √d_k keeps the pre-softmax logits at unit scale regardless of how you chose h. Every head runs this identical computation on its own slice, completely independently: there is no interaction between heads inside the attention step. Head 3’s softmax cannot see head 5’s scores. That independence is exactly what makes the heads parallelisable and what lets them specialise into different relationships — the mixing between them is deferred entirely to the output projection, which we come to shortly. Until then, think of h separate attention operations proceeding side by side.
Shapes end-to-end: [N, d_model] to heads and back
The whole layer is a round trip out of the residual stream and back into it. Tracing the shapes is the fastest way to internalise the mechanism:
X : [N, d_model] input tokens
Q,K,V : [N, d_model] one fused projection each
reshape : [N, h, d_k] split the width into heads
transpose : [h, N, d_k] heads become a batch dim
scores : [h, N, N] Q · K^T per head
softmax/AV : [h, N, d_v] per-head attention output
transpose : [N, h, d_v] heads back beside each other
concat : [N, d_model] h·d_v = d_model
· W_O : [N, d_model] final mixed outputNotice that the tensor enters as [N, d_model] and leaves as [N, d_model] — multi-head attention is shape-preserving, which is what lets you stack dozens of these layers with a residual connection around each. The head dimension appears only in the middle, where the reshape and transpose turn the width into a batch axis so that a single batched matmul does all h heads at once. No numbers are lost or duplicated in the reshape; it is a pure reinterpretation of the same N × d_model values as N × h × d_k.
Concatenation and the output projection W_O
After the parallel attention step you hold h output slices, each [N, d_v]. Concatenating them along the feature axis rebuilds a full-width tensor, and one more learned matrix finishes the layer:
concat(head_0, …, head_{h-1}) → [N, h·d_v] = [N, d_model]
MultiHead(X) = concat(…) · W_O , W_O : [d_model, d_model] → [N, d_model]W_O is not a formality — it is load-bearing. Straight after concatenation, dimensions 0..d_v-1 of the vector belong exclusively to head 0, dimensions d_v..2d_v-1 to head 1, and so on. The heads have produced their answers in isolation and those answers sit in disjoint blocks. W_O is the first and only place they are mixed: every output dimension becomes a learned linear combination of all heads’ contributions. Without it, the heads’ outputs would pass into the residual stream as unmixed, block-structured features and the model could never learn that, say, head 2’s coreference signal should modulate head 6’s topical signal. You can even read W_O as a stack of per-head output matrices W_O^i : [d_v, d_model] so that MultiHead(X) = Σ_i head_i · W_O^i — the layer is literally a sum of per-head contributions back into the full model width.
The batched implementation: one matmul, not a loop
Written as math the heads look like an explicit for loop, but no real implementation runs one. The four projections are fused into [d_model, d_model] matrices, so computing Q, K, and V is three big matmuls that already contain every head’s output interleaved. The split into heads is then a zero-copy reshape from [N, d_model] to [N, h, d_k] followed by a transpose to [h, N, d_k], which promotes the head index to a batch dimension.
From there the entire multi-head attention is two batched matrix multiplies: Q · K^T over the batch of h heads gives scores: [h, N, N], and after the scaled softmax, scores · V gives [h, N, d_v]. A transpose and reshape concatenate the heads back to [N, d_model], and a final matmul applies W_O. This is why multi-head attention costs essentially nothing extra to implement over single-head: on a GPU or a well-tuned CPU kernel it is the same handful of GEMMs, just with an extra leading batch axis of size h. The head structure is a way of reshaping the same arithmetic, not adding to it.
Counting the parameters
The learned weights of a multi-head attention layer are exactly four matrices (plus optional biases). Because the per-head projections stack into full-width matrices, the head count h vanishes from the count:
W_Q : d_model × (h·d_k) = d_model × d_model
W_K : d_model × d_model
W_V : d_model × d_model
W_O : d_model × d_model
------------------------------------------
params = 4 · d_model^2 (+ 4·d_model biases)Read that carefully: the parameter count is 4·d_model^2 regardless of how many heads you use. Splitting into 8 heads, 16 heads, or running a single full-width head all cost the same number of weights, because h·d_k is pinned to d_model. This is the first half of the ‘multiple heads are nearly free’ story. The choice of h changes only how those 4·d_model^2 parameters are carved into subspaces — it reshapes the projection matrices’ effective structure without adding or removing a single weight. Biases, when present, add a negligible 4·d_model terms. For the whole transformer block you would add the feed-forward network’s parameters (typically 8·d_model^2 for a 4× expansion), which usually dominates, but that is a separate story from the attention sublayer itself.
Counting the FLOPs
Compute splits into two parts: the projections and the attention core. For a sequence of length N, counting multiply-accumulate operations (MACs):
Projections (Q,K,V,O): 4 · N · d_model^2 MACs
Scores Q·K^T per head: h · N^2 · d_k = N^2 · d_model MACs
AV scores·V per head: h · N^2 · d_v = N^2 · d_model MACs
------------------------------------------------------------
total ≈ 4·N·d_model^2 + 2·N^2·d_model MACsThe critical line is the middle one. The per-head scores are N × N × d_k, and there are h of them, so summed over heads the cost is h · N^2 · d_k = N^2 · d_model — the h cancels against the 1/h hidden in d_k. The attention core therefore costs the same whether you use one head or many. Multiply MACs by two for FLOPs. Two regimes fall out: when N < d_model the projection term 4·N·d_model^2 dominates (attention is ‘cheap’, the layer is basically four linear layers); when N > d_model the quadratic 2·N^2·d_model term takes over and long context becomes the bottleneck. The crossover sits right around N ≈ 2·d_model.
Why multi-head is roughly free compute
Put the two counts together and the ‘free lunch’ becomes precise. Compare h heads of width d_k = d_model/h against a single hypothetical head of the full width d_model. The projection matrices are d_model × d_model in both cases — identical. The attention core is N^2 · d_model MACs in both cases, because summing h heads of size d_k reconstructs exactly d_model. So the two FLOP-dominant terms are not approximately equal, they are exactly equal.
Where does the ‘roughly’ come from, then? One genuine difference: the softmax now normalises h attention maps of N×N entries each — h·N^2 exponentials versus N^2 for a single head. That is h× more softmax work, but softmax has no ×d factor, so it is a lower-order term dwarfed by the matmuls at any realistic dimension. There is also a little extra memory-movement for the reshape/transpose. Net: multi-head attention delivers h independent representation subspaces for the same parameters and essentially the same FLOPs as one wide head — you gain expressive structure and lose only a negligible softmax and reshaping overhead. That is the whole reason the design is universal.
A fully worked example
Make it concrete with d_model = 512, h = 8, so d_k = d_v = 64, on a sequence of N = 100 tokens.
Shapes
X [100, 512]
Q = X·W_Q [100, 512] reshape → [100, 8, 64] → [8, 100, 64]
scores [8, 100, 100] (per head: 100×100)
head outputs [8, 100, 64] → concat → [100, 512]
· W_O [100, 512]
Parameters
each of W_Q,W_K,W_V,W_O = 512×512 = 262,144
total = 4 × 262,144 = 1,048,576 ≈ 1.05 M weights
Compute (MACs, N = 100)
projections = 4 · 100 · 512^2 = 104,857,600
attention = 2 · 100^2 · 512 = 10,240,000
total ≈ 115.1 M MACs ≈ 230 M FLOPsAt N = 100, well below d_model = 512, the projections are about 10× the attention cost — short sequences are projection-bound. Now push to N = 2000: projections scale linearly to 4·2000·512^2 ≈ 2.1 B MACs, but attention scales quadratically to 2·2000^2·512 ≈ 4.1 B MACs and overtakes them. Same 1.05 M parameters throughout; only the compute mix shifts with context length.
What the heads actually learn
The subspace story is not just motivation — it shows up in trained models. Probing studies of transformers find heads that specialise into interpretable roles: positional heads that attend to the previous or next token, syntactic heads that link verbs to their subjects or objects, coreference heads that connect a pronoun to its antecedent, rare-token or ‘induction’ heads that copy patterns seen earlier in the context, and heads that mostly attend to a delimiter or the first token as a no-op ‘rest’ position. None of this is hand-designed; it emerges because giving the model h parallel channels lets gradient descent allocate different relationships to different heads.
The flip side is redundancy. Research pruning attention heads has repeatedly shown that many heads can be removed after training with little accuracy loss — sometimes a majority — which tells you the h subspaces are not all doing distinct, essential work. This is a hint that head count is often set generously for training dynamics rather than because every head earns its keep at inference, and it is precisely the slack that later efficiency tricks exploit.
Implications for CPU and small models
On a CPU-hosted small language model, the multi-head structure interacts with the two things you care about most: cache behaviour and the KV cache. Because heads are a batch axis over small d_k-wide matmuls, they map to tidy, cache-friendly GEMM tiles — but a very large h with tiny d_k can leave individual matmuls too small to saturate SIMD lanes, so there is a practical floor on useful head width. More importantly, during autoregressive decode you cache the keys and values of every past token, and that cache scales as 2 · N · d_model · layers — proportional to h · d_k, i.e. the full width per layer.
This is what motivates multi-query and grouped-query attention (MQA/GQA): keep h query heads for expressiveness but share one, or a few, key/value heads across them. The KV cache — the dominant memory cost of long-context decode on a memory-bandwidth-bound CPU — then shrinks by up to h×, at a small quality cost. Understanding that the h in multi-head lives independently in the query and the key/value paths is what makes those optimisations legible.
Common pitfalls and misconceptions
A handful of errors recur when people first work through this math. First: scaling by √d_model instead of √d_k. The dot products live in the per-head space, so the scale is the per-head dimension; using the full width over-damps the logits. Second: believing more heads means more parameters or more FLOPs. As shown, both are pinned by d_model; h only repartitions them. Third: thinking each head sees a fixed slice of the input vector. Every head reads all of X; the ‘slice’ is in the output of the projections, not a chopping of the input.
Fourth: dropping W_O, or treating it as optional. Without it the heads never communicate and the concatenated blocks enter the residual stream unmixed — the layer is crippled. Fifth: assuming heads are truly independent computations you could split across devices for free; they share the same input projections and are re-mixed by W_O, so the independence lives only inside the softmax-and-AV core. Keeping these straight is the difference between reciting the formula and actually understanding why the layer is shaped the way it is.
Putting it together
Step back and the whole layer is one clean idea expressed in linear algebra. Take the token stream X: [N, d_model]; project it three ways with fused d_model × d_model matrices; reinterpret the width as h heads of size d_k = d_model/h; run scaled dot-product attention independently and in parallel in each head’s subspace; concatenate the results back to full width; and mix them with the output projection W_O so the heads can finally talk to each other. In, [N, d_model]; out, [N, d_model]; four weight matrices; 4·d_model^2 parameters; 4·N·d_model^2 + 2·N^2·d_model MACs.
The elegance is that the head count h is a free structural knob: it does not touch the parameter budget and barely touches the FLOP budget, yet it decides how many distinct relationships the layer can attend to at once. That is why every modern transformer uses it, and why the interesting engineering — head pruning, multi-query and grouped-query attention, sparse and windowed variants — all works by re-carving this same fixed budget rather than by spending more of it.
d_model into h heads of size d_k = d_model/h, gives each head its own W_Q/W_K/W_V projection into a distinct representation subspace, runs scaled dot-product attention independently and in parallel per head, then concatenates the outputs and mixes them with a single output projection W_O — the one place the heads communicate. The payoff is structural: because h × d_k = d_model is held fixed, the layer costs 4·d_model^2 parameters and roughly 4·N·d_model^2 + 2·N^2·d_model MACs no matter how many heads you use. The projection matmuls and the attention core are exactly as expensive as a single full-width head; only the softmax (now h maps) and a little reshaping differ, which is why total compute is roughly, not exactly, the same. You get h parallel views of the sequence for the price of one wide one — and the head count becomes a free knob that later tricks like grouped-query attention re-carve rather than enlarge.