Skip to content
AI360XpertBeta

Multi-Head Attention

Split the model width into several smaller subspaces and run an independent attention in each one, so different heads can track different kinds of relationship.

Width 8 splits into 2 heads of width 4, each runs its own attention with weights 0.665, 0.245 and 0.09, then concatenation restores width 8 before one mixing projection
Width 8 splits into 2 heads of width 4, each runs its own attention with weights 0.665, 0.245 and 0.09, then concatenation restores width 8 before one mixing projection

Why Does This Exist?

One self-attention pass gives each token exactly one weighted average of the sequence, from one set of query, key, and value projections. Language needs more than one kind of looking at once. In "the vase that the movers dropped shattered," one pattern must link "shattered" back to "vase" across the whole clause, while another must link "dropped" to its neighbors "movers" and "vase." A single softmax distribution per token cannot be sharply focused far away and broadly spread nearby at the same time. It is one distribution, not several.

The fix is embarrassingly direct: run several smaller attentions side by side, each with its own projections, and combine their outputs after. Nothing about any single attention gets smarter. There are just more of them, and they are free to disagree.

This page covers the split, the per-head computation, and the final mixing. The single-head mechanism underneath is on the self-attention page, and the transformers page shows where the whole block sits.

The one idea you need first. Attention is a weighted average where the weights come from the inputs themselves: each token scores every token, normalizes the scores with a softmax, and blends. If that sentence is hazy, read the self-attention page first, because every head below runs exactly that computation.

Think of It Like This

Three proofreaders with different checklists

One proofreader hunting spelling mistakes, verb agreement errors, and wrong dates in a single pass will miss things. The three concerns need different focus: spelling wants letter-by-letter scanning, agreement wants whole-sentence structure, dates want cross-checking against a calendar. Hand the same essay to three proofreaders, each with one checklist, and each one reads narrowly and sharply.

Multi-head attention is that panel. Each head reads the same sequence but through its own projections, so each is free to track a different pattern, and their notes get merged afterward into one marked-up page.

Where the analogy stops: the proofreaders are handed their checklists, while attention heads are handed nothing. Their specialties emerge from training, and some heads never specialize at all. A few turn out redundant enough to prune after training with little quality loss.

How It Actually Works

Split, attend, combine

Call the model width dd and the head count hh, where hh divides dd. Each head works in a slice of width dk=d/hd_k = d / h. Head ii owns three learned projection matrices, WQ(i)W_Q^{(i)}, WK(i)W_K^{(i)}, and WV(i)W_V^{(i)}, which map the input XX into that head's private queries, keys, and values. It then runs ordinary attention, fully independently of every other head:

headi=Attention(XWQ(i),XWK(i),XWV(i))\text{head}_i = \text{Attention}(X W_Q^{(i)}, X W_K^{(i)}, X W_V^{(i)})

Each head returns an n×dkn \times d_k output for a sequence of nn tokens. The hh outputs concatenate back into one n×dn \times d matrix, restoring the original width, and one learned output matrix WOW_O mixes across heads:

MultiHead(X)=Concat(head1,…,headh)WO\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W_O

WOW_O is load-bearing, not decoration. Without it each output dimension would belong to exactly one head forever, and heads could never combine evidence in the final representation.

Worked example

Take d=8d = 8 with h=2h = 2 heads, so dk=4d_k = 4, over n=4n = 4 tokens. Zoom into head 1 for a single token whose raw scores against three tokens are [2,1,0][2, 1, 0]. Exponentiate after subtracting the max: [e0,e−1,e−2]=[1.0,0.368,0.135][e^0, e^{-1}, e^{-2}] = [1.0, 0.368, 0.135], sum to 1.5031.503, divide to get [0.665,0.245,0.090][0.665, 0.245, 0.090]. The top token takes two-thirds of the weight without taking all of it, which is why attention blends rather than selects. Head 2 computes its own weights from its own slice at the same time, seeing none of head 1's numbers. Both 4×44 \times 4 outputs concatenate into 4×84 \times 8, exactly the input shape, and WOW_O mixes the two halves.

Splitting dd into hh slices adds roughly no parameters and no compute over one full-width attention. The budget is the same; it is just divided into hh independent degrees of freedom instead of one shared one.

Code

The point of this snippet is the independence claim made concrete: each head reads only its own slice, and the shape comes back unchanged.

import numpy as np

def softmax(z: np.ndarray) -> np.ndarray:    e = np.exp(z - z.max(axis=-1, keepdims=True))    return e / e.sum(axis=-1, keepdims=True)

def 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:    n, d = x.shape    d_k = d // n_heads    q, k, v = x @ w_q, x @ w_k, x @ w_v  # (n, d) each    heads = []    for i in range(n_heads):        sl = slice(i * d_k, (i + 1) * d_k)  # this head's slice only        scores = q[:, sl] @ k[:, sl].T / np.sqrt(d_k)        heads.append(softmax(scores) @ v[:, sl])  # (n, d_k), apart from other heads    return np.concatenate(heads, axis=-1) @ w_o  # back to (n, d), then mix

print([round(float(v), 3) for v in softmax(np.array([2.0, 1.0, 0.0]))])# -> [0.665, 0.245, 0.09]
rng = np.random.default_rng(0)x = rng.normal(size=(4, 8))w_q, w_k, w_v, w_o = (rng.normal(scale=0.3, size=(8, 8)) for _ in range(4))print(multi_head_attention(x, w_q, w_k, w_v, w_o, n_heads=2).shape)# -> (4, 8)

The weights match the hand computation above, and the output keeps the input shape regardless of head count.

Watch Out For

Raising head count while the model width stays fixed

The symptom is quality that improves from 4 heads to 8, then quietly degrades at 16, which gets blamed on the data or the learning rate. But doubling head count at fixed dd halves each head's width dkd_k, and a head with too little room cannot represent much. Head count and head width are one joint decision, not two independent knobs.

Before adding heads, check dkd_k. If it drops into the single digits while quality falls, the heads are starved, and the fix is a wider model or fewer heads, not more training.

Dropping the output projection in a hand-rolled implementation

The symptom is shapes that line up perfectly with no errors, plus a model that trains yet plateaus below every reference. Concatenation alone restores the width, so nothing complains, but each output dimension stays tied to whichever single head produced it. The model lost its only cross-head mixing, silently.

Always end with the WOW_O multiply. If a from-scratch implementation underperforms a library equivalent with identical settings, the missing projection is the first place to look.

The Quick Version

  • One attention pass gives one weighting scheme per token, so multi-head runs hh smaller attentions in parallel, each in its own slice of width d/hd / h.
  • Each head owns private QQ, KK, VV projections and computes independently of the other heads.
  • Head outputs concatenate back to the original width, then pass through a learned WOW_O that mixes information across heads.
  • Total compute roughly matches one full-width attention: multi-head redistributes capacity instead of adding it.
  • Heads often specialize during training, some stay redundant, and starving dkd_k by adding heads at fixed width hurts quality.