Multi-Head Attention

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. From Single-Head to Multi-Head
5. Mathematical Formulation
6. The Projection Geometry
7. Why Multiple Heads Help
8. Computational Complexity
9. Attention Patterns Across Heads
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive multi-head attention from single-head by partitioning the embedding space.
2. Show that multi-head attention with $H$ heads and dimension $d_k = d/H$ has the same FLOPs as single-head.
3. Prove that multi-head attention is strictly more expressive than single-head of the same dimension.
4. Analyze the rank of attention matrices and how heads increase effective rank.
5. Explain Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) as efficiency variants.

Notation

• $H$ — number of attention heads
• $d$ — model dimension (embedding size)
• $d_k = d/H$ — dimension per head
• $\mathbf{W}_i^Q, \mathbf{W}_i^K \in \mathbb{R}^{d \times d_k}$ — query/key projections for head $i$
• $\mathbf{W}_i^V \in \mathbb{R}^{d \times d_k}$ — value projection for head $i$
• $\mathbf{W}^O \in \mathbb{R}^{d \times d}$ — output projection
• $\text{head}_i$ — output of the $i$-th attention head

Core Intuition

A single attention head computes one weighted average of values based on one notion of "relevance" (one Q-K dot product). But language has multiple simultaneous relationships: syntactic (subject-verb), semantic (noun-modifier), positional (adjacent words). Multi-head attention runs multiple attention operations in parallel, each attending to different types of relationships, then combines the results.

From Single-Head to Multi-Head

Single-head attention with dimension $d$:

Problem: The softmax produces a rank-1 (per query) attention pattern. A single head can focus on only one "aspect" at a time.

Solution: Use $H$ independent attention operations in lower-dimensional subspaces:

where each head operates in $\mathbb{R}^{d_k}$:

Mathematical Formulation

Given input $\mathbf{X} \in \mathbb{R}^{T \times d}$ (sequence of $T$ tokens):

For each head $i \in \{1, \ldots, H\}$:

Final output:

The Projection Geometry

Each head projects the $d$-dimensional input into a $d_k$-dimensional subspace. The projections $\mathbf{W}_i^Q, \mathbf{W}_i^K$ define the "similarity metric" for head $i$:

The matrix $\mathbf{W}_i^Q(\mathbf{W}_i^K)^T \in \mathbb{R}^{d \times d}$ defines a learned bilinear form — each head learns a different notion of "similarity."

Why Multiple Heads Help

Theorem. Multi-head attention is strictly more expressive than single-head attention of the same total dimension.

Argument. Single-head with dimension $d$ computes one attention matrix $\mathbf{A} \in \mathbb{R}^{T \times T}$. Multi-head computes $H$ different matrices $\mathbf{A}_1, \ldots, \mathbf{A}_H$, each potentially attending to different positions. The output combines $H$ different weighted averages of $H$ different value projections — a richer representation.

Rank argument. A single softmax row has rank 1 (it's a probability vector applied to the value matrix). With $H$ heads contributing rank-$d_k$ outputs each, the total output can have rank up to $d = H \cdot d_k$, enabling richer representations.

Empirical observation: Different heads learn different functions:

• Some heads attend to adjacent tokens (local patterns).
• Some attend to syntactic dependencies (long-range).
• Some attend to specific positions (positional heads).

Computational Complexity

Single-head ($d$ dimensions): Computing $\mathbf{Q}\mathbf{K}^T$ costs $O(T^2 d)$. Total: $O(T^2 d)$.

Multi-head ($H$ heads, $d_k = d/H$ each): Each head costs $O(T^2 d_k)$. Over $H$ heads: $O(H \cdot T^2 \cdot d/H) = O(T^2 d)$.

Same total cost — multi-head attention does not increase FLOPs compared to single-head of the same dimension. The parallelism across heads maps naturally to GPU tensor operations.

Parameters: $\mathbf{W}^Q, \mathbf{W}^K, \mathbf{W}^V$ each $\in \mathbb{R}^{d \times d}$, plus $\mathbf{W}^O \in \mathbb{R}^{d \times d}$. Total: $4d^2$.

Attention Patterns Across Heads

Multi-Query Attention (MQA): Share $\mathbf{K}, \mathbf{V}$ projections across all heads; only $\mathbf{Q}$ is per-head. Reduces KV-cache by factor $H$.

Grouped-Query Attention (GQA): Group heads into $G$ groups; each group shares $\mathbf{K}, \mathbf{V}$. Interpolates between MHA ($G=H$) and MQA ($G=1$). Used in LLaMA-2, Mistral.

Parameter comparison:

• MHA: $4d^2$ parameters
• MQA: $2d^2 + 2d \cdot d_k = 2d^2(1 + 1/H)$
• GQA: $2d^2 + 2d^2/G$

Common Pitfalls

Pitfall 1. Thinking heads are independent after training. While initialized independently, heads communicate through the residual stream and $\mathbf{W}^O$ — they learn complementary roles.

Pitfall 2. Using too many heads with too small $d_k$. If $d_k < 16$, individual heads lack capacity to compute meaningful attention. GPT-3 uses $d=12288, H=96 \Rightarrow d_k=128$.

Pitfall 3. Not scaling by $1/\sqrt{d_k}$. Without this, dot products grow with dimension, pushing softmax into saturation (near-one-hot outputs).

Summary

• Multi-head attention runs $H$ parallel attention operations in $d_k$-dimensional subspaces.
• Each head learns a different "similarity function" via its Q, K projections.
• Same computational cost as single-head, but strictly more expressive.
• The output projection $\mathbf{W}^O$ mixes information across heads.
• GQA/MQA trade off expressiveness for KV-cache efficiency.

Exercises

Exercise 1. Verify that multi-head attention has the same FLOP count as single-head attention of dimension $d$.

Exercise 2. Show that if all heads use identical projections ($\mathbf{W}_i^Q = \mathbf{W}_j^Q$ for all $i,j$), multi-head reduces to single-head.

Exercise 3. Compute the total parameter count for a multi-head attention layer with $d=512, H=8$.

Exercise 4. Derive the gradient $\partial\mathcal{L}/\partial\mathbf{W}_i^Q$ through the softmax attention.

Exercise 5. For GQA with $G=4$ groups and $H=32$ heads, how many independent K/V projections are there? What is the KV-cache size relative to full MHA?