Describe the W_O projection in multi-head attention, its shape and what it mixes.

**Multi-head attention** runs several attention computations in parallel, one per head, each producing a `d_head`-wide context vector. With `h` heads and `d_head = d_model / h`, the per-head outputs concatenate into a single `d_model`-wide vector per token. After the concat, one more linear layer is applied: the output projection `W_O`. Its shape is `(d_model, d_model)`, so it takes the concatenated d_model vector in and produces a d_model vector out. The shape on either side is the same, which is why people sometimes dismiss it as a formality. The real role of W_O is **mixing across heads.** After the concat, head 0's output sits in the first d_head dimensions, head 1's in the next d_head, and so on. Without W_O, each head's signal stays locked in its own slice forever, and the heads cannot combine. W_O is a dense matrix, so every output dimension is a learned combination of every input dimension, which lets head outputs land on overlapping parts of the residual stream. This mixing is what allows the model to combine, for example, head 3's syntactic-attachment signal with head 7's coreference signal, into one mixed representation that the next layer's FFN and attention can both use.

**W_O** is the output projection at the end of each multi-head attention sub-layer. After the `h` heads each compute their `d_head`-wide context vectors and the results are concatenated into a `d_model`-wide vector, W_O applies one more linear map of shape `(d_model, d_model)` before the residual add. The shape is square, so W_O does not change the dimensionality. What it does change is the structure of the output. Without W_O, head `i`'s output occupies columns `[i*d_head, (i+1)*d_head)` of the concatenated vector and stays there. Each head writes into a fixed disjoint slice of the residual stream, and the heads can never combine. W_O is a dense matrix, which means every output dimension is a learned weighted combination of every input dimension. This breaks the slice structure: head outputs can now land on overlapping subspaces of the residual stream, and the model gets to learn which heads combine with which. A convenient mental decomposition: think of W_O as `[W_O^(0), W_O^(1), ..., W_O^(h-1)]` stacked along the input axis, where each `W_O^(i)` is the `(d_head, d_model)` block that maps head `i`'s output into the full residual stream. Each per-head write block is learned independently and can write anywhere in d_model. In other words, W_O is the only place in the attention sub-layer where the heads stop being independent and become a multi-head system.