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 / hA 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.