Multi-Head Attention: Splitting, Merging, and Projecting

Multi-head attention runs h independent attention computations in parallel, each on a slice of the model dimension, then concatenates the results and projects back to d. This lesson builds it from scratch in PyTorch: combined in-projection, head splitting, parallel attention, head merging, and output projection. We verify against nn.MultiheadAttention with identical weights.

Why not a single head?

Before we build the mechanics, it is worth understanding why we split into multiple heads at all.

A single attention head produces one attention distribution over the sequence — a single set of weights that says “how much should this token attend to each previous token.” This forces an either/or compromise. If a token like "bank" needs to attend to both its subject ("river") and a prepositional modifier ("near"), or if a verb needs to track both its subject (for agreement) and a tense clue, a single head must average or choose between them. It cannot maintain two independent focusing patterns simultaneously.

Multiple attention heads solve this by giving the model multiple independent projection subspaces. Each head learns its own query, key, and value projections. One head can specialize in syntactic agreement (subject–verb), another in semantic similarity (word–concept), another in positional recency (recent tokens), and another in long-range dependencies. The output projection then mixes these independent information streams back together.

This is not just a theoretical argument: ablation studies on trained models (e.g., Voita et al., 2019; Michel et al., 2019) show that different heads consistently specialize in different linguistic phenomena — some track syntax, some track coreference, some track positional patterns. Removing specific heads degrades specific capabilities, confirming that the parallel subspaces are functionally distinct.

The one-sentence version

Q/K/V are projected once, split into h heads, attention runs in parallel per head, outputs are concatenated, and a final projection mixes information across heads.

Everything else in this lesson makes that concrete, with shapes tracked at every step.

Lab: Run the companion CPU exercise at /labs/track1/01-attention/ (head split, merge, output projection parity).

A consistent shape trace

Use these dimensions throughout the lesson:

Symbol	Meaning	Value in this trace
B	Batch size	2
T	Sequence length	5
h	Number of attention heads	4
d_head	Width per head (d / h)	8
d	Model width (h × d_head)	32

The numbers are deliberately small. A real model uses the same pattern with larger values; the pattern is what matters.

1. Combined in-projection

Lesson 4 used separate W_Q, W_K, W_V each [d, d_head]. In multi-head (and PyTorch’s nn.MultiheadAttention), we combine them into a single matrix W_in of shape [d, 3d]:

QKV = X @ W_in        # [B, T, 3d]
Q, K, V = QKV.chunk(3, dim=-1)  # Each [B, T, d]

This is mathematically equivalent to three separate projections but more memory-efficient (one matmul instead of three).

2. Head splitting

Each of Q, K, V has shape [B, T, d] where d = h × d_head. We split the last dimension into h heads and move the head dimension to position 1 for parallel batched matmul:

Q = Q.view(B, T, h, d_head).transpose(1, 2)  # [B, h, T, d_head]
K = K.view(B, T, h, d_head).transpose(1, 2)  # [B, h, T, d_head]
V = V.view(B, T, h, d_head).transpose(1, 2)  # [B, h, T, d_head]

Why transpose(1, 2)?

Moving the head dimension to index 1 gives shape [B, h, T, d_head]. Batched matrix multiplication then computes attention for all h heads simultaneously via Q @ K.transpose(-2, -1) → [B, h, T, T]. Without the transpose, heads would not be separated in the matmul.

3. Parallel attention per head

With shapes [B, h, T, d_head], scaled dot-product attention runs identically for all heads:

scores = Q @ K.transpose(-2, -1) / (d_head ** 0.5)  # [B, h, T, T]

# Causal mask: [1, 1, T, T] broadcasts to [B, h, T, T]
causal = torch.tril(torch.ones(T, T, dtype=torch.bool))
scores = scores.masked_fill(~causal, float('-inf'))

attn = torch.softmax(scores, dim=-1)  # [B, h, T, T]
out  = attn @ V                       # [B, h, T, d_head]

Each head’s attention weights form a valid distribution over the sequence (row sums = 1).

4. Head merging

After attention, we merge heads back into the model dimension. The critical step is contiguous() before view():

out = out.transpose(1, 2).contiguous().view(B, T, d)  # [B, T, d]

• transpose(1, 2) moves head dimension back: [B, T, h, d_head]
• contiguous() copies to a new memory layout because transpose makes the tensor non-contiguous
• view(B, T, d) flattens the last two dimensions: [B, T, h × d_head] = [B, T, d]

