Concatenate Heads and Project Once

After attention, head results are returned to position-major order, concatenated along the feature axis, and mixed by one output projection. Calculate this reconstruction and show why head order and W_O row-block order must agree.

After each head has produced

Z^{(r)}\in\mathbb R^{B\times T\times d_v},

the results are concatenated along their feature axes:

Z=\operatorname{Concat}(Z^{(1)},\ldots,Z^{(h)}) \in\mathbb R^{B\times T\times hd_v}.

The output projection then returns to model width:

Y=ZW_O,\qquad W_O\in\mathbb R^{hd_v\times d_{model}}.

Concatenation preserves each head result as a separate block of features. It is not an average over heads. $W_O$ decides how those blocks are mixed into the model-width output.

Continue the Two-Head Trace

The final query produced

z_3^{(1)}\approx[0.752,0.752],\qquad z_3^{(2)}\approx[-0.546,0.926].

Head-order concatenation gives

z_3\approx[0.752,0.752,-0.546,0.926].

With identity $W_O$ , the layer returns that vector unchanged. To make the mixing visible, use

W_O= \begin{bmatrix} 1&0&0.5&0\\ 0&1&0&0.5\\ 0.5&0&1&0\\ 0&0.5&0&1 \end{bmatrix}.

Then

y_3=z_3W_O\approx[0.479,1.215,-0.170,1.302].

The first output coordinate combines head 1's first feature with half of head 2's first feature. The fourth combines half of head 1's second feature with head 2's second feature. Attention weights remain head-specific, while the output projection can mix the features they produced.

Head Order Is a Coordinate Convention

Suppose we swap the order of head blocks in $Z$ . If we also swap the matching row blocks of $W_O$ , $Y$ remains unchanged. In matrix form, if $P$ permutes the concatenated head features,

Z'=ZP,\qquad W_O'=P^\top W_O,

then

Z'W_O'=ZPP^\top W_O=ZW_O.

This symmetry means that “head 1” and “head 2” have no fixed meaning across independently trained models. A head index becomes meaningful only inside a specific parameterization and checkpoint.

Zeroing a Head Occurs Before the Projection

To ablate head 2, replace its result with zero before concatenation:

z_{3,\text{ablated}}=[0.752,0.752,0,0].

The same $W_O$ then produces a different $y_3$ . The difference measures the direct effect of removing that head's current contribution through the fixed output projection. In a complete network, later layers and residual paths may change the final effect, and retraining may let other components compensate.

Q1. Join two head readings

Two width-2 heads return $[1,-2]$ and $[3,4]$ . They are concatenated in that order, and $W_O=I_4$ . What is the third coordinate of the layer output?

Your answer

Compute it first, then check your number.

Hint

Write the complete concatenated vector before applying the identity projection.

Solution

The concatenation is

[1,-2,3,4]

, so the third coordinate is 3.

Treat Concatenation Order and Projection Rows as One Contract

An implementation can produce correct head outputs and still fail by joining them in one order while using $W_O$ trained for another. Audit the ordered head blocks immediately before the output projection, not only the final width.