Rotary Position Embedding Architecture in Depth

Attention has no sense of order

Self-attention computes softmax(QK^T / sqrt(d_k)) · V from X: [N, d]. Permute the rows of X by a permutation matrix P and every term moves with it:

Attn(PX) = softmax( P (QK^T) P^T ) P V = P · Attn(X)

The output is the same vectors, shuffled, and the position-wise feed-forward block does not help: without an explicit signal a transformer is a bag of tokens. There are exactly three places to break that symmetry — the residual stream (add to the input embeddings), the N × N score matrix (add to the logits), or the Q and K vectors themselves. Every scheme below picks one, and the choice — not the formula — determines the cost.

Two properties then separate them. Translation invariance: the logit contribution should depend on m - n alone, so a token pair scores the same at positions 5 and 7 as at 5005 and 5007 — absolute schemes approximate this from data, relative schemes get it by construction. Extrapolation: past L_train the scheme feeds the network numbers — indices, biases, angles — from a range the weights were never fit on. Every trick below keeps those numbers inside the range the model knows.

Absolute schemes: learned tables and fixed sinusoids

Learned absolute (BERT, GPT-2) allocates a table P: [L_max, d] and adds P[m] at the input: one gather and one add, once — the cheapest thing on this list. It cannot extrapolate at all, since there is no row P[L_max]; invariance must be learned per position pair; and the tail of the table is undertrained because long documents are rare.

Sinusoidal replaces the table with fixed sin/cos(m / 10000^(2i/d)). Its geometric ladder of wavelengths is the good idea, and it survives into rotary. The pitch was extrapolation, since the functions are defined for every m, but the encoding competes with semantic content in the residual stream and the model learns to decode the phase patterns it actually saw — novel phases past L_train are the same out-of-distribution problem.

Score-matrix schemes: T5 relative bias and ALiBi

T5 moved position out of the vectors entirely: logit(m,n) = q_m·k_n / sqrt(d_k) + b_h[bucket(m - n)], with roughly 32 log-spaced buckets per head. Invariance is exact and free, the log spacing encodes a real prior (fine nearby, coarse far away), and distances past the largest bucket collapse into one — soft, non-crashing extrapolation. ALiBi keeps the placement and drops the table: − s_h · (m - n), slopes in a geometric series (1/2, 1/4, … 1/256) that give each head its own window, at zero parameters.

Both pay for the placement: the bias lives on the N × N score matrix, so O(h · N^2) gathers and adds per layer, plus a tensor your fused kernel must know about. ALiBi’s penalty also grows without bound, imposing a permanent recency prior — it will not collapse on long inputs but is structurally poor at retrieving from the far past, which is what people buy long context for.

Rotary: position in Q and K, not in the score

Rotary picks the third placement: leave the residual stream and the score matrix alone, and rotate each consecutive pair of dimensions in q_m and k_n by mθ_i, reusing the sinusoidal frequency ladder via θ_i = base^(-2i/d). The result — derived in the companion article, taken as given here — is that <R_m q, R_n k> depends only on m - n.

That is why rotary won: T5’s exact relative property at the cheapest cost class, O(N · d) element-wise multiply-adds on Q and K per layer, with no N × N tensor anywhere and no parameters. You rotate before the attention kernel, so FlashAttention-style kernels work unmodified. And unlike ALiBi it imposes no fixed decay; each head learns its own distance profile from the frequency mixture.

The placement table, and the KV-cache consequence

Placement has one production consequence worth spelling out: the KV cache. Under a score-matrix scheme cached keys are position-free, so entries can be evicted or renumbered freely. Under rotary the cached key is already rotated — position is baked in, so dropping tokens from the middle leaves a hole in the position sequence, and eviction schemes must re-assign contiguous positions to the survivors.

Why naive extrapolation breaks rotary

Rotary’s angles are unbounded, so nothing crashes past L_train — the output simply becomes noise. The cause is a frequency-band asymmetry. Take d_head = 128, base = 10000, so θ_i = 10000^(-i/64):

band i=0 : θ=1.0       λ=6.3 tokens
band i=63: θ=1.155e-4  λ=54,400 tokens

trained at L=4096:  band 0 sweeps ~650 full turns  (well sampled)
                    band 63 sweeps 0.47 rad = 27° (barely sampled)
