Multi-Head Attention: Full Mechanics

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Multi-Head Formulation
5. Parameter Count
6. Why Multiple Heads?
7. Head Specialization
8. The Output Projection
9. Relation to Tensor Products
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive multi-head attention from single-head attention.
2. Compute the exact parameter count for MHA.
3. Prove that MHA with one head of dimension $d$ is more expressive than single-head attention.
4. Explain empirically observed head specialization patterns.
5. Derive the computational equivalence of parallel heads vs sequential computation.

Notation

• $d$ — model dimension
• $H$ — number of heads
• $d_k = d_v = d/H$ — per-head dimension
• $\mathbf{W}_i^Q, \mathbf{W}_i^K \in \mathbb{R}^{d \times d_k}$ — per-head projection matrices
• $\mathbf{W}^O \in \mathbb{R}^{d \times d}$ — output projection

Core Intuition

A single attention head can only compute one type of similarity between positions. But understanding language requires attending to multiple relationships simultaneously: syntactic structure, semantic similarity, positional proximity, coreference, etc. Multi-head attention runs $H$ independent attention operations in parallel, each in a lower-dimensional subspace, allowing the model to jointly attend to information from different representation subspaces.

The Multi-Head Formulation

where each head:

Step by step for input $\mathbf{X} \in \mathbb{R}^{T \times d}$:

1. Project: $\mathbf{Q}_i = \mathbf{X}\mathbf{W}_i^Q \in \mathbb{R}^{T \times d_k}$ (similarly for K, V).
2. Attend: $\text{head}_i = \text{softmax}(\mathbf{Q}_i\mathbf{K}_i^T/\sqrt{d_k})\mathbf{V}_i \in \mathbb{R}^{T \times d_v}$.
3. Concatenate: $[\text{head}_1; \ldots; \text{head}_H] \in \mathbb{R}^{T \times d}$.
4. Output project: Multiply by $\mathbf{W}^O \in \mathbb{R}^{d \times d}$.

Parameter Count

For each head:

• $\mathbf{W}_i^Q$: $d \times d_k$ parameters
• $\mathbf{W}_i^K$: $d \times d_k$ parameters
• $\mathbf{W}_i^V$: $d \times d_v$ parameters

Across $H$ heads:

• Q projections: $H \times d \times d_k = d \times d$ (since $H \cdot d_k = d$)
• K projections: $d \times d$
• V projections: $d \times d$
• Output projection: $d \times d$

This is identical to a single head with full dimension $d$ (same compute budget, split across heads).

Why Multiple Heads?

Theorem (Expressiveness). Multi-head attention with $H$ heads is strictly more expressive than single-head attention of the same total dimension.

Intuitive argument: Each head produces a different attention pattern $\mathbf{A}_i$. The output is a weighted combination of values retrieved using different similarity metrics. A single head can only produce one attention pattern per layer.

Formal argument: Consider the output at position $t$:

where $\mathbf{W}^O_i$ is the $i$-th column block of $\mathbf{W}^O$. Each term uses a different attention distribution $\alpha^{(i)}$ — this cannot be replicated by a single softmax.

Head Specialization

Empirical analysis of trained transformers reveals systematic specialization:

Positional heads: Attend to fixed relative positions (e.g., always attend to the previous token, or token 2 positions back). These implement a form of learned n-gram patterns.

Syntactic heads: Track grammatical structure — subject-verb agreement, modifier-noun relationships. Identifiable by attention aligning with dependency parse trees.

Induction heads: Two-head circuits that implement in-context learning:

1. Head A: copies previous token identity.
2. Head B: attends to positions where Head A's output matches current context → "if A follows B before, then A follows B again."

Retrieval heads: In later layers, attend to semantically relevant tokens regardless of position (long-range factual recall).

The Output Projection

Why $\mathbf{W}^O$? Without it, each head's contribution is restricted to its $d_v$-dimensional subspace. The output projection allows:

• Mixing information across heads.
• Projecting the concatenated output back to model dimension.
• Learning which head combinations are useful.

Alternative view: The full MHA operation can be written as:

where $\mathbf{W}_i^O \in \mathbb{R}^{d_v \times d}$ is the $i$-th row block of $\mathbf{W}^O$.

Relation to Tensor Products

Multi-head attention can be viewed through the lens of tensor decomposition. The attention operation for all heads simultaneously:

where $\otimes$ denotes the outer product in the token dimension, and $\mathbf{W}_i^{VO} = \mathbf{W}_i^V\mathbf{W}_i^O$ is the combined value-output projection.

This is a rank-$H$ decomposition of the full attention tensor — each head contributes one rank-1 component in the "attention pattern × value mixing" space.

Common Pitfalls

Pitfall 1. Assuming more heads = better. Beyond a threshold, increasing $H$ while keeping $d$ fixed makes each head dimension $d/H$ too small for meaningful attention patterns. Typical sweet spot: $d_k = 64$–$128$.

Pitfall 2. Ignoring the output projection. Removing $\mathbf{W}^O$ (using raw concatenation) significantly hurts performance because heads cannot communicate.

Pitfall 3. Treating heads as independent. While computed in parallel, heads interact through the residual stream — later layers read the combined output of all heads from previous layers.

Summary

• Multi-head attention runs $H$ parallel attention operations in $d_k$-dimensional subspaces.
• Total parameters: $4d^2$ (same budget as single-head, but more expressive).
• Heads specialize: positional, syntactic, induction, retrieval patterns emerge from training.
• Output projection $\mathbf{W}^O$ enables cross-head information mixing.
• MHA as rank-$H$ tensor decomposition of the full attention operation.

Exercises

Exercise 1. For a model with $d=1024, H=16$, compute $d_k$ and the total MHA parameters (including output projection).

Exercise 2. Prove that $\text{Concat}(\text{head}_1, \ldots, \text{head}_H)\mathbf{W}^O$ cannot be replicated by a single attention head of dimension $d$ (hint: different attention patterns per subspace).

Exercise 3. Compute the FLOPs for multi-head attention with $T=2048, d=4096, H=32$.

Exercise 4. Describe how an "induction head" circuit works using two attention heads. What attention patterns must each head learn?

Exercise 5. If we reduce $H$ from 32 to 8 while keeping $d=4096$ constant, what changes in terms of per-head expressiveness and total parameter count?