num_heads Is Not a Capacity Knob

Multi-head attention — MHA — is not several attention layers running side by side. It is one attention layer whose model dimension has been cut into h slices. The block holds the same number of weights and does the same number of multiplications whether you give it 1 head or 512. What the head count buys you is not size. It is how many different things the layer can pay attention to at once, and it is paid for out of how complicated any one of those things is allowed to be.

The panel is one attention layer of the original Transformer: d_model is 512, exactly as in Vaswani et al. Underneath it are four projection matrices — W_Q, W_K, W_V and the output projection W_O — and a job list. A job is an attention pattern the layer is required to produce: attend to the previous token, attend to the first token, and so on. Each job is handed to a head. If there are more jobs than heads, a head carries several and has to compromise between them.

Move heads across its whole range and watch two things at once: the parameter count in the top-left readout, and the bar chart. One of them never moves.

The weight count is 1,048,576 at 1 head, at 8 heads and at 512 heads. So is the projection arithmetic. The score and value arithmetic does not move either: at a context of 8,192 tokens it is 8,388,608 multiply- accumulates per token whatever the slider says. Vaswani et al. put this in one sentence in section 3.2.2: Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality. Every head you add makes every head narrower, and the two effects cancel exactly.

The bar chart does move, and not in the direction people expect. It rises from 0.471 at one head to 1.000 by four heads — and then falls back to 0.686 at 512. Nothing was made cheaper or more expensive anywhere along that curve.

Why the cost does not move

Multi-head attention is defined in section 3.2.2 of Vaswani et al. as MultiHead(Q,K,V) = Concat(head_1, …, head_h) W_O, where head_i = Attention(Q W_Q_i, K W_K_i, V W_V_i). The projections are per-head matrices: W_Q_i and W_K_i are d_model × d_k, W_V_i is d_model × d_v, and W_O is h·d_v × d_model. Then the choice that decides everything: In this work we employ h = 8 parallel attention layers, or heads. For each of these we use d_k = d_v = d_model/h = 64.

Put d_k = d_model/h into the shapes and the head count cancels. Stack the h query projections side by side and you have one d_model × (h · d_k) matrix, which is d_model × d_model. Same for keys and values. W_O is d_model × d_model the other way round. Four square matrices: 4 · d_model² weights, and with d_model = 512 that is 1,048,576 — the number the readout shows at every position of the slider. In a real implementation there is literally one fused matrix: vLLM builds a QKVParallelLinear of width total_num_heads · head_dim and reshapes it into heads afterwards. The heads are a view of a tensor, not a set of layers.

The score arithmetic cancels the same way. One head compares a query against n keys in d_k dimensions, so it costs n · d_k multiply-accumulates per query, and the value average costs another n · d_v. Across h heads that is 2 · n · h · d_k = 2 · n · d_model, with no h left in it. At 8,192 tokens the readout says 8,388,608 per token, and it says the same at 1 head and at 512.

This is worth holding next to the shape of the bill. At a context of 128 the score work is 131,072 multiply-accumulates per token against 1,048,576 for the projections — the projections dominate by 8 to 1. At 8,192 the score work is 8,388,608 and the ratio has inverted. That crossover, not the head count, is what makes long context expensive, and it is the subject of scaled dot-product attention and of FlashAttention, which is about the memory traffic of the same matmuls rather than their count.

One consequence people find surprising: because W_O is applied to the concatenation, and a matrix applied to a concatenation is the sum of the blocks applied to the pieces, the whole layer is Σ_i head_i · W_O_i — a plain sum over heads of independent per-head outputs. The concatenation is bookkeeping. That is why you can delete a head by zeroing its slice of W_O without touching anything else. Head pruning is therefore a masking operation, not surgery: Michel, Levy and Neubig switch each head off with a variable in front of its term in that sum, and report in Are Sixteen Heads Really Better than One? that a large percentage of attention heads can be removed at test time without significantly impacting performance and that some layers can even be reduced to a single head. It is also why tensor parallelism splits attention by head and finishes with one all-reduce: each GPU computes part of that sum.

What a head cannot express

Turn every job off except attend to the previous token and walk the slider up. The score is 1.00 all the way to 32 heads. At 64 it is 0.774, at 128 it is 0.487, and at 512 heads it is 0.302. Now switch to attend to the first token alone and do it again: 1.00 everywhere, including at 512 heads, where each head is one dimension wide.

The same head, the same width, and one job is destroyed while the other is untouched. The reason is a rank ceiling, and it is exact rather than approximate.

A head's pre-softmax score matrix is S = (X W_Q)(X W_K)ᵀ for a sequence X of n tokens. The left factor is n × d_head and the right is d_head × n, so S is an n×n matrix that is the product of two thin ones. Its rank cannot exceed d_head. That is the whole constraint, and it is the one Bhojanapalli et al. named the low-rank bottleneck: their paper's diagnosis is that the scaling between the number of heads and the size of each head in the current architecture gives rise to a low-rank bottleneck in attention heads, and their proposed fix is to set the head size to the input sequence length and stop tying it to d_model at all.

Take the extreme case, because it is provable in one line. If d_head = 1 then S[t][s] = a_t · b_s for two vectors. Row t of the attention is softmax(a_t · b) over the visible positions — the same fixed ranking of keys b for every query, just sharper or flatter depending on the size of a_t, and reversed if a_t is negative. A one-dimensional head has one opinion about which positions matter and can only turn its volume up and down. Attend to the first token is one opinion, held by every query, so it survives perfectly. Attend to the previous token requires the argmax to move with the query, and no rank-one matrix can do that at all. The 0.302 the panel reports is not a rounding error; it is a head reduced to attending broadly because the pattern it was asked for is not in its reach.

