Multi-Head Attention › The Complete Formula

Now that we've walked through each step conceptually, let's see how they combine into a complete formulation. Understanding this picture helps clarify how multi-head attention works: by running multiple specialized attention operations in parallel and learning to combine their insights.

The outer structure of multi-head attention takes the outputs from all heads and merges them:

$$\text{MultiHead}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \mathbf{W}^O$$

This formula reads from inside out: first compute each head, then concatenate their outputs, then apply the output projection. But what exactly is each head doing? Each head applies its own learned projections before computing attention:

$$\text{head}_i = \text{Attention}(\mathbf{Q} \mathbf{W}^Q_i, \mathbf{K} \mathbf{W}^K_i, \mathbf{V} \mathbf{W}^V_i)$$

where:

• $\mathbf{Q}, \mathbf{K}, \mathbf{V}$: the QKV inputs (in self-attention, all three equal the input $\mathbf{X}$)
• $\mathbf{W}^Q_i, \mathbf{W}^K_i, \mathbf{W}^V_i$: the per-head projection matrices for head $i$, each transforming the full $d_{\text{model}}$-dimensional input into a smaller $d_k$-dimensional subspace
• $\text{Attention}(\cdot)$: the scaled dot-product attention function from the previous chapter
• $\mathbf{W}^O$: the output projection that learns how to combine all heads' contributions

The design is flexible. In self-attention, $\mathbf{Q} = \mathbf{K} = \mathbf{V} = \mathbf{X}$, so the same input sequence projects to all three roles. But the same multi-head attention mechanism also supports cross-attention in encoder-decoder models, where queries come from the decoder and keys/values come from the encoder. The mathematical structure remains identical; only the inputs differ.

Why should projecting to lower-dimensional subspaces and computing attention there be useful? Think of it this way: the full $d_{\text{model}}$-dimensional embedding space contains many different types of information about each token. Its syntactic role, semantic meaning, position in the sentence, relationship to neighboring words. A single attention operation in the full space must somehow balance all of these considerations when deciding what to attend to.

By projecting to a lower-dimensional subspace, each head can focus on a particular "slice" of the information. One head's projection matrices might learn to emphasize dimensions that capture syntactic structure, causing that head to attend based on grammatical relationships. Another head's projections might emphasize semantic dimensions, causing it to attend based on meaning similarity. The output projection then learns to weigh and combine these different perspectives.

This is analogous to how ensemble methods work in machine learning: multiple weak learners, each capturing part of the pattern, combine to form a stronger predictor than any single model could achieve.

There is a subtle but important distinction between multi-head attention and a simple ensemble of independent attention modules. In an ensemble, each member operates on the same full representation and the combination is usually an average or a vote. In multi-head attention, the combination is a learned linear transformation through $\mathbf{W}^O$. This output projection is not a passive merger; it actively learns how to mix head contributions in a context-dependent way. Heads that are more relevant for a given input can have their contributions amplified, while heads that detected spurious patterns can be suppressed. The projection also allows interactions between head outputs, so the final representation can encode relationships that no single head captured on its own.

In practice, this means you should think of the output projection as a learned integration step rather than a simple concatenation. The concatenation preserves all head outputs separately so the projection can distinguish between them. If you replaced concatenation with summation, you would lose the ability for the projection to treat different heads differently, and you would also lose the dimensional expansion that gives the projection enough freedom to mix contributions flexibly.

A natural question arises: if we're running $h$ separate attention operations, doesn't that multiply our parameter count by $h$? No. The dimension splitting ensures that multi-head attention uses approximately the same number of parameters as single-head attention with the full dimension.

Let's trace through the calculation:

1. Per head: Each head has three projection matrices ($\mathbf{W}^Q_i$, $\mathbf{W}^K_i$, $\mathbf{W}^V_i$), each of size $d_{\text{model}} \times d_k$. That's $3 \times d_{\text{model}} \times d_k$ parameters per head.

2. All heads: With $h$ heads, we multiply by $h$: total QKV parameters are $h \times 3 \times d_{\text{model}} \times d_k$. Here's the key insight: since $d_k = d_{\text{model}} / h$, we have $h \times d_k = d_{\text{model}}$. This simplifies our expression to $3 \times d_{\text{model}}^2$.

3. Output projection: The matrix $\mathbf{W}^O$ has size $(h \cdot d_v) \times d_{\text{model}}$. With $d_v = d_k = d_{\text{model}} / h$, the input dimension is $h \cdot (d_{\text{model}} / h) = d_{\text{model}}$, giving us $d_{\text{model}} \times d_{\text{model}} = d_{\text{model}}^2$ parameters.

4. Grand total: $3 \times d_{\text{model}}^2 + d_{\text{model}}^2 = 4 \times d_{\text{model}}^2$ parameters.

This is exactly what we'd have with single-head attention using the full dimension, plus one additional output projection. The key benefit of multi-head attention is that we get $h$ different attention patterns, each potentially specializing in different relationships, for essentially the same parameter budget as a single larger head.