Multi-Head Attention as Independent Subspace Projections

A single attention head with width $d_v = d_{\text{model}}$ is, on paper, at least as expressive as any collection of narrower heads, because the wider head’s key and value matrices can be written as the concatenation of the narrower ones. The scores would differ, though: a wide head computes one softmax over the sum of all the per-head score matrices, while separate heads compute a softmax per score matrix and only then combine. That difference, namely where the normalization sits, is the entire justification for the multi-head construction, and it explains most of what follows.

The concatenation is a block-diagonal structure

The layer is defined by $h$ triples of projections. Head $i$ has $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$, and produces

$$\text{head}_i = \mathrm{softmax}!\left(\frac{X W_i^Q (X W_i^K)^\top}{\sqrt{d_k}}\right) X W_i^V$$

The $h$ results are concatenated along the feature axis to give a $n \times h d_v$ matrix, then multiplied by an output projection $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$. With the usual choice $d_k = d_v = d_{\text{model}} / h$, the concatenation has exactly $d_{\text{model}}$ columns and the shapes line up without any reduction step.

Writing the whole thing as one matrix product exposes the structure. Stack the per-head query projections into $W^Q = [W_1^Q; \dots; W_h^Q] \in \mathbb{R}^{d_{\text{model}} \times h d_k}$, and likewise $W^K$ and $W^V$. Then $X W^Q$ is a $n \times h d_k$ matrix that can be reshaped to $n \times h \times d_k$, and the scores become an $h \times n \times n$ tensor. The reshape is not cosmetic: it is the statement that the $h$ score matrices are computed from disjoint slices of the projected features and never interact before the softmax.

The output projection $W^O$ is where the heads meet. It can be partitioned into $h$ blocks $W^O_i \in \mathbb{R}^{d_v \times d_{\text{model}}}$, and the layer output is $\sum_i \text{head}i W^O_i$. So the heads are summed, not merely concatenated, and the sum happens in the $d{\text{model}}$-dimensional output space. Each head contributes a rank-$d_v$ term, and the sum of $h$ such terms has rank at most $\min(d_{\text{model}}, h d_v) = d_{\text{model}}$, no constraint at all in the typical case. That is a deliberate design property: nothing about the layer caps the rank of what it can write back to the residual stream.

Where a single head loses

Suppose instead that one head of width $d_{\text{model}}$ is used, with $W^Q$, $W^K$, $W^V$ each $d_{\text{model}} \times d_{\text{model}}$. The pre-softmax score matrix is $QK^\top/\sqrt{d_{\text{model}}}$, and since $Q = XW^Q$ and $K = XW^K$, this is $X W^Q W^{K\top} X^\top / \sqrt{d_{\text{model}}}$. The middle factor $W^Q W^{K\top}$ is a $d_{\text{model}} \times d_{\text{model}}$ matrix of rank at most $d_{\text{model}}$, so the scores have rank at most $d_{\text{model}}$, again no constraint. But every entry of that matrix is a single softmax argument, and all $n^2$ entries compete through one normalization per row.

The consequence is a competition for probability mass. A row of the score matrix must distribute a total of 1 across all $n$ positions. If the layer needs to attend to a previous-token position, a syntactic-head position, and a coreference position simultaneously, those three logits must all be high relative to everything else, and the softmax divides the mass among them. With $h$ separate heads, each row of each head’s score matrix distributes its own unit of mass, so three different heads can each place most of their mass on a different position. The wide single head can represent the same three-way split only by assigning each a probability near $1/3$, and the resulting output is a fixed mixture of three value vectors rather than three separate reads that the output projection can weight independently.

That last clause is the part that matters. Because $W^O$ is applied after the softmax, separate heads give the output projection a free hand in weighting the reads: it can take head 1’s read at full strength and suppress head 3’s entirely. A single head’s mixture is already committed before $W^O$ sees it, and $W^O$ can only apply one linear map to the mixture. The representational difference is real and it is not recoverable by widening the single head.

There is a second, softer argument from optimization rather than expressivity. Splitting the projections into $h$ independent matrices means $h$ independent gradient paths through the score computation, each normalized by its own softmax Jacobian. When one head saturates, its gradient vanishes locally but the others keep receiving signal, so the layer as a whole does not stall. A single wide head has one saturation event that affects the whole layer.

Rank, entropy, and what the heads actually specialize to

