Multi-head attention: split, merge and gradients

A reshape can preserve every tensor size while mixing tokens and heads. Forward and backward comparisons check that the batched code implements the intended operation.

implement

You can now

Implement head splitting, transposition, attention, merging and output projection without a high-level attention module.
Match outputs and input and projection-weight gradients against a per-head reference under identical weights.
Localize head-layout and scaling errors with a changed batch size, sequence length or head count.

Covers

Head split, transpose, merge and tensor layout
Multi-head attention forward and backward behavior

Study, in this order

4 steps

Multi-head data flow
Modern LLM Notebook - decoder construction path

Section 4, Multi-Head Attention, especially How Multi-Head Is Implemented in Code and Complete Data Flow

Done
Explicit head split and merge
LLMs from Scratch companion - attention and GPT construction

Section 3.6.2, Implementing multi-head attention with weight splits, including the separate-head matrix multiplication example

Done
Causal MHA contract
Stanford CS336 Assignment 1 - model contract

Section 3.4.5, Causal Multi-Head Self-Attention, printed pp. 25 to 26

Depth · Implement the no-RoPE adapter here. Leave Applying RoPE and the with_rope adapter to the RoPE unit.

Done
Independent batched attention comparison
PyTorch scaled dot product attention reference

Equivalent Python implementation, dropout warning, numerical backend note and Shape legend

Depth · Compare only after the explicit per-head oracle passes. Set dropout_p=0.0; kernel tuning is outside this unit.

Done

Focus on

With \(D=H \cdot d\), split \([B, T, D]\) to \([B, T, H, d]\), then transpose to \([B, H, T, d]\). A reshape alone does not exchange \(T\) and \(H\).
Each head normalizes its own \([T, T]\) score matrix. Concatenate head outputs before applying the output projection.
Reverse the transpose before merging. A transposed tensor may need contiguous storage or reshape rather than view.
Transfer identical weights into a per-head loop and the batched path. Backpropagate the same scalar contraction of outputs and compare gradients, not just shapes.

Reproduce closed-book

Implement causal two-head attention for \(B=2\), \(T=3\), \(D=8\). Reproduce its output and gradients with a loop over heads using sliced projection weights, then compare the attention core with SDPA at zero dropout.

Traps

Comparing separately initialized modules. Different random weights do not test implementation equivalence.
Merging \([B, H, T, d]\) directly into \([B, T, D]\). This silently scrambles head and token ordering.
Assuming eval() disables functional SDPA dropout. Its dropout_p argument must be zero explicitly.

Practice

CS336 Assignment 1: `multihead_self_attention`

Problem (multihead_self_attention), section 3.4.5; implement run_multihead_self_attention in tests/adapters.py and pass test_multihead_self_attention

The public test checks a forward snapshot, not gradients. Extend your own tests with a transferred-weight per-head oracle, input and all four projection gradients, and non-contiguous input.

Self-test

For \(B=2\), \(T=5\), \(D=12\) and \(H=3\), what shapes enter attention and leave the merge?
At fixed \(D\), does doubling \(H\) double the bias-free MHA projection parameter count?
Why compare projection-weight gradients as well as input gradients?

Completion criterion

Stop when

The named public test passes and your float64 forward and gradient comparisons agree within atol=1e-8 and rtol=1e-6 for one and multiple heads, with dropout disabled.

Mark passed, then open 2.4 Only mark this when you can do it without notes.

Leads to

2.4 MQA and GQA: sharing keys and values Change the number of KV heads without changing the query-head layout.
2.5 RoPE: positions as rotations Rotate \(Q\) and \(K\) within each head.

Comes after

2.2 Scaled dot-product attention from scratchBatch this operation across independent heads.

Resources for this unit

Modern LLM Notebook - decoder construction path
Walking Labsprimarycourse
Work through English notebooks 03-embedding, 04-position-encoding, 05-transformer-block, 06-mini-gpt and the selected normalization, RoPE, GQA and SwiGLU comparisons in 08-gpt2-to-modern-models. Re-derive shapes and rerun with changed dimensions.
LLMs from Scratch companion - attention and GPT construction
Sebastian Raschkarequiredimplementation
Use the ch03 attention notebooks and ch04 model code. Rebuild the model without a high-level Transformer module and compare one forward and backward path with an independent small reference.
Stanford CS336 Assignment 1 - model contract
Stanford Universityrequiredassessment
Use the named linear, embedding, rmsnorm, positionwise_feedforward, rope, softmax, scaled_dot_product_attention, multihead_self_attention, transformer_block, transformer_lm and transformer_accounting specifications and public tests. Keep implementation independent and respect the course AI and honor-code boundary.
PyTorch scaled dot product attention reference
PyTorchsupportreference
Use the functional contract only after the explicit attention oracle passes. Transfer identical weights and masks, compare outputs and gradients, and explain dtype, causal and dropout behavior from the documented semantics.

Also in this section

Deep-ML candidate - three transformer closing drills
Deep-MLoptionalpractice
After the platform quality sample passes, attempt at most problems 490, 53 and 904 as changed-prompt drills for attention shapes and multi-head composition. Preserve independent tests; a judge result does not replace the section checkpoint.
nanoGPT - residual projection initialization
Andrej Karpathysupportimplementation
Read model.py: MLP, Block, GPT.__init__'s special scaled initialization of c_proj.weight, and GPT._init_weights. Contrast the depth-dependent residual initialization with CS336's component initialization; do not adopt the training or generation code.