Multi-Head Attention

Why Does This Exist?

A single self-attention computation produces exactly one weighted average per token, using one set of $Q$, $K$, $V$ projections. That's a real limitation: language relationships aren't one-dimensional. The same sentence needs one kind of attention to track which pronoun refers to which noun, a different kind to track subject-verb agreement, and another to track which adjective modifies which noun. One softmax distribution per token can express only one weighting scheme at a time — it can't simultaneously be sharply focused on a distant antecedent and broadly averaged over nearby modifiers, because it's one distribution, not several.

The fix researchers landed on wasn't to make one attention computation smarter. It was to run several smaller, independent attention computations side by side, each free to learn a different pattern, and combine their outputs afterward.

Think of It Like This

How It Actually Works

Split, attend, concatenate, project

Take the diagram above literally. The model dimension $d$ splits into $h$ heads, each of dimension $d_k = d / h$. Each head gets its own $W_Q^{(i)}$, $W_K^{(i)}$, $W_V^{(i)}$ — independently learned projections into that head's smaller subspace — and runs the full self-attention computation entirely on its own:

Every head produces its own $n \times d_k$ output, computed independently and in parallel — head 1's attention weights have no influence on head 2's. Once all $h$ heads finish, their outputs concatenate back into a single $n \times d$ matrix, restoring the original width:

The final matrix $W_O$ is a learned linear mixing layer applied after concatenation — it's not optional. Without it, each output dimension would come from exactly one head with no way for heads to interact in the final representation; $W_O$ lets the model combine information across heads before passing it on.

Same total compute, different allocation

Here's the detail worth sitting with: splitting $d$ into $h$ heads of size $d/h$ doesn't add extra parameters or extra compute compared to one attention computation at the full dimension $d$ — the total work is roughly the same, just divided among $h$ smaller, parallel computations instead of one large one. What multi-head buys isn't more capacity in the aggregate; it's more independent degrees of freedom in how that capacity gets used, since each head learns its own projections and can specialize.

What heads actually learn

Empirically, when researchers inspect trained multi-head attention, individual heads often turn out to specialize — one head consistently attending to the previous token, another consistently attending to a syntactic dependency, another spreading attention broadly across the whole sequence. None of that specialization is designed in; it falls out of training the same way self-attention's linguistic structure does. Not every head specializes cleanly, and some heads turn out to be redundant enough that they can be pruned after training with little quality loss — a genuine, measured finding, not a hypothetical.

Show Me the Code

Splitting a model dimension into heads, running self-attention in each, and concatenating back — the entire mechanism in one function.

1import numpy as np2
3
4def softmax(z: np.ndarray) -> np.ndarray:5    e = np.exp(z - z.max(axis=-1, keepdims=True))6    return e / e.sum(axis=-1, keepdims=True)7
8
9def multi_head_attention(x: np.ndarray, w_q: np.ndarray, w_k: np.ndarray, w_v: np.ndarray, w_o: np.ndarray, n_heads: int) -> np.ndarray:10    n, d = x.shape11    d_k = d // n_heads12    q, k, v = x @ w_q, x @ w_k, x @ w_v                     # (n, d) each13    heads = []14    for i in range(n_heads):15        sl = slice(i * d_k, (i + 1) * d_k)                  # this head's slice of the dim16        scores = q[:, sl] @ k[:, sl].T / np.sqrt(d_k)17        heads.append(softmax(scores) @ v[:, sl])            # (n, d_k), independent of other heads18    return np.concatenate(heads, axis=-1) @ w_o              # concat back to (n, d), then mix19
20
21rng = np.random.default_rng(0)22x = rng.normal(size=(4, 8))23w_q, w_k, w_v, w_o = (rng.normal(scale=0.3, size=(8, 8)) for _ in range(4))24out = multi_head_attention(x, w_q, w_k, w_v, w_o, n_heads=2)25print(out.shape)  # -> (4, 8) — same shape as the input, regardless of head count

Notice each head only ever sees its own slice of q, k, v — head 0 never reads head 1's portion of the dimension, which is exactly the "independent, in parallel" claim above made concrete.

Watch Out For

The Quick Version

• One self-attention computation produces one weighting scheme per token; multi-head attention runs several in parallel, each in its own subspace.
• Each head gets its own $Q$, $K$, $V$ projections into a dimension $d/h$, computed independently of the other heads.
• Head outputs concatenate back to the original width, then pass through a learned output projection $W_O$ that mixes across heads.
• Total compute is roughly the same as one full-dimension attention computation — multi-head redistributes capacity, it doesn't add much of it.
• Individual heads often specialize during training, though not all of them, and some turn out to be prunable.

What to Read Next

• Self-Attention is the single computation multi-head attention runs several copies of, in parallel.
• Cross-Attention can also be split across multiple heads, using the same mechanism this page describes.
• Transformer Architecture is where multi-head attention sits as one of the two sub-layers in every block.
• Attention Complexity covers the cost of running attention, whether single-head or multi-head.