Multi-Head Attention — Multiple Parallel Attention Patterns

The split: partition d_model into h heads

The core trick is that heads share the width instead of each taking the full width. Pick a head count h that divides d_model, and give every head a per-head dimension

d_h = d_model / h

A common configuration is d_model = 512, h = 8, so d_h = 64; larger models keep the ratio, e.g. d_model = 4096 with h = 32 and d_h = 128. Each head projects the input into its own d_h-dimensional query, key, and value spaces, runs ordinary scaled dot-product attention there, and produces a d_h-dimensional output. Because the per-head width shrinks in exact proportion to the head count, the total projection work is unchanged — the split reorganizes the budget rather than enlarging it.

Per-head projections

Each head i owns three projection matrices that map from the full model width into its own narrow subspace:

X       : [N, d_model]
Q_i = X W_Q^i    W_Q^i : [d_model, d_h]   → Q_i : [N, d_h]
K_i = X W_K^i    W_K^i : [d_model, d_h]   → K_i : [N, d_h]
V_i = X W_V^i    W_V^i : [d_model, d_h]   → V_i : [N, d_h]

The heads do not see disjoint input coordinates — each projection reads all of X; what differs is the learned subspace each head projects into. In practice the h small matrices are stored as three big ones, W_Q, W_K, W_V each [d_model, d_model], and the result reshaped to [N, h, d_h]. That reshape is the entire implementation of the split: one large GEMM, then a view that carves the width into h contiguous bands.

Concatenate, then mix with the output projection

Running attention independently in each head gives h outputs of width d_h. Concatenating them along the feature axis rebuilds a d_model-wide vector, but that raw concatenation is a stack of isolated slices — head 3 knows nothing of head 5. The output projection W_O lets the heads talk:

head_i = Attention(Q_i, K_i, V_i)   : [N, d_h]
concat = [head_1 ; head_2 ; … ; head_h]  : [N, d_model]
MHA    = concat · W_O          W_O : [d_model, d_model]

W_O is not a formality. Without it the block would emit h independent sub-vectors wedged into fixed coordinate ranges; W_O mixes them into a shared representation, letting the model weight and combine what each head retrieved. Concatenation plus W_O is mathematically a sum of each head’s output passed through its own slice of W_O — the heads added back together, learnably.

Why several narrow heads beat one wide head

Make the comparison fair: one head of width d_model versus h heads of width d_model/h, same parameter count. The wide head computes a single score matrix — one mixing distribution per query. The h-head version computes h separate distributions, each an [N, N] map in its own subspace, and can attend to h different positions at once. It is strictly more expressive at the level of attention patterns.

The trade is resolution per head: at d_h = 64 each comparison lives in a smaller space than a full 512-wide dot product. In practice that is a good bargain — the relationships a head encodes (agreement, adjacency, coreference) are low-rank, so 64 dimensions are ample, and eight independent low-rank views beat one high-rank view that still collapses to a single blend.

Subspace specialization

Because each head projects into its own subspace and competes only within its own softmax, heads are free to specialize, and trained models show that they do. Probing studies find heads that reliably track particular relations: one attends to the previous token, another to the syntactic head of the current word, another to the matching bracket or quotation mark, another to rare or repeated tokens.

This is not designed in; it emerges because W_O rewards heads for supplying distinct information — two heads that learned the same pattern are redundant, and gradient descent has little pressure to keep both. The upshot is a division of labor: the sequence is read simultaneously through several complementary lenses, and the output projection decides how much each contributes.

A worked split

Take d_model = 4 and h = 2, so d_h = 2. A single token’s projected query is a 4-vector that the reshape carves into two heads:

q = [ 0.9, 0.1 | 0.2, 0.8 ]      split →  q_1 = [0.9, 0.1],  q_2 = [0.2, 0.8]

Suppose head 1’s keys favor position A (aligned with [1,0]) and head 2’s favor position B (aligned with [0,1]). Then q_1 attends mostly to A and q_2 mostly to B: one token, one forward pass, two positions read at once — something a single head cannot do without averaging A and B into one muddy distribution. The concatenation [H_1 ; H_2] is then a 4-vector again, which W_O blends.

Parameter accounting

The headline fact: multi-head attention with h heads costs the same parameters as one full-width head plus an output projection. Summed over heads, the query projections are h matrices of [d_model, d_h] that stack into one [d_model, d_model] — identical to a single wide head. The same holds for keys and values, and W_O adds one more [d_model, d_model]:

params(MHA) = 4 · d_model^2   (Q, K, V, O; ignoring biases)

So h is essentially free in parameter terms — it repartitions a fixed 4 d_model^2 budget into more, narrower heads. That is why changing the head count at fixed d_model barely moves the model size, and why head count is an architecture dial separate from width and depth.

FLOP and memory accounting

Compute is likewise conserved. The projection GEMMs cost O(N d_model^2) whether the width is sliced or not, and the scores and value-mixing cost Σ_i O(N^2 d_h) = O(N^2 d_model) in total — d_h times h recovers d_model. So multi-head attention has the same asymptotic cost as single-head attention of the same width; heads reshape the arithmetic, they do not add to it.

Memory tells a different story at inference. The KV cache stores keys and values for every head at every past position: 2 · N · h · d_h = 2 N d_model numbers per layer. On a CPU-bound small model that streaming cost, not the FLOPs, dominates single-token decoding — the pressure that motivates the head-sharing schemes below.

Head redundancy and pruning

The freedom that lets heads specialize also lets them go to waste. Empirically, many heads in a trained transformer can be removed at test time with little loss: some encode redundant patterns, and a handful of “important” heads do the bulk of the work in any given layer. Pruning studies routinely drop a large fraction of heads while keeping accuracy nearly intact.

This does not make multi-head attention wasteful during training — the redundancy appears to help optimization, giving the model several chances to find a useful pattern. But the trained head count often exceeds what the deployed model needs. For CPU inference that is an opportunity: prune low-importance heads to shrink the KV cache and projection work, or, better, share key/value heads from the start.

Grouped-query and multi-query attention

The most consequential modern variation attacks the KV-cache cost directly. In vanilla multi-head attention every head has its own keys and values. Multi-query attention keeps h separate query heads but shares a single key/value head across all of them, collapsing the KV cache by a factor of h. Grouped-query attention (GQA) is the middle ground: partition the h query heads into g groups, each group sharing one key/value head, so the cache shrinks by h/g.

GQA is now the default in large open models: it recovers most of multi-query’s memory savings while preserving nearly all of multi-head’s quality, because the query side keeps its full complement of specialized heads and only the expensive-to-cache key and value sides are consolidated. For a small CPU model, where every decoded token streams the whole cache from RAM, that is often the single highest-leverage change available.

Pitfalls worth avoiding

Several misreadings recur. “More heads means more parameters.” No — at fixed d_model, heads partition the width, so the parameter count is unchanged. “Each head sees a different part of the input.” No — every head reads all of X; they differ in the subspace they project into, not the coordinates they observe. “The concatenation is the output.” The concatenation is just a buffer; without W_O the heads never interact and the block is far weaker. “Scale by √d_model.” The divisor is √d_h, the per-head width. Keep these straight and the whole mechanism reads as one idea: slice the width, attend independently in each slice, then mix the slices back together.