Unit 7.3 · about 4 h
216 / 324Self-attention, cross-attention and multi-head shapes
Heads are parallel representation subspaces, and cross-attention is the bridge between sequences. Shape errors here run without error and fail silently as wrong mixing.
- implement
You can now
- Distinguish self-attention from cross-attention by which tensors supply queries, keys and values.
- Implement head splitting, per-head attention, concatenation and output projection with exact shapes.
- Explain why fixing total width while varying head count holds parameters and compute flat.
Covers
- Self-attention versus cross-attention
- Multi-head shapes, concatenation and output projection
Study, in this order
6 steps
- Multi-head self-attentionUnderstanding Deep Learning - Transformers
Section 12.3.3 (pp. 214 to 215): per-head equations 12.10 and 12.11, \(D/H\) sizing, and concatenation with output projection in equation 12.12
- Cross-attention for encoder-decoder modelsUnderstanding Deep Learning - Transformers
Section 12.8 with Figure 12.14 (pp. 226 to 227): queries from decoder embeddings, keys and values from encoder embeddings
- Notebook 12.2: multi-head self-attentionUnderstanding Deep Learning - Transformers
The complete notebook, referenced in the margin of the multi-head derivation
- Multi-head attention implementationD2L - attention and transformer implementation route
Section 11.5: the head equations with \(p\) set to
num_hiddensovernum_heads,transpose_qkvandtranspose_output, thevalid_lensrepeat per head, and thenum_hiddens100 withnum_heads5 shape check - Self-attention call and decoder cross-attentionD2L - attention and transformer implementation route
Section 11.6 Self-Attention: the
attention(X, X, X, valid_lens)toy with batch 2, 4 queries and width 100Depth · Stop before the CNN, RNN and self-attention comparison and before Positional Encoding; both belong to later units.
- Decoder self versus encoder-decoder attentionD2L - attention and transformer implementation route
TransformerDecoderBlock:attention1over(X, key_values, key_values)withdec_valid_lensversusattention2over(Y, enc_outputs, enc_outputs)withenc_valid_lens
Focus on
- D2L sets per-head width to
num_hiddensovernum_heads, so all four projections staynum_hiddensbynum_hiddens. The head count changes mixing, not size. - Heads fold into the batch axis:
(batch, steps, heads, width)becomes(batch times heads, steps, width)for the shared attention call, then unfolds back.valid_lensmust repeat once per head to stay aligned. - Cross-attention reuses the identical math with different sources: queries from the decoder state, keys and values from encoder outputs, with padding-only lengths.
- UDL Problem 12.5 is the cost intuition: narrow heads at fixed total width keep the implementation efficient.
Reproduce closed-book
(2, 4, 100), folded per-head tensors are (10, 4, 20), concatenation restores (2, 4, 100). State BERT per-head width 64 from 1024 over 16 heads.Traps
- Folding heads with one
reshape. Going straight from(batch, steps, hidden)to(batch times heads, steps, hidden/heads)runs without error but gives each head a mix of positions. Split hidden into(heads, head width), move the head axis next to batch, then merge, astranspose_qkvdoes. - Repeating
valid_lensin the wrong order. D2L folds batch-major, so lengths needrepeat_interleave([l0, l0, l1, l1]). Tiling withrepeat([l0, l1, l0, l1]) has the right length and silently masks heads with another example's length. Dropping therepeatfails loudly inmasked_softmax. - Assuming cross-attention needs a new formula. Only the sources change; calling self-attention with (Y, encoder outputs, encoder outputs) is the whole difference.
Practice
TrenTorch: Multi-Head Attention, splitting into heads
Implement split_heads and multi_head_attention_per_head; the tests check consecutive \(d_k\) chunks per head, a match with a per-head loop, and masking
TrenTorch keeps heads on their own axis, (batch, heads, seq, \(d_k\)); D2L folds them into the batch, (batch x heads, seq, \(d_k\)). Walk the reproduce shapes through both layouts.
D2L Multi-Head Attention: Exercise 1 (visualize heads)
Exercises in Section 11.5: visualize attention weights of multiple heads in the toy experiment
Plot one head per panel on the same input and confirm different heads attend differently; identical panels expose a split bug.
Self-test
-
BERT uses width 1024 with 16 heads. What is the per-head width, and what shape are its per-head query matrices?
64. Each per-head query, key and value matrix is \(64 \times 1024\). -
Why does fixing total width while varying head count hold parameters and compute flat?
All four projections staynum_hiddensbynum_hiddensregardless of head count, so QKVO parameters stay 4 times width squared and per-head matmuls sum back to the same total. -
In the D2L decoder block, which tensors supply queries, keys and values to the second attention call, and which lengths mask it?
Queries come from the decoder stateYafter self-attention; keys and values come from encoder outputs. Only encoder padding lengths mask it; there is no causal mask on this call. -
The
valid_lensrepeatper head is deleted. What breaks?The folded batch has batch times heads rows but only batch lengths, so D2Lmasked_softmaxfails with a mask shape error. The silent form of the bug is tiling withrepeatinstead ofrepeat_interleave: the count matches, but heads use another example's length.
Completion criterion
Stop when
Your multi-head block returns (2, 4, 100) on the D2L toy, per-head weights match a manual single-head call, and masking with valid lengths [3, 2] zeroes the same key slots in every head.
Leads to
- 7.5 Transformer blocks: residual stream, normalization and feed-forward Place multi-head attention inside a full residual block.
- 7.6 Parameter, FLOP and memory accounting for attention Count the projections and scores you just shaped.
Comes after
- 4.7 Decoding, variable-length batches and leak testsCross-attention replaces the fixed context.
- 7.1 Queries, keys, values and scaled dot-product attentionReuse single-head attention inside per-head projections.
- 7.2 Causal masks, padding masks and fully masked rowsCarry the same mask through folded head batches.
Resources for this unit
- Understanding Deep Learning - Transformers Use Chapter 12.1-12.11 and the matching notebooks for dot-product self-attention, extension to transformers, transformer layers and vision transformers.
- D2L - attention and transformer implementation route Use queries-keys-values, attention-scoring-functions, multihead-attention, self-attention-and-positional-encoding, transformer and vision-transformer. Implement masked single-head attention, multi-head attention and one pre-normalized block before comparing with PyTorch.
Also in this section
- PyTorch - scaled dot product attention Use the official function documentation as a differential oracle for shapes, masks, dropout and backend behavior after the manual version passes intervention tests.
- Deep-ML candidate - attention problems Use 490 Scaled Dot-Product Attention as the fifth platform admission sample. If the judge catches a seeded masking error, use at most one of 490, 491, 53, 904 or 905 as a changed-prompt drill.
- TrenTorch - attention and block questions Solve the multi-head split, sinusoidal positional encoding, causal mask and pre-norm versus post-norm questions the units name, then carry the same checks into your own block.