run at    L=32768:  band 63 sweeps 3.79 rad = 217°
                    → 88% of that arc was NEVER SEEN in training

The high-frequency bands are fine: they wrapped hundreds of times in training, so every angle is in-distribution. The low-frequency bands never completed a rotation, so the network saw a narrow arc of their values and extrapolation hands it arcs it cannot read. Every extension method is a rule for which bands to squeeze and which to leave alone.

Position Interpolation, then NTK-aware base scaling

Position Interpolation is the blunt fix: divide m by s = L_new / L_train before rotating. With s = 8, position 32767 becomes 4095.9 and every angle lands back inside the trained range. But it is uniform: the well-sampled fast bands are squeezed by the same 8, so adjacent tokens differ by only θ_i / 8 and short-range resolution degrades — a perplexity hit at short lengths that a brief fine-tune mostly repairs.

NTK-aware scaling stretches only the out-of-distribution bands, by scaling the base rather than the positions:

base’ = base · s^( d / (d - 2) )        d=128, s=8 → base’ ≈ 82,700

θ’_i = base’^(-2i/d),  so at the two ends:
  i = 0       : θ’ = 1      unchanged  — local detail preserved
  i = d/2 - 1 : θ’ = θ/s    fully interpolated — same as PI

Fast bands keep resolution, slow bands get the full PI treatment, the middle interpolates. This is why bumping rope_theta — the most common long-context tweak in the wild — buys real extension with no fine-tuning at modest s. Quality still drops as s grows, and the middle bands follow a closed form rather than evidence.

YaRN: a per-band ramp plus attention temperature

YaRN makes the band decision explicit. For each band it counts the full rotations completed in the original context, r_i = L_train / λ_i, and interpolates by ramp:

r_i > β (many turns seen, e.g. 32)  → no interpolation
r_i < α (under one turn,   e.g. 1)   → full interpolation (/s)
in between                            → linear ramp

plus attention temperature:  logits × 1/t,  sqrt(1/t) = 0.1·ln(s) + 1

The ramp is NTK-aware scaling with cutoffs taken from measurable band statistics rather than a closed form. The temperature fixes a separate problem: with far more keys in the softmax, entropy rises and the model attends to everything a little. Sharpening the logits counteracts that, and being a constant it folds into the precomputed tables for free. YaRN reaches a given context with roughly an order of magnitude less fine-tuning data than PI.

Effective context is not nominal context

A config claiming max_position_embeddings: 131072 asserts only that the model will not error out at 131k tokens. It says nothing about whether the model uses token 100,000.

Three measurements, weakest first. Perplexity versus position: loss on the nth token should keep falling as n grows, and a flat curve marks where context stopped paying — but perplexity is dominated by local prediction, so a model can look healthy while retrieving nothing. Needle-in-a-haystack plants a fact at a known depth and asks for it back; it exposes real retrieval failure but saturates easily. Multi-task long-context suites — several needles, variable tracking, aggregation over the whole input — are the honest test, and they routinely put effective context two to four times below the advertised number. Extend, then measure.

What long context actually costs on a CPU

On CPU the scheme is never the bottleneck; the arithmetic behind the context is. KV cache memory is 2 · n_layers · n_kv_heads · d_head · bytes per token. A 3B model with 28 layers, 8 grouped-query KV heads, d_head = 128, fp16: 2×28×8×128×2 = 114,688 bytes ≈ 112 KiB/token, so 32k tokens is 3.5 GiB, larger than the 4-bit weights. Decode rereads all of it every step: at 20 GB/s usable bandwidth that is ~190 ms per token, capping you near 5 tok/s on cache traffic alone. Prefill is worse: 2 · N · P = 2×32768×3e9 ≈ 2.0e14 FLOPs, about an hour at a realistic 50 GFLOP/s.

Extension makes long context possible on CPU, not usable; chunked retrieval, int8/int4 KV quantisation and sliding-window eviction are what make it practical. Two pitfalls: a checkpoint tuned with rope_theta = 500000 served by a runtime assuming 10000 produces fluent, confidently wrong text rather than an error, so check the serving config; and any change to the scaling factor invalidates every cached key, which was rotated under the old schedule.