The third job is the interesting middle. Leave it as the only job ticked — untick attend to the first token when you tick it, or you will be reading the average of two jobs rather than this one. Attend to tokens with the same identity over these twelve tokens is a sum of six rank-one blocks, one per distinct token identity, so it is a rank-6 pattern. It scores 1.00 down to head_dim of 8, then 0.858 at rank 4 and 0.444 at rank 1. A head is not simply good or bad at its width. It is exactly as good as the rank of what you asked it for.

Two honest limits on this demonstration. The approximation shown is the truncated singular value decomposition of the target scores, which minimises error before the softmax rather than after it, so a cleverer choice of weights could score somewhat higher than the panel shows. And a real head is further constrained: its S must lie in the span of the actual token representations, which the panel ignores. The rank ceiling itself is not an approximation. Nothing a head can do escapes it.

Which brings the other end of the slider into focus. Tick all four jobs back on and take the slider to its far left. With four jobs and one head, the mean score is 0.471 even though that head is 512 dimensions wide and could express any of the four patterns perfectly on its own. It has one score matrix and four things to say with it, so it says their average, and the softmax of an average of spikes is a spread. This is the sentence in the paper that everybody quotes and few people test: Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.

So the curve has two different causes at its two ends. On the left, too few score matrices for the number of jobs. On the right, each score matrix too low-rank for the job it holds. Between them, at four heads with four jobs, a plateau of exactly 1.00 that runs to 32 heads and buys nothing further — the idle-head count in the log rises to 28 there, which is the state Michel et al. measured in trained models and then pruned away.

Where the invariance breaks

The parameter count is constant because of the choice d_head = d_model/h, and that choice is a convention, not a law. Tick Pin head_dim at 128 and the hero readout starts moving: 262,144 weights at one head, 2,097,152 at eight, 134,217,728 at 512. With the head size fixed, adding a head really does add 4 · d_model · 128 weights, because the stacked projections are now d_model × (h · 128) and no longer square.

This is not hypothetical. vLLM's Llama attention reads self.head_dim = head_dim or self.hidden_size // self.total_num_heads — an explicit head_dim in the config wins, and the fallback is only a default. Qwen3-4B's config.json sets hidden_size 2560, num_attention_heads 32 and head_dim 128, so its query projection is 2560 × 4096 and its output projection is 4096 × 2560. Thirty-two times one hundred and twenty-eight is not 2560, and it was not meant to be. Qwen3-8B sets the same head_dim of 128 with hidden_size 4096, where the two happen to agree. If you are reasoning about a model's attention cost from hidden_size / num_attention_heads, check whether the config overrode it before you trust the arithmetic.

The second place the head count stops being free is deployment. Heads are the unit that tensor parallelism splits on, and vLLM's Llama implementation states it as an assertion: assert self.total_num_heads % tp_size == 0. A 64-head model runs on 1, 2, 4, 8, 16, 32 or 64 GPUs and on nothing else — which is why tensor parallelism with a non-divisible amount of attention heads is a feature request people file after buying three GPUs. Key/value heads are counted separately and get their own rule in the same constructor: when there are fewer of them than there are GPUs they are replicated rather than split, which is the arithmetic that grouped-query attention introduces and this lesson deliberately leaves alone.

The third is the kernels. head_dim is a tile dimension in every fast attention implementation, which is why 64 and 128 are everywhere and why unsupported head_dim is one of the commonest issue titles on GitHub. You are not free to pick 96 and expect the fused path.

None of this makes the head count a capacity dial. It makes it a shape constraint with three separate sets of teeth: the rank ceiling in the maths, the divisibility rule in the parallel layout, and the tile size in the kernel.

Checking it in a real system

1. Read the three numbers together, never one of them. hidden_size, num_attention_heads and head_dim. If the config has no head_dim, it is hidden_size / num_attention_heads and the product is square. If it has one, multiply it out and see whether it agrees. That single check tells you whether the attention block is 4 · d_model² weights or something larger.

2. Count the parameters, do not estimate them. For a decoder layer, sum the shapes of q_proj, k_proj, v_proj and o_proj. In PyTorch that is four lines over model.named_parameters(). Change num_attention_heads in a config you control, rebuild the model, and count again. If the total moved, your model sets head_dim explicitly. If it did not, you have just watched the invariance this lesson is about.

3. When a kernel refuses, read which dimension it named. Unsupported head_dim: 160 is not about your head count and not about your model size; it is about hidden_size / num_attention_heads landing somewhere the fused kernel has no tile for. The fix is a different head count that produces a supported head size, or a fallback path — and the fallback is usually the slow one, which is worth knowing before you conclude the model is inherently slow.

4. Before buying GPUs, factorise the head count. num_attention_heads and num_key_value_heads together decide which tensor-parallel widths exist. Three GPUs will not serve a 64-head model no matter how much memory they have. Check this at procurement, not at deployment.

5. If you are tempted to raise the head count for quality, price the alternative first. Nothing about the layer gets bigger, so whatever improvement you are hoping for has to come from the layer being able to hold more distinct patterns — and it costs you rank on each of them. If your sequences are long, head_dim is already far below the sequence length and you are on the falling side of the curve. Raising d_model, or decoupling head_dim from it the way Bhojanapalli et al. proposed and Qwen3-4B implements, are the changes that actually add capacity. Note that the second one does add parameters, which is the whole point of doing it deliberately rather than by accident.