Why contiguous()?

transpose() returns a view with swapped strides. view() requires a contiguous memory layout matching the new shape. Without contiguous(), PyTorch raises RuntimeError: view size is not compatible with input tensor's size and stride. This is a common bug; always call contiguous() after transpose() before view().

5. Output projection

The merged representation is projected back to d by W_O of shape [d, d]:

output = out @ W_O  # [B, T, d]

In the original Transformer (Vaswani et al.), W_O has no bias. Some implementations (e.g., PyTorch’s nn.MultiheadAttention) include a bias by default — set bias=False to match the paper.

6. Fully worked numeric example (verified, seed 123)

Setup: B=2, T=5, h=4, d_head=8, d=32, torch.manual_seed(123), torch.set_default_dtype(torch.float32).

import torch
torch.manual_seed(123)
torch.set_default_dtype(torch.float32)

B, T, h, d_head = 2, 5, 4, 8
d = h * d_head  # 32

X = torch.randn(B, T, d)
W_in = torch.randn(d, 3 * d)  # Combined in-projection
W_O = torch.randn(d, d)       # Output projection

# In-projection
QKV = X @ W_in                    # [2, 5, 96]
Q, K, V = QKV.chunk(3, dim=-1)   # Each [2, 5, 32]

# Split heads
Q = Q.view(B, T, h, d_head).transpose(1, 2)  # [2, 4, 5, 8]
K = K.view(B, T, h, d_head).transpose(1, 2)  # [2, 4, 5, 8]
V = V.view(B, T, h, d_head).transpose(1, 2)  # [2, 4, 5, 8]

# Causal mask
causal = torch.tril(torch.ones(T, T, dtype=torch.bool))  # [5, 5]
causal = causal[None, None, :, :]  # [1, 1, 5, 5]

# Scaled dot-product attention
scores = Q @ K.transpose(-2, -1) / (d_head ** 0.5)  # [2, 4, 5, 5]
scores = scores.masked_fill(~causal, float('-inf'))
attn = torch.softmax(scores, dim=-1)  # [2, 4, 5, 5]
out = attn @ V                        # [2, 4, 5, 8]

# Merge heads
out = out.transpose(1, 2).contiguous().view(B, T, d)  # [2, 5, 32]

# Output projection
output = out @ W_O  # [2, 5, 32]

Shape verification:

X:              [2, 5, 32]
QKV:            [2, 5, 96]
Q, K, V:        [2, 5, 32] each
Q, K, V split:  [2, 4, 5, 8] each
scores:         [2, 4, 5, 5]
attn:           [2, 4, 5, 5] (each row sums to 1)
out (per head): [2, 4, 5, 8]
out merged:     [2, 5, 32]
output:         [2, 5, 32]

Verification: Matches nn.MultiheadAttention(embed_dim=32, num_heads=4, batch_first=True, bias=False) with weights copied to in_proj_weight and out_proj.weight, using additive causal mask, to max_diff < 2e-5 (PyTorch 2.5, CPU, float32, seed 123).

Multi-head attention

In-projection → head split → parallel attention → head merge → output projection, with tensor shapes at each step.

Read this diagram as text

Text equivalent. Panel A: input X [B,T,d] projects via combined W_in [d,3d] to Q,K,V [B,T,d], then chunks to three [B,T,d]. Panel B: each of Q,K,V reshaped via view+transpose to [B,h,T,d_head]. Panel C: parallel attention computes scores [B,h,T,T], applies causal mask [1,1,T,T], softmax, then attn @ V → [B,h,T,d_head]. Panel D: heads merged via transpose+contiguous+view to [B,T,d], then output projection W_O [d,d] produces final [B,T,d].

7. PyTorch reference implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, bias=False):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads

        # Combined in-projection
        self.in_proj = nn.Linear(d_model, 3 * d_model, bias=bias)
        # Output projection
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)

    def forward(self, X, attn_mask=None):
        B, T, _ = X.shape

        # In-projection + chunk
        QKV = self.in_proj(X)  # [B, T, 3d]
        Q, K, V = QKV.chunk(3, dim=-1)  # Each [B, T, d]

        # Split heads: [B, T, h, d_head] -> [B, h, T, d_head]
        Q = Q.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
        K = K.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
        V = V.view(B, T, self.num_heads, self.d_head).transpose(1, 2)

        # Scaled dot-product attention
        scores = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5)  # [B, h, T, T]

        if attn_mask is not None:
            scores = scores.masked_fill(~attn_mask, float('-inf'))

        attn = torch.softmax(scores, dim=-1)  # [B, h, T, T]
        out = attn @ V                        # [B, h, T, d_head]

        # Merge heads: [B, h, T, d_head] -> [B, T, h, d_head] -> [B, T, d]
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)

        # Output projection
        output = self.out_proj(out)  # [B, T, d]
        return output, attn

