Multi-Head Attention

Learning Objectives

By the end of this chapter, you will be able to:

1. Explain why a single attention head creates a representational bottleneck and why multiple heads running in parallel overcome it.
2. Derive the multi-head attention formula $\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W^O$ from first principles, understanding every symbol and matrix dimension.
3. Compute per-head attention weights and outputs by hand for our shared 5-token sentence using $H = 2$ heads.
4. Compare how different heads learn different attention patterns from the same input — syntactic relationships in one head, semantic proximity in another.
5. Implement a complete, runnable multi-head attention class in both NumPy and PyTorch that you can use to simulate any number of heads on any input.
6. Connect multi-head attention to Flash Attention, KV-cache, positional encodings, and transformer scaling laws.

Where this appears: Multi-head attention is the standard attention module used in virtually every transformer — GPT-4, Claude, Gemini, LLaMA, ViT, DALL-E, AlphaFold, and Codex all use multi-head attention as their core computational building block. A single GPT-4 layer uses 96 or more heads. Understanding how and why heads split the representation is essential for understanding every modern architecture.

The Real Problem

The Single-Head Bottleneck

In Chapter 1, we computed scaled dot-product attention using a single set of $Q$, $K$, $V$ matrices. The result was a single $5 \times 5$ attention weight matrix — one pattern describing which tokens attend to which.

But a single pattern is a severe bottleneck. Consider the word “cat” in “The cat sat on mat”. A human instantly recognizes that “cat” is:

1. Syntactically the subject of the verb “sat” (subject-verb relationship)
2. Grammatically a noun that follows the determiner “The” (determiner-noun agreement)
3. Semantically an animate entity related to the location “mat” (entity-location relation)
4. Positionally close to “The” and “sat” (local context)

A single attention head computes one weighted combination of all keys for each query. In our Chapter 1 computation, “cat” assigned weight 0.4026 to “The”, 0.2442 to “sat”, and only 0.0898 to itself. That is one compromise pattern — it cannot simultaneously attend heavily to “The” (for grammar) and to “sat” (for semantics) through separate, specialized lenses. The softmax forces a single probability distribution over all keys.

Vaswani et al. (2017) articulated this limitation precisely in the original transformer paper: “Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.”

The Committee Analogy

Think of multi-head attention as a committee of experts reviewing a document. A single reviewer reads the entire document through one lens and produces one summary. A committee of $H$ experts — a grammarian, a semanticist, a coreference resolver, a positional tracker — each reads the same document but focuses on their specialty. Their reports are concatenated into a richer, more nuanced understanding than any single reviewer could produce.

Crucially, the experts do not talk to each other while reading. Each head operates independently and in parallel on a different subspace of the representation. Only after all heads produce their outputs are the results combined.

From Intuition to Mathematics

The Projection Idea

The core mathematical insight is that a $d_{\text{model}}$-dimensional vector space can be decomposed into $H$ lower-dimensional subspaces. If $d_{\text{model}} = 512$ and $H = 8$, then each head operates in a $d_k = 512/8 = 64$-dimensional subspace.

Each head $i$ gets its own set of learned projection matrices: $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. These projections are learned during training — the model discovers which subspaces are useful for which kinds of relationships.

Split vs. Learned Projection

In our teaching example, we split the dimensions directly: Head 1 gets dims $[0,1]$, Head 2 gets dims $[2,3]$. This is equivalent to using identity projections where $W_1^Q$ selects columns 0-1 and $W_2^Q$ selects columns 2-3. In real transformers, the projection matrices are dense learned parameters that can mix dimensions in arbitrary ways, discovering subspaces not aligned with any coordinate axis.

The Mathematical Definition

The multi-head attention formula from Vaswani et al. (2017):

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) \, W^O$

where $\text{head}_i = \text{Attention}(Q W_i^Q, \; K W_i^K, \; V W_i^V)$

Symbol-by-Symbol Breakdown

What the Formula Says in Plain English

