Multihead Attention Part 2
How do multiple head outputs recombine, and how do 4D tensors flow through multi-head attention?
How this lesson unfolds.
After individual attention heads compute their independent contextual vectors, they must be seamlessly recombined. Learn how concatenation and the output projection matrix W_O fuse multi-head signals, and trace 4D GPU tensor shapes step-by-step.
How does the final output projection matrix W_O synthesize isolated multi-head attention outputs back into a unified token stream?
Imagine a crime investigation where 8 specialist detectives investigate a case: one examines financial records, another inspects phone metadata, another analyzes ballistic trajectories, and another interviews witnesses. If each detective merely files a separate, isolated report in their own desk drawer, the police department never solves the crime. The chief of police gathers all 8 reports, binds them side-by-side, and writes a synthesized master briefing that connects the bank transfer to the ballistics timing. The output projection W_O is that chief of police.
- Learning goalExplain the subspace isolation problem of multi-head attention and formulate the concatenation and linear projection steps using W_O ∈ R^d_(model) × d_(model).
- Helpful foundationsMulti-head dimension splitting · Tensor concatenation and linear projections
- Common trapAssuming that head outputs are summed together. Summing head outputs would overwrite and average distinct head signals; concatenation preserves all distinct features before W_O performs controlled linear blending.
- Next connectionNext, trace the exact concatenation tensor shapes, linear mapping formulas, and parameter ledgers in Part B.
More context & mechanics
While multi-head attention allows parallel specialization, it introduces the Subspace Isolation Problem:
- The Subspace Isolation Dilemma:
Each attention head i computes an output vector Oᵢ = Aᵢ Vᵢ ∈ R^(dₖ). Because each head operates in its own private coordinate frame, Head 1 cannot directly communicate with Head 2.
- Horizontal Concatenation:
To assemble their findings without loss of information, all h head outputs are concatenated side-by-side along the channel dimension:
O_(concat) = [O₁, O₂, ṡ, O_h] ∈ R^(h · dₖ) = R^d_(model)
- Synthesis via Output Projection (W_O):
Simply concatenating vectors places them adjacent in memory, but does not allow feature interaction. Multiplying by trainable projection matrix W_O ∈ R^d_(model) × d_(model) computes linear combinations across all heads:
MultiHeadOutput = O_(concat) W_O
Any feature from Head 1 can now modulate or combine with features from Head 8.
In GPT-2 and Llama-3, W_O has the exact same dimensions and parameter count as a single Q, K, or V projection (d_(model) × d_(model)). During model pruning, analyzing the rows of W_O reveals which attention heads contribute most heavily to downstream performance.
The concatenation-and-projection synthesis design was formalized by Vaswani et al. in 2017. Prior multi-head proposals attempted to average head outputs; Vaswani's insight that a trainable linear projection W_O should learn how to weight and blend head outputs unlocked significant gains in linguistic expressivity.
- Concat binds, W_O blends: Concatenation glues heads together; the output projection allows them to talk.
- Back to d_(model): The output projection restores the exact dimensional width required by the residual stream.
How do we derive the parameter count, tensor dimensions, and matrix multiplication equations for the multi-head output projection?
Think of assembling a team of 12 musicians playing 12 different instruments. Each musician plays a 64-beat measure (12 × 64 = 768 notes total). You record all 12 audio tracks simultaneously into a 768-track mixing console. The master sound engineer (W_O) adjusts a 768 × 768 grid of mixing knobs to produce the final balanced orchestral chord that fits into the main stereo bus.
- $$Multi-head attention synthesizes $h$ tensors of shape $[B, T, d_k]$ into $[B, T, d_{model}]$ and multiplies by $W_O \in \mathbb{R}^{d_{model} \times d_{model}}$, consuming $d_{model}^2$ parameters.$$
- Learning goalDerive the tensor shape transitions, parameter accounting, and linear projection algebra for multi-head output synthesis using W_O.
- Helpful foundationsPyTorch memory layout and `.contiguous()` · Linear algebra matrix transformations
- Common trapAssuming that `.contiguous()` is just a boilerplate call. In PyTorch, transpose changes stride metadata without reordering bytes; calling `.contiguous()` is strictly required before `.view()` to physically rearrange memory in C-order.
- Next connectionNext, discover how multi-head attention and feed-forward networks assemble into complete transformer blocks in Chapter 19 Screen 1.
More context & mechanics
#### 1. General Mathematical Formulation
The complete multi-head attention operation is formulated as:
MultiHead(Q, K, V) = Concat( head₁, head₂, ṡ, head_h ) W_O
Where each individual head is:
headᵢ = Attention(Q Wᵢ^Q, K Wᵢ^K, V Wᵢ^V) ∈ R^(B × T × dₖ)
- Input from Heads: 4D tensor O ∈ R^(B × h × T × dₖ).
- Permutation: Transpose dimensions 1 and 2:
O_(perm) = O.transpose(1, 2) ∈ R^(B × T × h × dₖ)
- Contiguous Reshape (Flattening Heads):
O_(flat) = O_(perm).contiguous().view(B, T, h · dₖ) = R^B × T × d_(model)
- Output Linear Projection:
Y = O_(flat) W_O + b_O ∈ R^B × T × d_(model)
For GPT-2 small (d_(model) = 768, h = 12, dₖ = 64):
- W_O matrix shape: [768, 768]
- Weight parameters: 768 × 768 = 589,824
- Bias parameters: 768
- Total parameters in W_O: 590,592
- FLOPs per token: 2 × d_(model)² = 2 × 768 × 768 = 1,179,648 FLOPs
Let B = 1, T = 2, h = 2, dₖ = 2 (d_(model) = 4).
Let head outputs for token 1 be:
- head₁ = [1.0, 2.0]
- head₂ = [3.0, 4.0]
Concatenation yields:
O_(concat) = [1.0, 2.0, 3.0, 4.0]
Let output matrix W_O have diagonal weights [0.5, 0.5, 0.5, 0.5] and off-diagonal cross-head blending terms:
Y = O_(concat) W_O = [1.0, 2.0, 3.0, 4.0]
0.5 & 0.1 & 0.0 & 0.0 \
0.0 & 0.5 & 0.1 & 0.0 \
0.1 & 0.0 & 0.5 & 0.0 \
0.0 & 0.0 & 0.1 & 0.5
= [0.8, 1.1, 2.1, 2.0]
The output vector is a synthesized 4-dimensional representation combining signals across both heads.
In Megatron-LM tensor parallel implementations, the multi-head attention block is split across GPUs: QKV projections use ColumnParallelLinear, while W_O uses RowParallelLinear. This elegant mathematical duality means all-reduce communication is only required after W_O, minimizing inter-GPU interconnect traffic.
The Row/Column parallel decomposition of attention layers was discovered by Shoeybi et al. (NVIDIA) in 2019 in the Megatron-LM architecture, making W_O the central synchronization anchor for training 100B+ parameter models.
- [B, T, h · dₖ] → [B, T, d_(model)]: Flatten the head dimension, then project.
- Row-parallel anchor: W_O accumulates distributed head slices back into a single residual tensor.
How do 4D tensor shapes transform through Multi-Head Attention during GPU batched execution?
Imagine packing a shipping container. Instead of running 12 separate cargo flights for 12 heads, the GPU reshapes the memory layout into a 4-dimensional box: (Batch, Heads, Tokens, Head-Dimension). In PyTorch, a single batch matrix multiplication (bmm) computes all attention dot products across all batches and all heads simultaneously at lightning speed, before transposing and flattening back to standard sequence format.
- Learning goalTrace the exact 4D tensor shapes (B, T, d_model) -> (B, h, T, dₖ) -> (B, h, T, T) -> (B, T, d_model) throughout GPU multi-head execution.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · PyTorch Tensor Reshaping: view, transpose, contiguous · Batched Matrix Multiplication (torch.matmul)
- Common trapUsing a Python `for` loop over heads. A Python loop is catastrophically slow on GPUs. Modern frameworks reshape tensors to (B, h, T, dₖ) and execute all heads in parallel via a single batched GEMM call.
- Next connectionWith Multi-Head Attention fully assembled and tensor-tracked, we integrate it into the complete Transformer Block alongside Layer Normalization and Feed-Forward Networks.
More context & mechanics
Trace the tensor journey through Multi-Head Attention where B = Batch Size, T = Sequence Length, h = Heads, d = d_(model), and dₖ = d / h:
- Input:
X ∈ (B, T, d)
- Linear Projections (Q, K, V):
Q = X W_Q, K = X W_K, V = X W_V ∈ (B, T, d)
- Reshape & Transpose for Multi-Head Parallelism:
- Reshape to (B, T, h, dₖ)
- Transpose dims 1 and 2:
Q, K, V ∈ (B, h, T, dₖ)
- Batched Scaled Dot Products:
S = (Q · K^\top) / (√(dₖ)) ∈ (B, h, T, dₖ) × (B, h, dₖ, T) ⟹ (B, h, T, T)
- Causal Mask & Softmax:
A = softmax(S + M, dim=-1) ∈ (B, h, T, T)
- Batched Value Aggregation:
Z = A · V ∈ (B, h, T, T) × (B, h, T, dₖ) ⟹ (B, h, T, dₖ)
- Transpose & Contiguous Reshape:
- Transpose back: (B, T, h, dₖ)
- Contiguous flatten: Z_(concat) ∈ (B, T, h · dₖ) = (B, T, d)
- Output Projection (W_O):
Output = Z_(concat) W_O ∈ (B, T, d)
This shape pipeline is universal across PyTorch, Hugging Face Transformers, JAX/Flax, and TensorRT-LLM.
The `(B, h, T, dₖ)` tensor layout was pioneered in the original `tensor2tensor` library by Vaswani, Kaiser, and Shazeer in 2017 to fully exploit NVIDIA Volta tensor cores for batched matrix multiplication.
- The 4D Tensor Ledger:
- (B, T, d) → (B, h, T, dₖ) (Split)
- (B, h, T, T) (Attention matrix)
- (B, h, T, dₖ) (Context values)
- (B, T, d) (Recombined output)
How does the final output projection matrix W_O synthesize concatenated head outputs back into the unified residual stream?
Imagine 12 specialized detectives investigating a crime. Detective 1 found the weapon; Detective 2 checked the alibi; Detective 3 traced the phone records. Right now, each detective holds their own separate notepad. If they simply tape their notepads together side by side (concatenation), they have a long pile of paper, but no combined verdict. The Output Projection Matrix W_O acts as the lead prosecutor: it reads all 12 notepads simultaneously, cross-examines their findings, and writes a single unified case summary that fits neatly into the official case file.
- $$Multiply concatenated head vectors by W_O in R^(h*d_v x d_model) to linearly recombine isolated subspace channels into the residual highway.$$
- Learning goalAnalyze the block-matrix formulation of the output projection W_O and compute synthesized feature vectors from concatenated head outputs.
- Helpful foundationsMatrix block multiplication · Attention Is All You Need (Vaswani et al., 2017)
- Common trapThinking concatenation alone is sufficient for multi-head attention; without W_O, features in head 1 and head 2 cannot linearly interact until the next layer.
- Next connectionAs features flow from attention into the residual stream, how does LayerNorm reshape their geometric distribution? We examine hyperplanes and hyperspheres.
More context & mechanics
Let headᵢ = Attention(Q Wᵢ^Q, K Wᵢ^K, V Wᵢ^V) ∈ R^(T × d_v).
The concatenation across all h heads yields:
H_(concat) = [head₁, head₂, ṡ, head_h] ∈ R^(T × (h · d_v))
Because h · d_v = d_(model), this concatenated tensor has the correct dimension, but each feature channel remains strictly isolated within its head slice.
The output projection matrix W_O ∈ R^d_(model) × d_(model) performs a full-rank linear synthesis:
y = H_(concat) W_O
Expanding W_O as h vertical blocks W_O^((i)) ∈ R^d_v × d_(model):
y = ∑_(i=1)^h headᵢ W_O^((i))
This allows features discovered in Head 1 to linearly interact with and modulate features discovered in Head 12 before being added to the residual stream.
Without W_O, the subsequent MLP feed-forward layer would be forced to perform cross-head synthesis. Having W_O directly inside the attention block ensures the attention layer outputs a fully integrated representation.
Vaswani et al. (2017) included W^O as an integral component of Multi-Head Attention, ensuring linear expressiveness before the residual addition step.
- Concatenate to gather; project through W_O to synthesize.
- W_O blends 12 separate detective notes into one cohesive story.