Beyond shape preservation, name two things W_O does after head concat

Short answer·Medium·4.0 · 0·~3 min·Asked atCanva Databricks Notion·Relevant atMeta

Attempt it

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 why removing it tanks quality even when shapes still line up.

Free · 2 AI evals / day

**Multi-head attention** runs `h` independent attention computations in parallel, each producing a `d_head`-wide context vector per token. These get concatenated back into a single `d_model`-wide vector (`d_model = h * d_head`), and then a final linear projection `W_O` of shape `(d_model, d_model)` is applied before the residual add. If you only think of `W_O` as 'keeping shapes consistent', you miss two important things it actually does. First, **cross-head mixing.** After the concat, head 0's output occupies dimensions 0 to d_head-1, head 1 occupies d_head to 2*d_head-1, and so on. Without `W_O`, each head's output writes into its own slice of the residual stream and stays there. Heads cannot ever combine. `W_O` is a dense matrix that lets every output dimension be a learned combination of every input dimension, so the heads can finally talk. Second, **learned per-head weighting.** During training, `W_O` learns which heads carry useful signal and which are noise. It can amplify the useful ones and suppress the noisy ones. You can literally read off head importance by inspecting the norm of the relevant columns of `W_O` after training. Remove `W_O` and the shapes still line up, but the model loses both the mixing and the weighting, and perplexity collapses.

**Cross-head mixing** is the first answer. After the concat, head `i`'s `d_head`-wide output sits at columns `[i*d_head, (i+1)*d_head)` of the concatenated vector. Without `W_O`, that block writes straight into the same slice of the residual stream and stays there for the rest of the forward pass. Heads live in disjoint sub-spaces of the residual stream forever, and the FFN downstream has to compose them without any inter-head structure. `W_O` is a dense `(d_model, d_model)` matrix, so every output dimension is a learned combination of every input dimension. Heads can land on overlapping subspaces of the residual stream and contribute to the same directions. This is the substrate that makes phenomena like induction heads composing with previous-token heads possible at all. **Learned per-head weighting** is the second answer. `W_O` is free to up-weight useful heads and down-weight noisy or redundant ones. You can read off head importance by looking at the column norms of the `W_O` columns that read from each head's slice; pruning research consistently finds that low-importance heads have small `W_O` column norms. Removing `W_O` preserves the shape but loses both effects. Perplexity collapses on every standard benchmark. The cost of `W_O` is exactly one `(d_model, d_model)` matrix per layer, roughly equal to the combined parameter count of Q, K, and V. That is a real fraction of the model, and it earns its keep.

TL;DR

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.

Memory aid

Easy to grasp

Imagine eight researchers writing reports into separate columns of a shared spreadsheet. Without W_O, every researcher's findings stay locked in their own column forever; nobody can build on anyone else's work. W_O is the editor who reads all eight columns and rewrites them into a single mixed document where insights from researcher 3 can reinforce insights from researcher 7, and where the editor can quietly 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 quality is wildly different. Take the editor away and you lose both the cross-pollination and the quality control.

Key concepts

Multi-head attention Residual stream Head specialization

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, and it is the only place where the model can 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.

This deep dive walks the per-head slice structure that motivates W_O, then explains the two roles (cross-head mixing and learned weighting) in mechanistic detail, surveys the empirical evidence from pruning and interpretability research, and closes with the parameter cost vs benefit math.

Mental model: the concat is bookkeeping. W_O is where the heads actually 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, which is what lets us batch 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, and no operation has mixed it with the other heads.

Why this matters 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 residual stream itself would partition into h disjoint sub-streams, one per head, and downstream layers would have to do all the cross-head composition themselves through the FFN.

The architectural choice between 'mix at the attention sub-layer via W_O' and 'mix at the FFN sub-layer' is not a draw. The former is dramatically more effective empirically, which is why every production transformer has W_O.

Role 1: cross-head mixing

Role 2: learned per-head weighting

What happens when you remove W_O

Situations where this technique stops working.

2–4 min · Everything important, quickly.

python

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) showed that composition between heads in different layers depends on W_O writing to shared residual subspaces.
Voita et al. 2019, 'Analyzing Multi-Head Self-Attention', identified prunable heads via W_O column-norm analysis.

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 you would lose cross-head mixing entirely while keeping the parameter cost. Empirically this performs barely better than no W_O at all, because the dense off-diagonal blocks are where the cross-head composition lives.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

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.