Empirical analyses of trained models converge on a consistent picture, and it is worth stating without overclaiming what the numbers mean. A probing study of trained translation models found that a small number of heads carry identifiable roles, such as attending to the position immediately before the current token, to the least frequent token in the source, or to syntactically dependent positions, while a larger fraction could be pruned with a modest loss in translation quality. One such measure combines the gradient of the loss with respect to a head’s output and the magnitude of that output, computed on a development set, and heads are removed in order of increasing score.

The relevant caveat is that importance is context-dependent. A head that is dispensable when pruned alone may be load-bearing when another head has already been removed, because the two were redundant rather than useless. Pruning studies that remove heads greedily and then measure report much smaller losses than the naive reading of per-head scores would suggest; the ordering matters, and the final count of removable heads is a property of the whole set rather than of individual heads.

Another line of analysis looks at the attention matrices themselves rather than at performance. A head whose attention distribution is nearly uniform over the row contributes a value that is close to the mean of all value vectors, which is a low-information operation that the output projection can approximate cheaply with a constant term. Heads with sharply peaked distributions contribute a specific token’s value. The distribution of row entropies across heads in a trained layer is usually broad: some heads sit near maximal entropy, some near zero. That spread is consistent with the design argument above, since the layer gets to use different temperatures and different selectivity per head without any explicit mechanism to enforce it.

import numpy as np

def split_heads(x, h):
    # x: (batch, seq, d_model) -> (batch, h, seq, d_k)
    b, n, d = x.shape
    return x.reshape(b, n, h, d // h).transpose(0, 2, 1, 3)

def attention(q, k, v, mask=None):
    d_k = q.shape[-1]
    scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(d_k)   # (b, h, n, n)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    scores = scores - scores.max(axis=-1, keepdims=True)   # shift before exp
    p = np.exp(scores)
    p /= p.sum(axis=-1, keepdims=True)
    return p @ v                                          # (b, h, n, d_v)

def multi_head(x, wq, wk, wv, wo, h):
    q, k, v = split_heads(x @ wq, h), split_heads(x @ wk, h), split_heads(x @ wv, h)
    out = attention(q, k, v)                              # heads stay separate
    b, _, n, d_v = out.shape
    return out.transpose(0, 2, 1, 3).reshape(b, n, h * d_v) @ wo

Head count and head width are not independent knobs

Fixing $d_{\text{model}}$ and increasing $h$ shrinks $d_k = d_{\text{model}}/h$. At $d_{\text{model}} = 512$ with 8 heads, $d_k = 64$; with 32 heads, $d_k = 16$. The per-head score matrix $Q_i K_i^\top$ has rank at most $d_k$, so at 16 dimensions the head cannot represent more than 16 independent directions of query-key interaction, and the softmax over $n$ positions is driven by a 16-dimensional inner product. Two distinct keys whose projections differ only in the remaining $d_{\text{model}} - 16$ directions become indistinguishable to that head.

The variance argument for the $1/\sqrt{d_k}$ divisor also changes with $h$: at fixed input scale, smaller $d_k$ means smaller logit spread and a softer softmax. Whether that is desirable depends on the task, and it is one reason head count is usually chosen together with model width rather than tuned in isolation. Common configurations keep $d_k$ in a band from roughly 64 to 128 and scale $h$ in proportion to $d_{\text{model}}$, which holds both the softmax temperature and the per-head rank roughly constant as the model grows.

The other constraint is the $n \times n$ score tensor. Its size grows with $h$ and quadratically with sequence length, and it is the tensor that dominates attention’s memory traffic; flash attention and IO-aware tiling describes the reorganization that avoids materializing it. When heads share key and value projections instead of having their own, the cache cost drops but the per-head score matrices remain, which is the trade examined in grouped-query attention and KV cache sharing.

A final structural point concerns what happens when heads are averaged rather than concatenated, a variant that appears in some architectures to save parameters. Averaging $h$ head outputs gives a $n \times d_v$ matrix whose rank is at most $d_v = d_{\text{model}}/h$, so the layer’s write-back to the residual stream is rank-deficient by a factor of $h$. Concatenation followed by $W^O$ avoids this exactly because the sum happens after each head has been lifted back into $d_{\text{model}}$ dimensions. The distinction is easy to miss when reading shape-annotated diagrams, since both variants end with a $d_{\text{model}}$-wide output.