Unit 2.5 · about 4 h
301 / 324RoPE: positions as rotations
RoPE makes position differences enter attention through rotations of \(Q\) and \(K\). An implementation that silently restarts positions at zero can pass a short demo and fail incremental decoding.
- derive
- implement
You can now
- Derive \((R_i q)^\top(R_j k)=q^\top R_{j-i} k\) for paired rotations.
- Implement RoPE on \(Q\) and \(K\) with explicit position IDs, head-local pairs and preserved output dtype.
- Detect pairing and offset errors using norm, common-offset and nonzero-position tests.
Covers
- Rotary Position Embedding as paired rotations
- RoPE offsets and relative-position identity
Study, in this order
2 steps
- Paired rotations and relative dot productsModern LLM Notebook - decoder construction path
Rotary Position Encoding (RoPE): 2D rotation mathematical intuition, From 2D to higher dimensions, and Full RoPE implementation
Depth · Stop before Experimental comparison. The rotation identity, not a plotted attention pattern, is the result to reproduce.
- Explicit position IDs and rotary attention contractStanford CS336 Assignment 1 - model contract
Section 3.4.3, Relative Positional Embeddings and problem (rope), printed pp. 22 to 23; section 3.4.5, Applying RoPE, printed p. 25
Focus on
- For zero-based pair \(r\), frequency is \(\theta^{-2r/d}\). Rotate \([x_{2r}, x_{2r+1}]\) by
position*frequencyusing radians, not degrees. - \(R_i^\top R_j=R_{j-i}\). The positional factor depends on the difference; content vectors \(q\) and \(k\) still matter.
- Apply the rotation after projection and head splitting to \(Q\) and \(K\), not \(V\). Require an even rotary width and agree on adjacent-pair versus split-half layout.
- Index \(\sin/\cos\) with the supplied positions, including nonzero offsets. Compute angles and rotations in
float32when needed and cast the result back; buffers are not trainable weights.
Reproduce closed-book
Traps
- Mixing adjacent pairs with a split-half reference without permuting channels and projection weights. Both conventions can work, but their raw tensors are not interchangeable.
- Using
arange(T)instead of supplied position IDs. This discards offsets and per-sequence positions. - Inferring monotonic distance decay or reliable long-context quality from the identity. Rotations are periodic, and the identity alone guarantees neither.
Practice
CS336 Assignment 1: `rope` and rotary `multihead_self_attention`
Problem (rope), section 3.4.3, and Applying RoPE, section 3.4.5; implement run_rope and run_multihead_self_attention_with_rope and pass test_rope and test_multihead_self_attention_with_rope
Add your own position-zero, nonzero-offset, common-offset dot-product, norm-preservation and odd-width rejection tests. Test float32 and a supported low-precision dtype. For positions beyond a precomputed table, test the documented extension or rejection policy.
Self-test
-
Why do a common offset \(c\) on query and key positions cancel?
\(R_{i+c}^\top R_{j+c}=R_{j-i}\), assuming unchanged \(q\), \(k\) and frequencies. The rotary positional factor is unchanged. -
With \(\theta=10000\) and \(d=4\), what frequencies do the two adjacent pairs use?
1 and 0.01 radians per position. Position \(p\) gives angles \(p\) and \(p/100\). -
Should rotating a vector change its norm?
Not in exact arithmetic. Each 2D rotation is orthogonal; numerical tests allow a dtype-appropriate tolerance. -
For a full paired rotation, what should happen with an odd head width?
Reject it unless a separate partial-rotation contract explicitly leaves a channel unrotated. The assignment's full paired construction needs an even width.
Completion criterion
Stop when
Both named public tests and your offset and pairing tests pass, and you derive the relative-position identity without looking at the notebook.
Mark passed, then open 2.6 Only mark this when you can do it without notes.Leads to
- 2.6 The modern decoder block: residuals, RMSNorm and SwiGLU Insert rotary attention into the pre-norm residual block.
- 5.2 Building a KV cache Preserve absolute positions when appending rotated keys.
Comes after
- 2.1 Embeddings, positions and the vocabulary axisPut relative-position structure into attention scores.
- 2.3 Multi-head attention: split, merge and gradientsRotate \(Q\) and \(K\) within each head.
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.
- 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.
Also in this section
- 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.
- 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.
- 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.