Multihead Attention Part 1
Why are multiple attention heads better than one single head, and how is the embedding dimension partitioned?
How this lesson unfolds.
A single attention head can only attend to one relationship at a time, forcing compromise. Explore how multi-head attention splits the embedding dimension into parallel subspaces to capture syntax, semantics, and context simultaneously.
Why does a model need multiple concurrent attention heads instead of a single large attention head?
Imagine reading a detective novel. If you only had one highlighter, you would have to choose: do you highlight grammatical subjects, chronological clues, emotional motives, or character locations? A single attention head produces a single probability distribution per token—it can only look in one direction. With 12 or 32 heads, each head holds a distinct colored highlighter: Head 1 tracks grammatical verbs, Head 2 tracks pronoun referents, Head 3 tracks punctuation boundaries, and Head 4 tracks semantic synonyms, all in parallel.
- Learning goalExplain the expressive limitations of single-head attention and prove how multiple heads capture distinct linguistic subspaces simultaneously.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · Linear Subspaces & Direct Sum Decompositions · Scaled Dot-Product Attention Mechanics
- Common trapAssuming multiple heads simply repeat the same computation for redundancy. Because each head has independent projection matrices initialized randomly and optimized via gradient descent, each head specializes into distinct linguistic and relational features.
- Next connectionTo implement multiple heads without multiplying memory and FLOP requirements by the number of heads, we split the model embedding dimension across heads.
More context & mechanics
In a single attention head, the attention distribution α_(i) is a single convex combination over the sequence. If token i needs to attend to both its governing verb at position j and its qualifying adjective at position k, a single head can only produce an average blend:
zᵢ = α_(ij)vⱼ + α_(ik)vₖ
This averaging washes out distinct geometric signals into an indistinct composite vector.
Multi-Head Attention (MHA) resolves this by deploying h parallel heads:
headᵢ = Attention(X Wᵢ^Q, X Wᵢ^K, X Wᵢ^V)
Each head possesses its own linear projection matrices Wᵢ^Q, Wᵢ^K ∈ R^d_(model) × dₖ and Wᵢ^V ∈ R^d_(model) × d_v. This allows:
- Head 1 to assign 95% attention weight to the syntactic subject.
- Head 2 to assign 90% attention weight to the historical antecedent.
- Head 3 to attend strictly to local adjacent tokens.
None of these heads compromise their attention distributions.
Empirical studies of BERT and GPT models reveal that specific heads naturally specialize: Head 2 in Layer 4 tracks direct objects, while Head 8 in Layer 7 tracks long-range coreference links across paragraphs.
Vaswani et al. (2017) introduced Multi-Head Attention in Attention Is All You Need. They observed that multi-head attention significantly outperformed single-head attention with the same overall parameter count and computational budget.
- Golden Axiom: "One head averages; multiple heads specialize."
- Multiple heads provide multiple distinct perspectives on the same text.
Why is multi-head attention superior to a single massive attention head, and how does splitting the hidden dimension maintain compute budgets?
Imagine you are tasked with judging an international culinary competition. If you have only one super-expert judge, that judge must taste each dish and simultaneously evaluate taste, presentation, aroma, texture, and cultural authenticity all in one single bite; their attention gets overwhelmed and compromises are inevitable. Now imagine instead that you take the same judging budget and appoint 8 specialized judges: Judge 1 focuses solely on aroma, Judge 2 on visual plating, Judge 3 on seasoning, and Judge 4 on texture. Each judge operates in parallel within their narrow specialty, and their reports are compiled together. Multi-head attention is that panel of specialized judges.
- Learning goalExplain the representational advantages of multi-head attention and prove that partitioning d_(model) into h heads preserves total computational complexity.
- Helpful foundationsScaled dot-product attention · Matrix dimension partitioning
- Common trapAssuming that having 12 attention heads makes the model 12 times slower or requires 12 times more parameters. Because each head operates on dimension d_(model)/12, the total FLOPs and parameter counts remain identical to a single 768-dimensional head.
- Next connectionNext, trace the exact 4D tensor reshaping, transposition, and permutation contracts in Part B.
More context & mechanics
A single full-dimensional attention head suffers from Averaging Blindness:
- The Single-Focus Trap: A single attention head produces a single distribution α over the context. If token i needs to attend to its subject for syntactic agreement and simultaneously to its object for semantic meaning, a single head is forced to average these two signals, blunting both.
- The Multi-Head Solution: Instead of one attention head over R^d_(model), we split the budget into h independent heads, each operating over head dimension:
dₖ = (d_(model)) / (h)
- The Invariant Compute Budget:
- Single head: 1 head × (d_(model) × d_(model)) = d_(model)² operations.
- Multi-head: h heads × (dₖ × dₖ) = h × ((d_(model)) / (h))² = (d_(model)²) / (h) per head ⟹ h × (d_(model)²) / (h) = d_(model)² total!
The computational complexity and parameter budget remain exactly identical to a single-head model, while expressivity increases h-fold!
In GPT-2 small, d_(model) = 768 is partitioned across h = 12 heads (dₖ = 64). In Llama-3-70B, d_(model) = 8192 is split across h = 64 heads (dₖ = 128). Probing studies show that individual heads spontaneously specialize into distinct linguistic detectors: some track open quotation marks, others trace Coreference pronouns, and others track indentation depth in Python code.
Multi-Head Attention was invented by Ashish Vaswani et al. in 2017. Earlier attention mechanisms (Bahdanau 2014, Luong 2015) used a single attention vector. Vaswani realized that partitioning the hidden dimension into parallel heads created an ensemble of diverse representational subspaces with zero computational penalty.
- Partition the budget, multiply the perspectives: h heads of size d/h cost the exact same FLOPs as 1 giant head.
- Specialized subspaces: Different heads track different relationships across the sequence.
How do 3D batched tensors reshape and transpose into 4D multi-head tensors [B, h, T, dₖ] without copying memory?
Think of a deck of 52 playing cards laid out in a long single row. If you want to deal 4 hands of 13 cards each, you do not need to print new cards or buy new card stock. You simply mentally group them into 4 piles of 13. In computer memory, tensor reshaping is that mental regrouping: the underlying bytes stay frozen in place while the shape descriptor changes how coordinates are read.
- Learning goalDerive the 4D tensor reshaping and transposition pipeline [B, T, d_(model)] → [B, h, T, dₖ] and trace batched parallel matrix multiplications.
- Helpful foundationsPyTorch `view` vs `transpose` vs `contiguous` · Batched 4D tensor matrix multiplication
- Common trapBelieving that `transpose(1, 2)` can be skipped. Without transposing to `[B, h, T, dₖ]`, the matrix multiplication cannot treat h as an independent batch dimension, preventing standard batched BLAS execution.
- Next connectionNext, learn how parallel head outputs are concatenated and projected back into a single unified stream in Chapter 18 Screen 1.
More context & mechanics
#### 1. The 4D Shape Transition Pipeline
Let input projected tensor be X ∈ R^B × T × d_(model).
Let h be the number of attention heads, and dₖ = d_(model) / h.
The transition executes in two strict steps:
- Reshape (View):
X_(view) = view(B, T, h, dₖ)
This operation merely updates the tensor stride metadata, taking O(1) time with zero data copying in memory.
- Permutation (Transpose):
X_(perm) = transpose(1, 2) ⟹ R^(B × h × T × dₖ)
Swapping dimension 1 (T) and dimension 2 (h) places head dimension h as an independent batch dimension.
With Q, K ∈ R^(B × h × T × dₖ), computing attention across all B batches and all h heads simultaneously reduces to a single standard batched matrix multiplication:
S = (Q K^T) / (√(dₖ)) ⟹ [B, h, T, dₖ] × [B, h, dₖ, T] = [B, h, T, T]
Applying row-wise softmax across the final dimension yields attention weights:
A = softmax(S) ∈ R^(B × h × T × T)
Multiplying by V ∈ R^(B × h × T × dₖ):
Out = A V ∈ R^(B × h × T × dₖ)
Let B = 2, T = 4, d_(model) = 768, h = 12:
- Head dimension: dₖ = 768 / 12 = 64.
- Input shape: [2, 4, 768].
- After `.view(2, 4, 12, 64)`: shape is [2, 4, 12, 64].
- After `.transpose(1, 2)`: shape is [2, 12, 4, 64].
- Q K^T: [2, 12, 4, 64] × [2, 12, 64, 4] = [2, 12, 4, 4] (12 attention matrices of size 4 × 4 per batch!).
- Output A V: [2, 12, 4, 4] × [2, 12, 4, 64] = [2, 12, 4, 64].
In PyTorch, this is written concisely:
```python
B, T, C = x.shape
q = self.q_proj(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
```
Modern CUDA libraries (cuBLAS, CUTLASS) treat the first two dimensions `[B, h]` as a combined batch stride, executing all heads in parallel on streaming multiprocessors.
The tensor permutation convention `[B, h, T, dₖ]` was introduced in tensor library frameworks like PyTorch and TensorFlow in 2017 to map multi-head attention directly onto existing batched BLAS Level 3 operations (`gemmStridedBatched`).
- View splits, transpose aligns: `view(B, T, h, dk)` cuts the channels; `transpose(1, 2)` brings heads forward.
- `[B, h, T, dₖ]`: The sacred four-dimensional shape of modern multi-head attention.
How does splitting embedding vectors across multiple attention heads allow simultaneous tracking of distinct semantic relationships?
Imagine a 768-lane mega-highway. If all cars must travel in one single collective swarm, everyone is forced to drive at the exact same average speed and focus on the same exit. But if you divide the highway into 12 dedicated lanes (64 lanes per head), Lane 1 can track grammatical syntax ('is this word a verb?'), Lane 2 can track pronouns ('who does she refer to?'), Lane 3 can track rhyming patterns, and Lane 4 can track emotional sentiment—all traveling at full speed in parallel without colliding.
- $$Partition d_model dimensions into h orthogonal subspaces of dimension d_k = d_model / h to perform h parallel self-attention operations.$$
- Learning goalDeconstruct high-dimensional embedding tensors into multi-head query-key-value slices and assemble concatenated outputs via projection.
- Helpful foundationsMatrix concatenation · Attention Is All You Need (Vaswani et al., 2017)
- Common trapBelieving multi-head attention multiplies the total compute cost by h; because each head operates on d_model / h dimensions, the total FLOPs are virtually identical to a single full-rank attention head.
- Next connectionAs signals pass through dozens of deep multi-head layers, their activations can drift wildly; how do we keep layer activations stable and calibrated? We apply Layer Normalization.
More context & mechanics
Instead of performing a single attention function with dimension d_(model), Multi-Head Attention projects queries, keys, and values h times with learned projection matrices:
headᵢ = Attention(Q Wᵢ^Q, K Wᵢ^K, V Wᵢ^V)
Where Wᵢ^Q, Wᵢ^K ∈ R^d_(model) × dₖ and Wᵢ^V ∈ R^d_(model) × d_v, typically with dₖ = d_v = d_(model) / h.
The independent head representations are concatenated and projected back:
MultiHead(Q, K, V) = Concat(head₁, ṡ, head_h) W^O
Where W^O ∈ R^h d_v × d_(model).
In GPT-3 (175B), d_(model) = 12288 is divided across h = 96 attention heads, with each head operating in a dₖ = 128 subspace. This allows the model to simultaneously monitor 96 distinct semantic and syntactic channels.
Introduced by Vaswani et al. in 2017 ('Attention Is All You Need'), multi-head attention replaced complex recurrent cells (LSTMs) and convolutional sequence models across all NLP tasks.
- Split the vector into lanes; each lane watches a different feature.
- Concatenate the lanes; project through W_O to blend their findings.
How does the direct sum of attention head subspaces decompose the full embedding space into orthogonal feature manifolds?
Imagine building a luxury apartment complex with 768 total rooms. If all 768 rooms share one giant open floor plan, noise from the kitchen drowns out the bedroom, the home theater, and the office. Multi-head attention builds soundproof walls, dividing the building into 12 private apartments of 64 rooms each. Apartment 1 handles cooking; Apartment 2 handles sleeping; Apartment 3 handles video editing. Each apartment operates in total privacy, yet together they span the entire complex.
- $$Decompose R^(d_model) into the internal direct sum of h disjoint subspaces H_1 + ...$$
- $$+ H_h, each of dimension d_k = d_model / h.$$
- Learning goalFormulate the direct sum decomposition of multi-head attention and prove the uniqueness of subspace representations.
- Helpful foundationsDirect sum of vector spaces · What Does BERT Look At? (Clark et al., 2019)
- Common trapBelieving attention heads compete for the same dimensions; by projecting into disjoint orthogonal subspaces, heads operate in parallel without destructive interference.
- Next connectionAfter each head finishes computing its private context in its own subspace, how do we fuse them back together? We inspect the output projection matrix W_O.
More context & mechanics
Let V = R^d_(model). Multi-Head Attention establishes h linear projections:
Pᵢ: V → Hᵢ ⊂ R^(dₖ), dₖ = (d_(model)) / (h)
The global representation space is isomorphic to the direct sum of head subspaces:
V ≅ H₁ ⊕ H₂ ⊕ ṡ ⊕ H_h
Where:
- Disjointness: Hᵢ ∩ Hⱼ = 0 for all i ≠ j.
- Dimension Sum: ∑_(i=1)^h dim(Hᵢ) = h × dₖ = d_(model).
- Unique Decomposition: Every vector v ∈ V can be uniquely written as:
v = h₁ + h₂ + ṡ + h_h, hᵢ ∈ Hᵢ
Probing experiments (e.g. Clark et al., 2019, 'What Does BERT Look At?') prove that individual attention heads specialize into dedicated geometric linguistic functions: Head 8-10 tracks direct objects, Head 7-6 tracks possessive pronouns, and Head 9-6 tracks punctuation delimiters.
The internal direct sum of vector spaces is a foundational concept of modern abstract algebra, formalized by Hermann Grassmann and Giuseppe Peano in the late 19th century.
- R⁷68 is the direct sum of 12 distinct 64D subspaces.
- Specialized heads attend to independent features without cross-talk.
How does Multi-Head Attention: Splitting Subspaces for Multiple Perspectives operate mathematically at scale?
Visual parallel pipeline:
- $$Subtitle: "Splitting embedding dimension D into h independent subspaces of size d_k = D/h."$$
- Learning goalMaster the algebraic and geometric derivation of Multi-Head Attention: Splitting Subspaces for Multiple Perspectives.
- Helpful foundationsFoundations of Build Llm From Scratch · Coordinate Invariance & Geometric Projections · Vaswani et al. (2017), Attention Is All You Need, NeurIPS
- Common trapAssuming Multi-Head Attention: Splitting Subspaces for Multiple Perspectives is an arbitrary computational trick rather than an invariant spatial law.
- Next connectionConnects upstream foundational lemmas to modern high-dimensional computational systems.
More context & mechanics
A multi-subspace architectural diagram in square 1:1 format on warm cream parchment (#fcf0db) with an 8px border.
Top section: Title in navy serif: "Multi-Head Attention: Parallel Semantic Lenses". Subtitle: "Splitting embedding dimension D into h independent subspaces of size dₖ = D/h."
Middle section: Visual parallel pipeline:
Input representation X ∈ R^(L × D) splits into h parallel attention heads via learned projection matrices Wᵢ^Q, Wᵢ^K, Wᵢ^V:
- Head 1 (Syntactic Role): Tracks subject-verb agreement (colored blue).
- Head 2 (Positional Proximity): Focuses on immediately adjacent words (colored teal).
- Head 3 (Coreference): Resolves pronouns to antecedent entities (colored gold).
- Head 4 (Semantic Topic): Associates related topical nouns (colored coral).
All h outputs head₁, ..., head_h are concatenated: [head₁, ..., head_h] ∈ R^(L × D).
Multiplied by output projection matrix W^O ∈ R^(D × D) to fuse all heads.
Bottom section: Formula: "MHA(Q, K, V) = Concat(head₁, ..., head_h) W^O. Allows simultaneous attention to information from different representation subspaces."
Pristine parallel pipeline layout, distinctive head colors, balanced square composition.
Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.
Formal derivation and peer-reviewed analysis:
"Multi-Head Attention: Splitting Subspaces for Multiple Perspectives: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."