Numerical check against nn.MultiheadAttention:

nn.MultiheadAttention uses a different API:

• in_proj_weight (not in_proj.weight) for the combined projection
• attn_mask is additive (added to scores), so use -inf for masked positions
• Call signature for self-attention: ref(X, X, X, attn_mask=...)

# Weight copying for exact match
ref = nn.MultiheadAttention(embed_dim=32, num_heads=4, batch_first=True, bias=False)
with torch.no_grad():
    ref.in_proj_weight.copy_(W_in.T)      # [96, 32] <- [96, 32]
    ref.out_proj.weight.copy_(W_O.T)      # [32, 32] <- [32, 32]

# Create additive causal mask: 0 for attend, -inf for don't attend
causal_mask = torch.where(
    torch.tril(torch.ones(T, T, dtype=torch.bool)),
    torch.tensor(0.0),
    torch.tensor(float('-inf'))
)

ref_out, _ = ref(X, X, X, attn_mask=causal_mask)
# Measured numerical tolerance with seed 123, float32: max_diff ~ 1.5e-5
assert (output - ref_out).abs().max() < 2e-5

8. Failure modes

Failure mode	Symptom	Root cause	Test
Wrong split order	Silent wrong attention; shapes look right	view(B, h, T, d_head) instead of view(B, T, h, d_head)	Assert Q.shape == [B, T, h, d_head] before transpose
Missing transpose	Heads not separated; attention mixes across heads	No transpose(1, 2) after split	Assert Q.shape == [B, h, T, d_head] after transpose
Missing contiguous()	RuntimeError: view size not compatible	Transpose makes tensor non-contiguous	Always call .contiguous() before .view() after transpose
Wrong merge order	Heads interleaved incorrectly	view(B, T, d) before transpose(1, 2)	Assert merged shape is [B, T, d] with d == h * d_head
Mask shape mismatch	Broadcast error or masks wrong positions	Mask not shaped [1, 1, T, T] or [B, 1, T, T]	Print mask.shape; must broadcast to scores [B, h, T, T]
Output proj bias	Extra parameters; diverges from paper	nn.Linear defaults to bias=True	Use bias=False in both in_proj and out_proj

9. What this lesson does not cover

• FlashAttention / kernel fusion — optimization, not mechanism
• Grouped-query / multi-query attention — advanced variants (fewer K/V heads)
• Relative position / ALiBi / RoPE in multi-head — covered in positional encoding lesson
• Dropout on attention weights — training concern
• KV caching for multi-head — covered in KV cache lesson

10. Exercises

Exercise	Prompt
1	Given B=3, T=10, h=8, d_head=16: what is d? What are the shapes of W_in and W_O?
2	After Q.view(B, T, h, d_head).transpose(1, 2), what are the strides? Why does contiguous() matter before final view?
3	Write the single line that creates a causal mask broadcastable to scores [B, h, T, T].
4	If W_in is replaced by three separate W_Q, W_K, W_V each [d, d], how many total parameters in the in-projection? Is this more or less than combined W_in [d, 3d]?

11. Related lessons

Related lessons

• LLM Internals & Architecture · Attention mechanisms from scratchSelf-Attention: Queries, Keys, and Values14 min
• LLM Internals & Architecture · Attention mechanisms from scratchCausal and Padding Masks: Implementation and Tests15 min

12. Sources

Source	Relevance
Vaswani et al., “Attention Is All You Need” (arXiv:1706.03762), §3.2.2	Original multi-head attention: MultiHead(Q,K,V) = Concat(head_1,...,head_h)W^O where head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)
PyTorch nn.MultiheadAttention source (torch/nn/modules/activation.py)	Reference implementation: combined in_proj, head split via view+transpose, scaled_dot_product_attention, out_proj
PyTorch F.scaled_dot_product_attention docs (v2.5+)	Supports multi-head via [B, h, T, d_head]; is_causal flag for causal masking

---