Unit 2.2 · about 5 h
298 / 324Scaled dot-product attention from scratch
Attention is a weighted read from permitted keys, not a framework black box. Correct shapes and masks make the decoder's causal information boundary testable.
- derive
- implement
You can now
- Derive \(\operatorname{softmax}(QK^\top/\sqrt{d_k})V\) with explicit projection, score and output shapes.
- Implement stable softmax and combine causal and padding masks before normalization.
- Diagnose scale, mask-polarity and fully masked-row failures using a small explicit reference.
Covers
- Query, key and value projections with explicit shapes
- Scaled dot-product attention and the scaling term
- Stable softmax and causal mask construction
- Padding masks versus causal masks
Study, in this order
4 steps
- Attention scores and causal maskingModern LLM Notebook - decoder construction path
Sections 2, Scaled Dot-Product Attention, and 3, Causal Masking
- Trainable \(Q\), \(K\) and \(V\) projectionsLLMs from Scratch companion - attention and GPT construction
Section 3.4.1, Computing the attention weights step by step
- Stable softmax and masked attention contractStanford CS336 Assignment 1 - model contract
Section 3.4.4, Scaled Dot-Product Attention, including problems (
softmax) and (scaled_dot_product_attention), printed pp. 23 to 25 - Mask shape and polarity referencePyTorch scaled dot product attention reference
Equivalent Python implementation, boolean-mask semantics note, and Parameters:
attn_maskandis_causalDepth · Read this after your explicit oracle works. Use the contract for padding-mask broadcasting; skip kernel selection.
Focus on
- For \(X=[B, T, D]\), \(Q\) and \(K\) are \([B, T, d_k]\), \(V\) is \([B, T, d_v]\), scores are \([B, T, T]\), and the output is \([B, T, d_v]\).
- With independent zero-mean unit-variance \(Q\) and \(K\) components, their dot product has variance \(d_k\). Dividing by \(\sqrt{d_k}\) keeps the variance near one.
- Subtract each score row's maximum before exponentiating. Mask forbidden entries before softmax and normalize over keys.
- An allow mask combines \(j\le i\) with
valid_key[b,j]. Broadcast key validity as \([B, 1, T]\); exclude padded query outputs separately and define a policy for rows with no valid key.
Reproduce closed-book
Traps
- Using \(\sqrt{D}\) after splitting into heads. The scale uses the query/key width of one head.
- Copying a boolean mask between APIs without checking polarity. CS336 and SDPA use
Truefor allowed entries;MultiheadAttentionpadding masks useTruefor blocked entries. - Assuming max subtraction fixes an all-negative-infinity row. It gives undefined subtraction; reject that row or explicitly return zero for an ignored padded query.
Practice
Problem (softmax), section 3.4.4; implement run_softmax in tests/adapters.py and pass test_softmax_matches_pytorch
CS336 Assignment 1: `scaled_dot_product_attention`
Problem (scaled_dot_product_attention), section 3.4.4; implement run_scaled_dot_product_attention and pass test_scaled_dot_product_attention and test_4d_scaled_dot_product_attention
Work in your own repository and respect the course AI policy. Add causal-plus-padding and fully masked-row tests; the public snapshots do not cover every mask edge case.
Self-test
-
If all allowed logits gain 1000, should the attention output change?
No. Softmax is invariant to a common additive shift; max subtraction avoids overflow. -
Why does replacing a forbidden score with zero fail?
\(\exp(0)\) is one, so the forbidden key receives probability. Use negative infinity before softmax. -
Does a causal mask remove padding keys?
No. It only forbids future positions. A padding key earlier than the query is still causally permitted and needs a separate validity mask. -
With \(Q=[2, 3, 4]\), \(K=[2, 5, 4]\) and \(V=[2, 5, 6]\), what are scores and output?
Scores are \([2, 3, 5]\), softmax runs over the five keys, and output is \([2, 3, 6]\).
Completion criterion
Stop when
The three named public tests pass, your explicit row matches the batched result, and changing a forbidden key or value cannot affect a valid query's output.
Mark passed, then open 2.3 Only mark this when you can do it without notes.Leads to
- 2.3 Multi-head attention: split, merge and gradients Batch this operation across independent heads.
Comes after
- 1.5 From documents to training examples: packing and loss masksEnforce causal and document visibility in attention.
- 2.1 Embeddings, positions and the vocabulary axisProject the token vectors into queries, keys and values.
Resources for this unit
- Modern LLM Notebook - decoder construction path 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 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 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 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 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 Read
model.py:MLP,Block,GPT.__init__'s special scaled initialization ofc_proj.weight, andGPT._init_weights. Contrast the depth-dependent residual initialization with CS336's component initialization; do not adopt the training or generation code.