1. Project the input into $H$ different subspaces using learned weight matrices $W_i^Q, W_i^K, W_i^V$.
2. Attend independently in each subspace using the scaled dot-product attention from Chapter 1: $\text{softmax}(Q_h K_h^\top / \sqrt{d_k}) \, V_h$.
3. Concatenate the $H$ output vectors into a single vector of the original dimensionality.
4. Project through $W^O$ to allow the model to mix and recombine information from different heads.

Why Multiple Heads Work

The Representational Subspace Argument

Each head operates on a $d_k$-dimensional subspace. In this subspace, the dot product $Q_h K_h^\top$ measures a different notion of similarity than the full-space dot product. A token pair that has a high score in Head 1's subspace (capturing syntactic proximity) may have a low score in Head 2's subspace (capturing semantic type).

In our example, consider the “on” token. In Head 1 (dims 0-1), $Q_{\text{on}} = [0.0, 0.0]$ — the zero vector. Every dot product is zero, so Head 1 assigns uniform attention (0.2 to each token). But in Head 2 (dims 2-3), $Q_{\text{on}} = [1.0, 1.0]$ actively attends to “on” itself (0.3673) where the key strongly matches. Two completely different behaviors from the same token, enabled by subspace specialization.

Capacity vs. Specialization Trade-off

There is a fundamental trade-off: more heads means more specialization, but each head has fewer dimensions to work with. With $H = 1$, the single head has $d_k = 4$ dimensions for computing similarity. With $H = 4$, each head has only $d_k = 1$ — a scalar dot product with very limited discriminative power. The sweet spot depends on the task and model size. In practice, $d_k = 64$ or $d_k = 128$ per head is typical.

Interactive: Head Splitting Pipeline

Step through the multi-head attention pipeline below. Click each step to see how the full matrices are split into per-head subspaces, processed independently, and concatenated back.

Step-by-Step Calculation

We now compute multi-head attention with $H = 2$ heads on our shared example. Head 1 operates on dims $[0, 1]$ (so $d_k = 2$) and Head 2 operates on dims $[2, 3]$. The scaling factor for each head is $\sqrt{d_k} = \sqrt{2} \approx 1.4142$.

Token: “The” (row 0)

Head 1 — $Q_{h1}[0] = [1.0, 0.0]$, operating on $K_{h1}$ (columns 0-1 of K):

Head 1 output for “The”: $O_{h1}[0] = [0.2491, 0.3763]$

Head 2 — $Q_{h2}[0] = [1.0, 0.0]$, operating on $K_{h2}$ (columns 2-3 of K):

Head 2 output for “The”: $O_{h2}[0] = [0.2289, 0.3663]$

Key observation: Head 1 distributes attention equally across “cat”, “sat”, and “mat” (all 0.2509). Head 2 focuses on “cat” and “on” (both 0.2711) while giving “sat” much less weight (0.1337). These are genuinely different attention patterns extracted from different dimensions of the same representation.

Concatenation: $O[0] = [O_{h1}[0], O_{h2}[0]] = [0.2491, 0.3763, 0.2289, 0.3663]$

Token: “cat” (row 1)

Head 1: $Q_{h1}[1] = [0.0, 2.0]$ — this query strongly activates dim 1, matching keys with high values in dim 1.

Head 1: “cat” strongly attends to “The” (0.3664) and “sat” (0.3664) — its grammatical neighbours. Output: $O_{h1}[1] = [0.4109, 0.1336]$.

