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

Short answer·Medium·4.0 · 0·~3 min

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.

Quick reference

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

Detailed answer & concept explanation

~6 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.

Quick scan · 2–4 min

Everything important, quickly.

Explain the per-head slice structure of the concatenated output, walk through how W_O lets heads mix across slices, show how column norms encode per-head weighting, connect to induction-head composition and head pruning, and close with parameter cost vs benefit.

Code example

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-world examples

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.
Michel et al. 2019, 'Are Sixteen Heads Really Better Than One?', pruned ~40% of heads with minimal loss using importance scores derived from W_O.
Llama 4 Maverick, Claude Opus 4.7, and GPT-5.5 all retain W_O at full d_model x d_model rank; nobody has shipped a frontier model without it.
TransformerLens and other mechanistic-interpretability tooling decomposes W_O into per-head write matrices for analysis.

Interviewer follow-ups

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?

QWhat does W_O look like after head pruning?

QHow does W_O interact with GQA?

Don't say this

Red flags and common mistakes 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 cheat sheet

The night-before-the-interview bullets. Scan these on the way to the call.

References

Primary sources. Skim if you want the original framing.