How does a Transformer know which word came first? We derive Rotary Position Embeddings (RoPE) from scratch, work a 2D example with real numbers, and verify the relative-position property in PyTorch.
Intro
The Order Problem
The dog bit the man. The man bit the dog. Same words, completely different story. Here is the surprising part: the attention mechanism you learned about, left alone, cannot tell these two sentences apart. Today we fix that, and we derive the trick almost every modern open model uses: rotary position embeddings, or RoPE.
- Same words, different meaning
- Attention sees a set
What We Will Cover
In the last episode we covered sampling: how a model picks the next token from a probability distribution. Now we look at what the model can actually see when it does that. We will derive RoPE, work a numerical example, implement it in PyTorch, then see why long context is hard and how models stretch it.
- Derive RoPE
- Why long context is hard
- Practical limits
Why Attention Needs Positions
Attention Is Order-Blind
Recall from the ML course's attention lesson: each token makes a query, a key and a value, scores are query dot key, and the output is a weighted sum of values. Nothing in that recipe mentions where a token sits. If you shuffle the input tokens, the outputs shuffle the same way. The model sees a bag of tokens, not a sequence.
- Scores use only content
- Permute tokens, permute outputs
Injecting Position
So we must inject position information. There are two big ideas. Absolute schemes tell each token its own index, like seat number seven. Relative schemes tell attention how far apart two tokens are, like three seats to my left. Language mostly cares about the second one: what matters is that an adjective is right before its noun, not that it is at position four hundred.
- Make vectors position-aware
- Absolute vs relative
Older Schemes
Sinusoidal and Learned
Briefly recalling the attention lesson: the original Transformer added fixed sine and cosine waves of different frequencies to the token embeddings. GPT-2 style models instead learned one vector per position. Both add position to the input once, at the bottom of the network, and hope that attention can decode it later.
- Sinusoidal: added to embeddings
- Learned: one vector per index
Their Weaknesses
Both have problems. A learned table simply has no row for position five thousand if you trained to four thousand. And with added positions, the relative offset is only implicit, tangled with content inside the dot product. We want a scheme where the attention score itself depends on the offset, by construction.
- Learned table stops at max length
- Relative offset is indirect
Deriving RoPE
The Goal
Here is the goal stated precisely. We want a function that takes a query vector q at position m, and a key vector k at position n, and transforms them so that the new dot product depends on q, k, and only the difference m minus n. No absolute positions allowed to leak through.
- Find f(q, m) and f(k, n)
- Dot product depends on m minus n
Rotation Preserves Dot Products
Think geometrically. A dot product is length times length times the cosine of the angle between vectors. If you rotate two vectors by the same angle, the angle between them does not change, so the dot product is unchanged. Now the clever twist: rotate them by different angles, and the dot product changes by exactly the difference in angles.
- Rotate both vectors equally
- Angle between them unchanged
The 2D Rotation Matrix
In two dimensions, rotation by angle phi is the matrix with cosine phi and minus sine phi on the top row, sine phi and cosine phi on the bottom row. Two facts matter. Rotating by phi then by psi equals rotating by phi plus psi. And the transpose of a rotation is rotation by minus phi, which is its inverse.
- R(φ) = [cos φ, −sin φ; sin φ, cos φ]
- Transpose is the inverse
Rotate by Position
Now define the rule. Take a two-dimensional query q at position m and rotate it by the angle m times theta. Take a key k at position n and rotate it by n times theta. Theta is a fixed constant for now. Each position gets its own rotation, growing linearly with the position index.
- q at m rotated by mθ
- k at n rotated by nθ
The Key Algebra
Compute the new score. The rotated query transposed is q transposed times R of minus m theta. Multiply by R of n theta times k. The two rotations combine into a single rotation by n minus m, times theta. So the score equals q transposed, rotation by n minus m theta, k. Only the offset survives. That is the whole trick.
- (R_m q)ᵀ(R_n k) = qᵀ R_{n−m} k
- Depends only on n − m
Many Dimensions, Many Frequencies
Real heads have dimension d, like sixty-four or one hundred twenty-eight. So we split the vector into d over two pairs, and rotate each pair with its own frequency: theta sub i equals ten thousand to the power of minus two i over d, for i from zero up to d over two minus one. The first pair spins fast, the last pair barely moves.
- Split d dims into d/2 pairs
- θ_i = 10000^(−2i/d)
What Those Frequencies Mean
Take d equals sixty-four. Theta zero is one radian per token, so that pair completes a turn about every six tokens. Theta thirty-one is about one point three times ten to the minus four, a full turn only after roughly forty-seven thousand tokens. Fast pairs resolve nearby offsets, slow pairs encode long-range distance, like the second hand and hour hand of a clock.
- d=64: θ0 = 1
- θ31 ≈ 0.000133
- Fast pairs: local detail
Matrix Form and Efficiency
The full rotation matrix is block diagonal, with one two-by-two rotation per pair. You never build it. For each pair you compute x times cosine minus y times sine, and x times sine plus y times cosine. That is just elementwise multiplications. Also note RoPE is applied to queries and keys only, not values, and inside every layer rather than once at the input.
- Block-diagonal rotation
- Apply with cos and sin elementwise
Worked Example
Setup With Real Numbers
Let us test this with numbers. Use one pair, with theta equal to thirty degrees. Let the query be two, zero, and the key be one, one. Without any rotation, the dot product is two times one plus zero times one, which is two. Now we place them at positions and watch what changes.
- q = (2, 0), k = (1, 1)
- θ = 30 degrees
Positions 3 and 1
First case: query at position three, key at position one. The query rotates by three times thirty, which is ninety degrees, giving zero, two. The key rotates by thirty degrees: cosine thirty minus sine thirty is about point three six six, sine plus cosine is about one point three six six. The dot product is two times one point three six six, which is two point seven three two.
- q at m=3: rotate 90°
- k at n=1: rotate 30°
Positions 5 and 3
Second case: shift both tokens forward by two, so query at five, key at three. The offset is still two. Now the query rotates by one hundred fifty degrees, landing at minus one point seven three two, one. The key rotates by ninety degrees, landing at minus one, one. Dot product: one point seven three two plus one, which is two point seven three two. Identical.
- q at m=5: rotate 150°
- k at n=3: rotate 90°
- Same offset of 2
Why It Matches
Why did they match? In both cases the angle between the rotated vectors is fifteen degrees: the key started at forty-five degrees, the query at zero, and the offset of two tokens adds sixty degrees to the gap, minus forty-five. The score is length two times length root two times cosine of fifteen degrees, which is 2.732. Move the pair anywhere and the score holds.
- Angle gap is 15° both times
- 2√2 × cos 15° = 2.732
Score Versus Distance
Now zoom out. For a fixed query and key, the score as a function of offset is a sum of cosines at many frequencies. One cosine would oscillate forever, but a blend of many frequencies tends to wash out as the offset grows, so on average distant tokens interact less. This gives a gentle built-in locality bias, though not a hard cutoff.
- Scores oscillate with offset
- Multi-frequency decay
PyTorch Implementation
Building the Angles
Now the code. First compute the inverse frequencies: base to the power of minus arange zero to d step two, divided by d. Then take the outer product of positions with those frequencies to get an angle for every position and every pair. Take cosine and sine of that table. This is the entire precomputation.
- inv_freq = base^(−2i/d)
- angles = positions ⊗ inv_freq
The Rotate Function
The rotate function reshapes the last dimension into pairs, calls them x one and x two, and returns x one cosine minus x two sine, and x one sine plus x two cosine, stacked back together. Be careful about layout: some libraries pair neighbouring dimensions, others pair the first half with the second half. Weights trained with one layout break silently under the other.
- Split pairs: x1, x2
- Return x1·cos − x2·sin, x1·sin + x2·cos
Verifying Relative Position
Now verify the property. Draw a random query and key of dimension sixty-four. Rotate the query at position m and the key at position n, and take the dot product. Repeat with both shifted by the same amount s, say a hundred and seventeen. Print the difference. You should see something around ten to the minus six in float32, which is just rounding noise, not a real gap.
- Random q, k
- Score(m,n) = Score(m+s, n+s)
Where RoPE Lives in the Model
In a real model, RoPE goes right after the query and key projections and before the score computation, in every layer and every head. The key at position n is rotated by its own position, once, and never needs to change. That detail matters a lot for the next episode, where we cache keys to avoid recomputing them.
- After Q and K projections
- Before attention scores
- Cached keys are stored rotated
Why Long Context Is Hard
Quadratic Cost
Positions are only half the story. Attention compares every token with every other token, so the score matrix has n squared entries per head per layer. At four thousand ninety-six tokens, that is about sixteen point eight million entries. At one hundred thirty-one thousand seventy-two tokens, it is about seventeen point two billion. Thirty-two times longer costs a thousand times more.
- n² score matrix
- 4,096² ≈ 16.8 million
- 131,072² ≈ 17.2 billion
Efficient Kernels Help
Fused kernels such as FlashAttention avoid writing the full matrix to memory, which saves memory enormously, but the arithmetic still grows quadratically. And there is a second cost: storing keys and values for every token in every layer. That memory bill is the subject of the next episode.
- Fused kernels avoid storing the matrix
- Compute still grows with n²
Extrapolation Failure
The second problem is extrapolation. A model trained on sequences up to four thousand tokens has only ever seen rotation angles up to four thousand times theta. Beyond that, the slow pairs reach angles they never saw, and attention scores become unfamiliar values. In practice quality often collapses rather than degrades gracefully. The model is not forgetting; it is in unexplored territory.
- Trained to length L
- Beyond L: unseen angles
Stretching the Window
Position Interpolation
The simplest fix is position interpolation. If you trained to length L and want length L prime, divide every position by the ratio. To go from four thousand to sixteen thousand, position eight thousand becomes two thousand. All angles stay inside the trained range. A short fine-tune then teaches the model that neighbouring tokens are now closer together on the dial.
- Squeeze positions: m → m·L/L′
- Stay inside seen angles
The Cost of Squeezing
But uniform squeezing has a cost. The fastest pairs, which distinguish adjacent tokens, now differ by only a quarter of the angle. Two neighbouring words look more similar to the model than before, which can hurt fine local detail. Slow pairs, on the other hand, benefit from squeezing. Different frequencies deserve different treatment.
- Neighbours become blurry
- High frequencies lose resolution
NTK-Aware and YaRN Ideas
That insight leads to NTK-aware scaling and YaRN, described conceptually. Instead of squeezing positions uniformly, enlarge the base so slow frequencies stretch a lot while fast ones barely change. YaRN goes further: it blends interpolation and no-change per frequency band, and adds a small temperature correction to the attention scores. The details vary between implementations, so check your model's documented configuration.
- Change the base, not positions
- Blend per frequency
- Adjust attention temperature
ALiBi in One Comparison
One comparison to finish the schemes. ALiBi skips position vectors entirely and simply subtracts a penalty from each attention score, proportional to the distance, with a different slope per head. It is simple and extrapolates well to longer sequences, but it bakes in a strong preference for nearby tokens. RoPE rotates; ALiBi penalises. RoPE became the dominant choice in many open models.
- No rotation, no added vectors
- Penalty −slope × distance on scores
Using Context Well
Advertised Versus Effective
A model advertising a very large window is telling you what fits, not what it uses well. Effective context, the length at which the model reliably retrieves and reasons over information, is often shorter than the advertised one, and depends on the task. Finding a single fact is easier than combining ten scattered facts. I will not quote numbers here, because they change with every release; test your own model.
- Window size is a limit, not a guarantee
- Quality varies with task
Lost in the Middle
Researchers have reported a lost-in-the-middle effect: when relevant information sits in the middle of a long prompt, models often use it less reliably than information at the beginning or end. It resembles how you remember the first and last items of a long list. Practical lesson: put key instructions and the question where they will be seen, and keep irrelevant text out.
- Start and end are used best
- Middle details get missed
Mistake: Perfect Memory
The first common mistake is treating the window as perfect memory. Stuffing an entire codebase or a hundred documents into the prompt does not guarantee the model uses all of it. Irrelevant text competes for attention, raises latency and cost, and can mislead. Retrieve the relevant parts and give the model a focused prompt, as we will do in the RAG lessons.
- More tokens ≠ better answers
- Noise dilutes attention
Mistake: Forgetting Output Tokens
The second mistake is forgetting that output tokens count toward the window. If the window is eight thousand and your prompt is seven thousand five hundred tokens, you have at most five hundred left for the answer, and it may be cut off mid-sentence. Also remember that in a chat, every previous turn is resent, so the history grows each turn. Budget input, history and output together.
- Window = input + output
- Truncation or cut-off answers
Practical Checklist
Here is a short checklist. Count tokens with the model's own tokenizer before sending. Reserve an explicit output budget. Trim or summarise old chat history. Put the most important material near the start or end. And run your own needle tests: hide a fact at different depths in your real documents and see where the model fails.
- Count tokens before sending
- Reserve output budget
- Test needle placement
Recap and Next Episode
Recap
Recap. Attention is order-blind, so we must inject position. RoPE rotates each query and key pair by position times a per-pair frequency, so the score depends only on the offset. Long context costs quadratic compute and risks extrapolation failure, which interpolation, NTK and YaRN-style scaling soften. And effective context is smaller than advertised, so budget your tokens.
- Attention needs positions
- RoPE: rotate q and k by mθ
- Long context: cost and extrapolation
Next: The KV Cache
Next time we ask: when a model generates token after token, why not recompute everything each step? We will meet the KV cache, see how much memory it takes with real arithmetic, and learn why decoding is memory-bound.
- Decoding repeats work
- Cache keys and values
- Memory-bound generation