← All Posts
Deep Learning · Transformers· Foundations

Multi-Head Attention: Separate Reads, One Residual Update

A head is a separate learned weighting of the available values. Multiple heads can retrieve different information for the same query position. The output projection then combines their results into one update in the model's shared feature space.

Inside a block → attention → combine the heads. View the model diagram →

Let the same token make more than one read

One head produces one mixture of values. If a token needs information from different places, several heads let it make several reads at once. Every head can see the same allowed positions, but its own query, key, and value projections let it produce a different mixture.

The single-head chapter explained one weighted read. Here we will run two of them, place their outputs next to each other, and project the joined vector into one update for the block.

Two heads read the same two positions

Give each head a one-number value at each of two source positions. The weights and values below are chosen for illustration. Both heads use the same weighted-sum operation; their separate learned projections can produce different weights and values.

Separate reads become one update
Head 1

Weights: (0.8, 0.2)

Values: (10, 0)

0.8 × 10 + 0.2 × 0

Read = 8
Head 2

Weights: (0.1, 0.9)

Values: (0, 20)

0.1 × 0 + 0.9 × 20

Read = 18
Join the outputs: (8, 18)

Then apply a learned output projection.
In this example: (8 + 18, 8 − 18) = (26, −10).

Both reads contribute to the same token’s update. Joining keeps their coordinates; the output projection can then combine or contrast them.

The projection here is $W_O=\begin{pmatrix}1&1\\1&-1\end{pmatrix}$, so $(8,18)W_O=(26,-10)$ using row vectors. It is an ordinary learned matrix, not another softmax that selects a winning head.

Repeat the single-head formula

For head $h$, form $Q_h=XW_Q^{(h)}$, $K_h=XW_K^{(h)}$, and $V_h=XW_V^{(h)}$. Each head has its own learned projections. Its output is:

$$O_h=\operatorname{softmax}_{\mathrm{row}}\!\left(\frac{Q_hK_h^\top}{\sqrt{d_h}}+M\right)V_h.$$

Join all $H$ outputs, then apply the output projection:

$$Y=\operatorname{Concat}(O_1,\ldots,O_H)W_O.$$

In the usual equal-width setup, $d_h=d/H$, each head returns width $d_h$, and concatenation restores width $d$. The projection $W_O\in\mathbb R^{d\times d}$ maps these head coordinates back into the residual stream. It does not compute a second softmax over heads.

The split is an arrangement of projected coordinates

Use $B$ for batch size, $n$ for token count, $d$ for model width, and $H$ for head count. For $B=2$, $n=3$, $d=8$, $H=2$, one combined query projection gives $[2,3,8]$. Reshape to $[2,3,2,4]$, then transpose token and head axes to obtain $[2,2,3,4]$. K and V follow the same layout. Batched score multiplication produces $[2,2,3,3]$.

The output is $[2,2,3,4]$. Transpose back to $[2,3,2,4]$, then reshape to $[2,3,8]$. Flattening the first layout directly mixes token and head coordinates. This is a common bug because the final shape can still look plausible.

# x: [B, n, d], and d is divisible by H
q = (x @ Wq).reshape(B, n, H, dh).transpose(0, 2, 1, 3)
k = (x @ Wk).reshape(B, n, H, dh).transpose(0, 2, 1, 3)
v = (x @ Wv).reshape(B, n, H, dh).transpose(0, 2, 1, 3)
# Per-head scaled, masked softmax attention returns [B, H, n, dh].
o = attend(q, k, v, allowed)
y = o.transpose(0, 2, 1, 3).reshape(B, n, d) @ Wo

Head count is not the projection parameter count

Each combined Q/K/V matrix is $d\times d$ in the standard setup, regardless of how those $d$ output coordinates are partitioned into heads. Adding the output projection gives $4d^2$ parameters before biases. Increasing $H$ at fixed $d$ narrows each head rather than multiplying total projected width.

The attention arithmetic sums $H$ products with head dimension $d/H$, yielding an $O(n^2d)$ term. A naïve implementation stores $Hn^2$ probabilities, however, so changing head count can affect memory and kernel behavior even when the leading arithmetic and projection parameter counts remain unchanged.

Separate head count from key/value sharing

Standard MHA gives each query head its own key/value projections. MQA shares one K/V head across all queries; GQA shares within groups. These change the number of K/V vectors that must be cached during generation. The number of distinct query weighting patterns can remain $H$ in all three cases.

This distinction is why “fewer K/V heads” does not mean “only one attention distribution.” Read MHA, MQA, and GQA for the cache and parameter calculations. MLA introduces a different shared latent representation and is explained separately.

Read an attention visualization with care

Heads are not assigned fixed jobs such as “syntax” and “meaning”. Training determines their behavior, and different heads can be redundant or hard to interpret.

A heatmap shows coefficients in a value mixture. It does not directly show the effect of a source token on the final prediction. Value vectors, output projections, residual routes, MLPs, and later layers all intervene. A large attention coefficient attached to a nearly zero value may contribute little; a small coefficient attached to a large relevant value may matter substantially.

Try it: Could two heads have identical attention weights but contribute different outputs?

Yes. Their value projections can differ. The same weights can mix different value vectors, and the output projection can use those results differently.

Decide which positions each head may read

The heads now return one combined update per position. Next, causal masking explains how every head avoids reading future tokens. The multi-head construction originates in Attention Is All You Need.