Review

Review the complete multi-head operation, axis trace, separate and combined projection forms, output reconstruction, parameter conventions, ablations, research evidence, and nested audit order.

The Complete Multi-Head Operation

For head $r\in\{1,\ldots,h\}$ ,

\begin{aligned} Q^{(r)}&=XW_Q^{(r)},\\ K^{(r)}&=XW_K^{(r)},\\ V^{(r)}&=XW_V^{(r)},\\ A^{(r)}&=\operatorname{softmax}_{last}\!\left( \frac{Q^{(r)}K^{(r)\top}}{\sqrt{d_k}}+M \right),\\ Z^{(r)}&=A^{(r)}V^{(r)}. \end{aligned}

The head results are joined and projected once:

Z=\operatorname{Concat}(Z^{(1)},\ldots,Z^{(h)}), \qquad Y=ZW_O.

With batch-major input, the shapes are

X:B\times T\times d_{model},

Q,K:B\times h\times T\times d_k,\qquad V:B\times h\times T\times d_v,

A:B\times h\times T\times T,\qquad Z:B\times T\times(hd_v),\qquad Y:B\times T\times d_{model}.

Why Multiple Heads Differ from One Wide Head

One head forms one attention distribution per query. Every value coordinate in that head uses the same weights. Multiple heads form several independently normalized distributions and may route different value subspaces from different positions.

This is a capability, not a specialization guarantee. Heads may learn distinct patterns, redundant patterns, input-dependent mixtures, or negligible contributions.

Separate and Combined Projections

Separate matrices make the head equations explicit. Combined matrices place all head projections side by side:

W_Q:d_{model}\times(hd_k),\quad W_K:d_{model}\times(hd_k),\quad W_V:d_{model}\times(hd_v).

The implementation then traces

(B,T,hd_k)\to(B,T,h,d_k)\to(B,h,T,d_k)

before attention and reverses the head-position transpose after attention. Reshape splits or joins axes; transpose reorders named axes. Their jobs are not interchangeable.

Parameter and Work Conventions

Without biases, the projection count is

2d_{model}hd_k+2d_{model}hd_v.

Under $hd_k=hd_v=d_{model}=d$ , it becomes $4d^2$ , independent of how the fixed total width is divided among heads. The dense score tensor still contains $BhT^2$ entries, one $T\times T$ map for every batch item and head.

Parameter entries, stored score entries, arithmetic, memory traffic, and measured runtime are different quantities. A careful comparison names which one it reports.

Head Identity and Ablation

Head indices are coordinate labels. Permuting head blocks and applying the matching inverse permutation to $W_O$ preserves the layer function. Similar roles may appear at different indices across training runs.

Zeroing a head for one input, pruning it at evaluation, fine-tuning after pruning, and training fewer heads from the start are different experiments. Published pruning and specialization results are evidence about stated models, tasks, metrics, and procedures—not universal laws of head count.

Audit Order

State $B,T,d_{model},h,d_k,d_v$ and the bias convention.
Verify combined projection widths.
Label axes before split and transpose.
Audit scores, masks, softmax rows, and values inside each head.
Compare separate and combined head results.
Reverse the head-position transpose.
Concatenate along the head-feature axis in the declared order.
Verify $W_O:(hd_v,d_{model})$ and selected numerical rows.
Test head zeroing and head permutation at explicitly named locations.

If a Step Is Unclear

Return to Self-Attention in Matrix Form for the computation inside one head.
Return to Add the Batch Axis Without Mixing Sequences for batch-major products and mask broadcasting.
Revisit Combined Projections Introduce a Head Axis when the equations are clear but the transforms are not.

Check Yourself

Can you explain why $d_v=128$ in one head still gives one weight row?
Can you derive $(B,h,T,T)$ without memorizing a diagram?
Can you convert separate Q/K/V matrices into combined matrices?
Can you state when $d_{model}$ must be divisible by $h$ and when that is only an API convention?
Can you count parameters without silently adding or removing biases?
Can you define the exact location of a head ablation?
Can you keep a specialization claim within its model and evidence?

Several heads can now produce several content-dependent reading patterns. The operation still lacks an independent representation of absolute or relative position. That missing information is the problem Chapter 3 addresses.