Beyond shape preservation, name two things W_O does after head concat
Multi-head attention concatenates h head outputs of dim d_head back to d_model, then multiplies by an output projection W_O of shape (d_model, d_model). Many people see W_O as a 'shape preservation' bookkeeping step, but it has substantive roles. Name two things W_O contributes beyond just preserving d_model, and explain what is lost when you remove it even though shapes still line up.
W_O does cross-head mixing (lets heads share residual directions) and learned per-head weighting (amplifies useful heads, suppresses noisy ones). Removing it preserves shape but kills both.
Imagine eight researchers writing reports into separate columns of a shared spreadsheet. Without W_O, every researcher's findings stay in their own column. Later readers can still look at every column, but no two researchers can add to the same line. W_O is the editor who reads all eight columns and rewrites them into a single mixed document. In it, insights from researcher 3 can reinforce insights from researcher 7. The editor can also suppress the column from researcher 5 if their work is mostly noise. The shape of the document is the same with or without the editor, but the content is not. Take the editor away and you lose both the cross-pollination and the quality control.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
The output projection W_O is the most underappreciated component of multi-head attention. It is the only place where the outputs of different heads can mix with each other before the residual add. It is also a natural place for the model to learn that some heads carry more signal than others. Treating it as a 'shape-preservation' formality is a common reading that misses what makes multi-head attention work.
Mental model: the concat is bookkeeping. W_O is where the heads become a multi-head system instead of h independent attention computations stapled together.
The per-head slice structure after concat
Multi-head attention with h heads and d_head = d_model / h per-head dimension produces, for each token position, an output tensor of shape (h, d_head). Concatenating along the last axis flattens this to (d_model,). The first d_head entries come from head 0; the next d_head from head 1; and so on.
The slice structure is a feature, not a bug
The slice structure is what makes parallel multi-head attention work in the first place. Each head computes its own softmax(Q_i K_i^T / sqrt(d_k)) V_i independently, with no inter-head communication. That is what allows batching all h heads into a single big matmul on the GPU.
But after the concat, that same slice structure means head i's output sits in a fixed d_head-wide region of the d_model vector. No operation has mixed it with the other heads.
What this means for the residual add
Without any mixing, the post-concat vector would be added to the residual stream as-is. Head 0's output would land in residual stream dimensions 0 to d_head-1; head 1 in d_head to 2*d_head-1; and so on. The attention writes would partition the residual stream into h fixed slices, one per head. Downstream layers could still read every slice through their dense W_Q, W_K, W_V and the FFN, but no two heads could write to the same direction.
Mixing at the attention sub-layer via W_O lets each head choose where it writes. Without it, each head's write slice is fixed by its index. Every production transformer keeps W_O.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model: int, n_heads: int):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.d_head = d_model // n_heads
self.W_q = nn.Linear(d_model, d_model, bias=False)
self.W_k = nn.Linear(d_model, d_model, bias=False)
self.W_v = nn.Linear(d_model, d_model, bias=False)
# W_O does the cross-head mixing and per-head weighting.
self.W_o = nn.Linear(d_model, d_model, bias=False)
def head_importance(self) -> torch.Tensor:
# Per-head write strength: norm of the d_head-wide columns of W_O
# that read from each head's slice. Low norm = effectively pruned.
W = self.W_o.weight # (d_model, d_model)
return torch.tensor([
W[:, h * self.d_head : (h + 1) * self.d_head].norm()
for h in range(self.n_heads)
])Real products, models, and research that use this idea.
- Anthropic's induction-heads paper (2022) describes composition between heads in different layers through each head's OV circuit (W_V with W_O) writing to residual subspaces that later heads read.
- Voita et al. 2019, 'Analyzing Multi-Head Self-Attention', identified prunable heads using layer-wise relevance, head confidence, and L0-gated pruning.
What an interviewer would ask next. Try answering before peeking at the approach.
QCould you replace W_O with a block-diagonal matrix that only mixes within each head's d_head slice?
You could, but each head could then write only into its own slice. Heads would still never share a write direction, while you still pay d_model^2 / h parameters for the diagonal blocks. The off-diagonal blocks are what let a head write outside its own slice.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Treating W_O as a 'shape-preservation' formality and missing that it is the only point in the attention sub-layer where head outputs can mix with each other.
60 second bullets to scan on the way to the call.
Two distinct roles of W_O beyond shape preservation
Why cross-head mixing requires a dense d_model x d_model matrix, not block diagonal
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.