Multi-head attention

1

Concept

A single attention operation creates one score distribution for each query and returns one weighted mixture of values. That is expressive, but it asks one set of projections to serve every useful relationship. Multi-head attention creates several attention operations in parallel. Each head receives its own learned query, key, and value subspace, allowing the layer to construct multiple routing patterns before combining them.

Let the model width be $d_{model}$ and the number of heads be $h$. A common design chooses head width $d_{head}=d_{model}/h$. Implementations often use three large linear projections, then reshape rather than instantiate separate modules for every head:

Here $B$ is batch size and $T$ is sequence length. Attention runs independently across the head axis:

The mask $M$ can enforce causality or padding rules. The head outputs are transposed and concatenated back into shape $B\times T\times d_{model}$, then multiplied by an output matrix $W_O$. This final projection is important: it lets the layer mix features delivered by different heads before adding the result to the residual stream.

Why split the width? Suppose one query needs information about nearby punctuation, a coreferent noun, and the start of a quoted span. Separate heads can form different similarity spaces, so “relevant” need not mean the same thing in each. One head can emphasize relative position while another emphasizes content features. Later layers can compose these routes into more complex behavior.

That description is a capability, not a promise of neat interpretability. Research has found that many trained heads can be pruned in particular models and tasks with limited immediate damage, while a smaller subset is important. Redundancy can be useful for optimization or robustness, and a meaningful computation can also be distributed across heads and layers. Naming a head “the syntax head” from a few examples is a hypothesis that needs causal testing, not an architectural fact.

The parameter count has a subtle property. If total model width stays fixed and $h d_{head}=d_{model}$, the large $Q$, $K$, $V$, and output projections have roughly the same parameter count across different head counts. More heads do not automatically mean proportionally more parameters; they change how the projected width is partitioned. They can still change runtime behavior, kernel efficiency, score-matrix storage, and representational bottlenecks.

Head count also constrains dimensions. The model width must be divisible by the chosen head count in the usual layout. Very narrow heads may lose useful per-head capacity. Very wide heads reduce the number of independent routing patterns. Modern architectures sometimes use different numbers of query heads and key/value heads for faster inference; those MQA and GQA variants preserve multiple query routes while sharing cached keys and values.

From an implementation perspective, axis order is a common source of bugs. The score multiplication must pair every query position with key positions inside the same batch item and head. The softmax must run over the key-position axis. After applying weights to values, the code must restore token order before concatenation. Shapes are not bookkeeping around the algorithm; they express the algorithm.

The enduring picture is a bank of learned communication channels. Each head decides where to read in its own coordinate system, produces a contextual message, and hands that message to a shared output mixer. The architecture offers specialization, redundancy, and composition while retaining efficient batched matrix operations.