Head 2: $Q_{h2}[1] = [0.0, 1.0]$ — activates dim 3 (the second dim in Head 2's subspace).

Head 2: “cat” attends to “The” (0.2711) and “on” (0.2711) — different from Head 1. Output: $O_{h2}[1] = [0.2289, 0.3663]$.

Concatenation

For each token, concatenate the two head outputs to recover the original $d_{\text{model}} = 4$ dimensionality:

Full Attention Weights and Output

Head 1 Attention Weights ($5 \times 5$)

Head 1 patterns: “cat” strongly attends to “The” and “sat” (both 0.3664). “sat” attends most to itself (0.3673). “on” is perfectly uniform (0.2000) because $Q_{\text{on}} = [0, 0]$ in this subspace.

Head 2 Attention Weights ($5 \times 5$)

Head 2 patterns: “on” now strongly attends to itself (0.3673) — completely different from its uniform pattern in Head 1. “sat” attends to “cat” and “on” (both 0.2711), not to itself.

Averaged Attention Weights (for comparison)

Output Matrix — Multi-Head ($5 \times 4$)

Compare with the single-head output from Chapter 1. The multi-head output has dims 0-1 shaped by Head 1's attention pattern and dims 2-3 shaped by Head 2's pattern. Each pair of dimensions carries contextual information computed through a different lens.

Interactive: Head Comparison Heatmap

Hover any cell to compare attention weights across Head 1, Head 2, single-head, and the difference between heads. Switch between side-by-side and individual views to see how each head specializes.

Interactive: Head Count Explorer

Change the number of heads to see how splitting affects attention patterns. With $H = 1$, you get one 4D pattern (Chapter 1). With $H = 4$, each head has only $d_k = 1$ — observe how limited scalar dot products produce repetitive patterns.

Applications Across Domains

Natural Language Processing

In large language models, different heads specialize for different linguistic phenomena. Clark et al. (2019) found that in BERT: certain heads track subject-verb agreement (head 8-10 in layer 6), others resolve coreference (“she” → “Mary”), and others attend to the previous token (positional heads). Voita et al. (2019) showed that pruning task-irrelevant heads has minimal impact on accuracy, confirming that heads truly specialize.

Computer Vision

In Vision Transformers (ViT), image patches become tokens. Dosovitskiy et al. (2021) observed that some heads attend to nearby patches (local texture), others attend to distant patches with similar color (global color coherence), and still others track edges and contours. This mirrors how CNNs learn different filters — but the multi-head mechanism learns these patterns from data rather than requiring architectural inductive biases.

Code Generation

In code models like Codex and CodeLLaMA, heads specialize for code-specific patterns: variable-definition tracking (where was $x$ defined?), bracket matching (which opening bracket does this closing bracket match?), type flow (what type does this expression evaluate to?), and import resolution. The parallel nature of multi-head attention lets the model track all of these simultaneously.

Scientific Modeling

In AlphaFold's structure module, multi-head attention over amino acid residues lets different heads capture: spatial proximity in 3D (nearby in folded structure), sequence proximity (nearby in primary sequence), evolutionary co-variation (residues that mutate together), and chemical interactions (hydrogen bonds, hydrophobic contacts). The 48-head architecture captures fundamentally different notions of “relevance” between residue pairs.

Connection to Modern Systems

Flash Attention

Flash Attention (Dao et al., 2022) does not change the mathematics of multi-head attention — it produces bit-identical results. What it changes is the memory access pattern. Standard multi-head attention materializes the full $(N \times N)$ score matrix for each head, requiring $O(H \cdot N^2)$ memory. Flash Attention tiles the computation so that each SRAM block computes a partial softmax, never materializing the full matrix. For $H = 96$ heads and $N = 8192$, this reduces peak memory from ~6 GB to ~50 MB.

KV-Cache Optimization

During autoregressive generation, the KV-cache stores previously computed key and value tensors. With multi-head attention, each head has its own KV-cache of shape $(\text{seq\_len}, d_k)$. The total cache size is $2 \times H \times L \times N \times d_k$ (2 for K and V, $L$ layers, $N$ tokens). For GPT-3 with 96 heads, 96 layers, and 2048 tokens, the KV-cache alone requires ~3 GB in fp16. This memory pressure is what motivates Multi-Query Attention (Chapter 3) and Grouped-Query Attention (Chapter 4), which share keys/values across heads.

Positional Encodings (RoPE)

Rotary Position Embeddings (Su et al., 2021) apply per-head position-dependent rotations to Q and K before the dot product. Because each head operates in its own$d_k$-dimensional subspace, RoPE rotates pairs of dimensions within each head at different frequencies. Lower-frequency rotations in some heads capture long-range position, while higher-frequency rotations in others capture local position. The multi-head structure is what makes this frequency decomposition possible.

Transformer Scaling

As model size grows, both $d_{\text{model}}$ and $H$ increase. GPT-3 (175B) uses $d_{\text{model}} = 12288$ with $H = 96$, giving $d_k = 128$. LLaMA-2 70B uses $d_{\text{model}} = 8192$ with $H = 64$ heads, also $d_k = 128$. The scaling laws (Kaplan et al., 2020) show that increasing both width and heads together improves sample efficiency, but the per-head dimension $d_k$ rarely drops below 64 in practice to maintain per-head expressiveness.

Complexity Analysis

The key insight: multi-head attention has the same asymptotic time complexity as single-head attention because $H \cdot d_k = d_{\text{model}}$. What changes is that the $H$ per-head computations can run in parallel on GPUs, and the model gains $H$ independent attention patterns instead of one.

Python Implementation

The full Python class below implements multi-head attention with the same structure as Chapter 1. Click any line to see its execution trace with actual values.

      The   cat   sat    on   mat
The  0.00  1.00  1.00  0.00  1.00
cat  2.00  0.00  2.00  0.00  0.00
sat  1.00  1.00  2.00  0.00  1.00
on   0.00  0.00  0.00  0.00  0.00
mat  0.00  1.00  1.00  0.00  1.00

      The   cat   sat    on   mat
The  0.00  1.00  0.00  1.00  0.50
cat  1.00  0.00  0.00  1.00  0.50
sat  0.00  1.00  0.00  1.00  0.50
on   1.00  1.00  0.00  2.00  1.00
mat  1.00  0.00  0.00  1.00  0.50

       The     cat     sat      on     mat
The  0.0000  0.7071  0.7071  0.0000  0.7071
cat  1.4142  0.0000  1.4142  0.0000  0.0000
sat  0.7071  0.7071  1.4142  0.0000  0.7071
on   0.0000  0.0000  0.0000  0.0000  0.0000
mat  0.0000  0.7071  0.7071  0.0000  0.7071

       The     cat     sat      on     mat
The  0.0000  0.7071  0.0000  0.7071  0.3536
cat  0.7071  0.0000  0.0000  0.7071  0.3536
sat  0.0000  0.7071  0.0000  0.7071  0.3536
on   0.7071  0.7071  0.0000  1.4142  0.7071
mat  0.7071  0.0000  0.0000  0.7071  0.3536

       The      cat      sat       on      mat
The  0.1237   0.2509   0.2509   0.1237   0.2509
cat  0.3664   0.0891   0.3664   0.0891   0.0891
sat  0.1811   0.1811   0.3673   0.0893   0.1811
on   0.2000   0.2000   0.2000   0.2000   0.2000
mat  0.1237   0.2509   0.2509   0.1237   0.2509

       The      cat      sat       on      mat
The  0.1337   0.2711   0.1337   0.2711   0.1904
cat  0.2711   0.1337   0.1337   0.2711   0.1904
sat  0.1337   0.2711   0.1337   0.2711   0.1904
on   0.1811   0.1811   0.0893   0.3673   0.1811
mat  0.2711   0.1337   0.1337   0.2711   0.1904

       d0       d1
The  0.2491   0.3763
cat  0.4109   0.1336
sat  0.2717   0.2717
on   0.3000   0.3000
mat  0.2491   0.3763

       d2       d3
The  0.2289   0.3663
cat  0.2289   0.3663
sat  0.2289   0.3663
on   0.1799   0.4579
mat  0.2289   0.3663

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

      d0   d1
The  1.0  0.0
cat  0.0  2.0
sat  1.0  1.0
on   0.0  0.0
mat  1.0  0.0

      d2   d3
The  1.0  0.0
cat  0.0  1.0
sat  1.0  0.0
on   1.0  1.0
mat  0.0  1.0

      d0   d1
The  0.0  1.0
cat  1.0  0.0
sat  1.0  1.0
on   0.0  0.0
mat  1.0  0.0

      d2   d3
The  0.0  1.0
cat  1.0  0.0
sat  0.0  0.0
on   1.0  1.0
mat  0.5  0.5

      d0   d1
The  1.0  0.0
cat  0.0  1.0
sat  0.0  0.0
on   0.0  0.0
mat  0.5  0.5

      d2   d3
The  0.0  0.0
cat  0.0  0.0
sat  1.0  0.0
on   0.0  1.0
mat  0.5  0.5

        d0       d1       d2       d3
The  0.2491   0.3763   0.2289   0.3663
cat  0.4109   0.1336   0.2289   0.3663
sat  0.2717   0.2717   0.2289   0.3663
on   0.3000   0.3000   0.1799   0.4579
mat  0.2491   0.3763   0.2289   0.3663

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       The      cat      sat       on      mat
The  0.1237   0.2509   0.2509   0.1237   0.2509
cat  0.3664   0.0891   0.3664   0.0891   0.0891
sat  0.1811   0.1811   0.3673   0.0893   0.1811
on   0.2000   0.2000   0.2000   0.2000   0.2000
mat  0.1237   0.2509   0.2509   0.1237   0.2509

       The      cat      sat       on      mat
The  0.1337   0.2711   0.1337   0.2711   0.1904
cat  0.2711   0.1337   0.1337   0.2711   0.1904
sat  0.1337   0.2711   0.1337   0.2711   0.1904
on   0.1811   0.1811   0.0893   0.3673   0.1811
mat  0.2711   0.1337   0.1337   0.2711   0.1904

        d0       d1       d2       d3
The  0.2491   0.3763   0.2289   0.3663
cat  0.4109   0.1336   0.2289   0.3663
sat  0.2717   0.2717   0.2289   0.3663
on   0.3000   0.3000   0.1799   0.4579
mat  0.2491   0.3763   0.2289   0.3663

1import numpy as np
2import math
3
4class MultiHeadAttention:
5    """
6    Multi-Head Attention (Vaswani et al., 2017)
7
8    Splits Q, K, V into H heads, runs scaled dot-product attention
9    on each head independently, then concatenates the results.
10
11    MultiHead(Q, K, V) = Concat(head_1, ..., head_H) @ W_O
12    head_i = Attention(Q @ W_i^Q, K @ W_i^K, V @ W_i^V)
13    """
14
15    def __init__(self, d_model: int, num_heads: int):
16        """
17        Args:
18            d_model:   Total model dimension (4 in our example)
19            num_heads: Number of attention heads H (2 in our example)
20        """
21        self.d_model = d_model
22        self.num_heads = num_heads
23        self.d_k = d_model // num_heads
24        self.scale = math.sqrt(self.d_k)
25
26    def _softmax(self, x: np.ndarray) -> np.ndarray:
27        """Numerically stable softmax along last axis."""
28        x_shifted = x - np.max(x, axis=-1, keepdims=True)
29        exp_x = np.exp(x_shifted)
30        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
31
32    def split_heads(self, M: np.ndarray) -> list:
33        """Split matrix M (N, d_model) into H matrices of (N, d_k)."""
34        heads = []
35        for h in range(self.num_heads):
36            start = h * self.d_k
37            end = start + self.d_k
38            heads.append(M[:, start:end])
39        return heads
40
41    def attention(self, Qh: np.ndarray, Kh: np.ndarray, Vh: np.ndarray):
42        """Single-head scaled dot-product attention."""
43        scores = Qh @ Kh.T
44        scaled = scores / self.scale
45        weights = self._softmax(scaled)
46        output = weights @ Vh
47        return weights, output
48
49    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
50        """
51        Full multi-head forward pass.
52
53        Args:
54            Q: Query matrix  (N, d_model)
55            K: Key matrix    (N, d_model)
56            V: Value matrix  (N, d_model)
57
58        Returns:
59            all_weights: List of H attention matrices, each (N, N)
60            output:      Concatenated output (N, d_model)
61        """
62        Q_heads = self.split_heads(Q)
63        K_heads = self.split_heads(K)
64        V_heads = self.split_heads(V)
65
66        all_weights = []
67        all_outputs = []
68
69        for h in range(self.num_heads):
70            w, o = self.attention(Q_heads[h], K_heads[h], V_heads[h])
71            all_weights.append(w)
72            all_outputs.append(o)
73
74        output = np.hstack(all_outputs)
75        return all_weights, output
76
77    def explain(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray,
78                tokens: list, query_idx: int = 0):
79        """Print a detailed trace for a specific query token."""
80        all_weights, output = self.forward(Q, K, V)
81        Q_heads = self.split_heads(Q)
82        K_heads = self.split_heads(K)
83
84        token = tokens[query_idx]
85        print(f"\n=== Multi-Head trace for '{token}' (row {query_idx}) ===")
86
87        for h in range(self.num_heads):
88            start = h * self.d_k
89            end = start + self.d_k
90            print(f"\n--- Head {h+1} (dims {start}:{end}, d_k={self.d_k}) ---")
91            print(f"  Q_h{h+1}[{token}] = {Q_heads[h][query_idx]}")
92            for j, t in enumerate(tokens):
93                dot = Q_heads[h][query_idx] @ K_heads[h][j]
94                print(f"  Q[{token}] . K[{t}] = {dot:.4f} -> /{self.scale:.2f} = {dot/self.scale:.4f}")
95            print(f"  weights = {np.round(all_weights[h][query_idx], 4)}")
96
97        print(f"\nConcatenated output[{token}] = {np.round(output[query_idx], 4)}")
98
99
100# ── Shared Example (used in every chapter) ──
101tokens = ["The", "cat", "sat", "on", "mat"]
102
103Q = np.array([
104    [1.0, 0.0, 1.0, 0.0],   # The
105    [0.0, 2.0, 0.0, 1.0],   # cat
106    [1.0, 1.0, 1.0, 0.0],   # sat
107    [0.0, 0.0, 1.0, 1.0],   # on
108    [1.0, 0.0, 0.0, 1.0],   # mat
109])
110
111K = np.array([
112    [0.0, 1.0, 0.0, 1.0],   # The
113    [1.0, 0.0, 1.0, 0.0],   # cat
114    [1.0, 1.0, 0.0, 0.0],   # sat
115    [0.0, 0.0, 1.0, 1.0],   # on
116    [1.0, 0.0, 0.5, 0.5],   # mat
117])
118
119V = np.array([
120    [1.0, 0.0, 0.0, 0.0],   # The
121    [0.0, 1.0, 0.0, 0.0],   # cat
122    [0.0, 0.0, 1.0, 0.0],   # sat
123    [0.0, 0.0, 0.0, 1.0],   # on
124    [0.5, 0.5, 0.5, 0.5],   # mat
125])
126
127# ── Run ──
128mha = MultiHeadAttention(d_model=4, num_heads=2)
129all_weights, output = mha.forward(Q, K, V)
130
131print("Head 1 Attention Weights (5x5):")
132print(np.round(all_weights[0], 4))
133
134print("\nHead 2 Attention Weights (5x5):")
135print(np.round(all_weights[1], 4))
136
137print("\nConcatenated Output (5x4):")
138print(np.round(output, 4))
139
140# Detailed trace for "The" (token 0)
141mha.explain(Q, K, V, tokens, query_idx=0)

PyTorch Implementation

The PyTorch version adds GPU support, automatic differentiation, batched inputs, and includes the learned projection matrices ($W^Q, W^K, W^V, W^O$) that real transformers use. We set $\texttt{use\_projections=False}$ to verify results match the NumPy version exactly.

       The      cat      sat       on      mat
The  0.1237   0.2509   0.2509   0.1237   0.2509
cat  0.3664   0.0891   0.3664   0.0891   0.0891
sat  0.1811   0.1811   0.3673   0.0893   0.1811
on   0.2000   0.2000   0.2000   0.2000   0.2000
mat  0.1237   0.2509   0.2509   0.1237   0.2509

       The      cat      sat       on      mat
The  0.1337   0.2711   0.1337   0.2711   0.1904
cat  0.2711   0.1337   0.1337   0.2711   0.1904
sat  0.1337   0.2711   0.1337   0.2711   0.1904
on   0.1811   0.1811   0.0893   0.3673   0.1811
mat  0.2711   0.1337   0.1337   0.2711   0.1904

        d0       d1       d2       d3
The  0.2491   0.3763   0.2289   0.3663
cat  0.4109   0.1336   0.2289   0.3663
sat  0.2717   0.2717   0.2289   0.3663
on   0.3000   0.3000   0.1799   0.4579
mat  0.2491   0.3763   0.2289   0.3663

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class MultiHeadAttention(nn.Module):
7    """
8    Multi-Head Attention (Vaswani et al., 2017) — PyTorch
9
10    Splits Q, K, V into H heads, runs scaled dot-product attention
11    on each, then concatenates and projects with W_O.
12
13    Supports GPU, automatic differentiation, and batched inputs.
14    """
15
16    def __init__(self, d_model: int, num_heads: int):
17        super().__init__()
18        self.d_model = d_model
19        self.num_heads = num_heads
20        self.d_k = d_model // num_heads
21        self.scale = math.sqrt(self.d_k)
22
23        # Learned projection matrices (not used in this example,
24        # but included for production completeness)
25        self.W_Q = nn.Linear(d_model, d_model, bias=False)
26        self.W_K = nn.Linear(d_model, d_model, bias=False)
27        self.W_V = nn.Linear(d_model, d_model, bias=False)
28        self.W_O = nn.Linear(d_model, d_model, bias=False)
29
30    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
31        """Reshape (B, N, d_model) -> (B, H, N, d_k)."""
32        B, N, _ = x.shape
33        return x.view(B, N, self.num_heads, self.d_k).transpose(1, 2)
34
35    def forward(
36        self,
37        Q: torch.Tensor,
38        K: torch.Tensor,
39        V: torch.Tensor,
40        mask: torch.Tensor | None = None,
41        use_projections: bool = False,
42    ) -> tuple[torch.Tensor, torch.Tensor]:
43        """
44        Args:
45            Q: (B, N, d_model) or (N, d_model)
46            K: (B, N, d_model) or (N, d_model)
47            V: (B, N, d_model) or (N, d_model)
48            mask: Optional (B, 1, N, N) or (N, N) boolean
49            use_projections: If True, apply W_Q/W_K/W_V/W_O
50
51        Returns:
52            all_weights: (B, H, N, N) attention weights
53            output: (B, N, d_model) multi-head output
54        """
55        # Handle unbatched input
56        if Q.dim() == 2:
57            Q = Q.unsqueeze(0)
58            K = K.unsqueeze(0)
59            V = V.unsqueeze(0)
60
61        # Optional learned projections
62        if use_projections:
63            Q = self.W_Q(Q)
64            K = self.W_K(K)
65            V = self.W_V(V)
66
67        # Split into heads: (B, N, d_model) -> (B, H, N, d_k)
68        Q_h = self.split_heads(Q)
69        K_h = self.split_heads(K)
70        V_h = self.split_heads(V)
71
72        # Scaled dot-product per head
73        scores = torch.matmul(Q_h, K_h.transpose(-2, -1)) / self.scale
74
75        if mask is not None:
76            scores = scores.masked_fill(mask, float("-inf"))
77
78        weights = F.softmax(scores, dim=-1)  # (B, H, N, N)
79        head_out = torch.matmul(weights, V_h)  # (B, H, N, d_k)
80
81        # Concatenate heads: (B, H, N, d_k) -> (B, N, d_model)
82        B, H, N, d_k = head_out.shape
83        output = head_out.transpose(1, 2).contiguous().view(B, N, H * d_k)
84
85        if use_projections:
86            output = self.W_O(output)
87
88        return weights.squeeze(0), output.squeeze(0)
89
90
91# ── Shared Example ──
92tokens = ["The", "cat", "sat", "on", "mat"]
93
94Q = torch.tensor([
95    [1.0, 0.0, 1.0, 0.0],
96    [0.0, 2.0, 0.0, 1.0],
97    [1.0, 1.0, 1.0, 0.0],
98    [0.0, 0.0, 1.0, 1.0],
99    [1.0, 0.0, 0.0, 1.0],
100])
101K = torch.tensor([
102    [0.0, 1.0, 0.0, 1.0],
103    [1.0, 0.0, 1.0, 0.0],
104    [1.0, 1.0, 0.0, 0.0],
105    [0.0, 0.0, 1.0, 1.0],
106    [1.0, 0.0, 0.5, 0.5],
107])
108V = torch.tensor([
109    [1.0, 0.0, 0.0, 0.0],
110    [0.0, 1.0, 0.0, 0.0],
111    [0.0, 0.0, 1.0, 0.0],
112    [0.0, 0.0, 0.0, 1.0],
113    [0.5, 0.5, 0.5, 0.5],
114])
115
116# ── Run (without learned projections for exact comparison) ──
117mha = MultiHeadAttention(d_model=4, num_heads=2)
118weights, output = mha(Q, K, V, use_projections=False)
119
120print("Head 1 weights (5x5):")
121print(weights[0].round(decimals=4))
122
123print("\nHead 2 weights (5x5):")
124print(weights[1].round(decimals=4))
125
126print("\nOutput (5x4):")
127print(output.round(decimals=4))

Key Takeaways

1. The problem: A single attention head produces one attention pattern — a single compromise that cannot simultaneously specialize for syntax, semantics, and position.
2. The solution: Run $H$ independent heads in parallel, each operating on a $d_k$-dimensional subspace of the representation.
3. Same cost: Total computation is identical to single-head because $H \cdot d_k = d_{\text{model}}$. The work is just distributed across heads.
4. Concatenation recovers dimensionality: Each head outputs $(N, d_k)$; concatenating $H$ heads gives $(N, d_{\text{model}})$.
5. $W^O$ is essential: The output projection mixes information across heads. Without it, each output dimension carries information from only one head.
6. Heads specialize: In trained models, different heads learn to capture syntactic dependencies, semantic similarity, positional proximity, and coreference — empirically verified by Clark et al. (2019) and Voita et al. (2019).

Exercises

Exercise 1: Compute for “sat”

Compute the multi-head attention output for the token “sat” (row 2) by hand. Show the per-head dot products, scaled scores, softmax weights, and per-head outputs. Verify that Head 1 attends most to “sat” itself (0.3673) while Head 2 distributes attention to “cat” and “on” (both 0.2711). Concatenate the outputs and check against the output matrix.

Exercise 2: $H = 4$ Heads

Repeat the computation with $H = 4$ heads, each with $d_k = 1$. What happens to the attention patterns when each head computes only a scalar dot product? Compare the expressiveness with $H = 2$. Use the Head Count Explorer to verify your answers.

Exercise 3: Identical Subspaces

What happens if you assign both heads the same columns (e.g., both use dims $[0, 1]$)? Compute the attention weights for both heads and verify they are identical. What does this tell you about the role of diverse subspaces?

Exercise 4: $W^O$ Projection

Create a $4 \times 4$ output projection matrix $W^O$ (e.g., a random orthogonal matrix or a permutation matrix) and apply it to the concatenated output. How does the final output change? Why is this step necessary for mixing information across heads?

Exercise 5: Head Pruning

Set one head's output to zero (simulating pruning). How much does the final output change? Compute the L2 norm of the difference. Voita et al. (2019) found that many heads can be pruned with minimal accuracy loss — does this match your observation?

References

1. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., & Polosukhin, I. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems, 30. The paper that introduced multi-head attention as part of the transformer architecture.
2. Clark, K., Khandelwal, U., Levy, O., & Manning, C.D. (2019). What Does BERT Look At? An Analysis of BERT's Attention. BlackboxNLP Workshop at ACL 2019. Empirical analysis showing that different BERT heads specialize for different linguistic phenomena.
3. Voita, E., Talbot, D., Moiseev, F., Sennrich, R., & Titov, I. (2019). Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. ACL 2019. Showed that a small number of heads are important for specific tasks and the rest can be pruned.
4. Dao, T., Fu, D.Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. IO-aware algorithm for exact attention computation without materializing the full score matrix.
5. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. Per-head rotary embeddings that encode position through rotation in each head's subspace.
6. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021. Vision Transformer showing multi-head attention on image patches.
7. Kaplan, J., McCandlish, S., Henighan, T., Brown, T.B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., & Amodei, D. (2020). Scaling Laws for Neural Language Models. Empirical scaling laws showing how performance improves with model size, dataset size, and compute.
8. Michel, P., Levy, O., & Neubig, G. (2019). Are Sixteen Heads Really Better than One? NeurIPS 2019. Systematic study of head importance, showing that many heads can be removed at test time with negligible accuracy loss.