Extending RoPE to Longer Contexts

The problem: context is a hard wall, not a soft limit

When we say a model has a ‘4K context window,’ we usually mean it was trained on sequences up to 4096 tokens long. Nothing in the transformer’s matrix multiplies actually forbids a longer sequence — attention is defined for any N. What breaks is the position information. The model learned to interpret positional signals only across the range it saw, and feeding it a token at position 8000 when it never trained past 4095 hands it a signal outside its learned domain.

With RoPE this shows up dramatically. Perplexity is flat and healthy up to the training length, then explodes almost vertically the instant you cross it. There is no gentle shoulder. That cliff is what ‘extending context’ has to defeat: not adding capacity in some abstract sense, but making positions beyond the training length produce positional signals the model already knows how to read. Every method in this article is a different answer to one question — how do we present a position of, say, 8000 to a model that was only ever trained on positions 0 through 4095, without lying to it in a way it can’t recover from?

RoPE in one paragraph: position as rotation

RoPE encodes position by rotating the query and key vectors rather than adding a learned vector to them. Split a head’s dimension d into d/2 pairs. For the pair at index i, RoPE rotates that 2D sub-vector by an angle θ_i(m) = m · b^(-2i/d), where m is the token’s absolute position and b is a fixed base (almost always 10000). Each pair spins at its own rate; low i spins fast, high i spins slowly.

The elegance is what happens inside attention. The score between a query at position m and a key at position n depends only on their relative offset m - n, because the two rotations combine to R(m) · R(n)^T = R(m - n). So RoPE injects absolute position on the way in but the dot product reads out relative position — exactly what language needs. That single property, score(m, n) = f(q, k, m - n), is both why RoPE works and why extending it is subtle: change how angles map to positions and you change how every relative offset is perceived.

The frequency spectrum: fast and slow dimensions

The exponent b^(-2i/d) sweeps a geometric series of frequencies. At i = 0 the factor is 1: the fastest dimension, advancing one radian per token. At i = d/2 - 1 the factor is close to 1/b = 1e-4: the slowest dimension, barely turning across thousands of tokens. In between sits a smooth ramp from high frequency to low.

It is more intuitive to talk in wavelengths. The wavelength of dimension i is λ_i = 2π / θ_i = 2π · b^(2i/d) — the number of tokens it takes that pair to complete one full revolution. For b = 10000 and d = 128, the fastest pair has λ ≈ 6.3 tokens and the slowest has λ ≈ 54,000 tokens. This spread is the whole game. High-frequency dimensions turn many times within any context and encode fine, local offsets; low-frequency dimensions may not even finish one turn across the trained window and encode coarse, long-range position. Every extension method is really a policy for which frequencies to touch and which to leave alone.

Why naive extrapolation breaks

Suppose we just run the model at position 8000 with the original angles. The slow, low-frequency dimensions are the ones that suffer. A pair with wavelength 54,000 tokens saw, during 4K training, only angles in [0, 0.47] radians — a thin sliver of the circle. At position 8000 it is asked to produce angles up to 0.95 radians: still a small number, but a region of its input space the attention weights were never fit on.

The failure is out-of-distribution, not out-of-range. The dot product that reads relative position was tuned only on the (cos, sin) values that occur for offsets up to the training length. Push the offset past that and the query/key interaction lands on angle combinations the network has no learned response to, so attention scores become erratic — large where they should be small, flat where they should be sharp. Because attention is a softmax, a few corrupted logits poison the whole distribution, and the error compounds layer over layer. That is the perplexity cliff. The fix is never to let any dimension exceed the angular range it was trained on — which is precisely what interpolation buys.

Linear Position Interpolation: squeeze positions inward

Position Interpolation (PI, Chen et al. 2023) is the blunt, beautiful idea. If the model is comfortable with positions [0, L] and we want to serve [0, L’] with L’ > L, don’t extend the angles — compress the positions. Divide every position by the scale factor s = L’ / L before computing its rotation, so the whole longer sequence is folded back into the angular range the model already knows.

Concretely, replace θ_i(m) = m · b^(-2i/d) with θ’_i(m) = (m / s) · b^(-2i/d). A model trained at 4K, extended to 16K, uses s = 4: position 16000 is presented as if it were position 4000, position 8000 as 2000, and so on. No angle ever exceeds what training covered, so the OOD cliff is gone. The cost is resolution: neighboring tokens that used to be one full unit apart in angle are now a quarter unit apart, so the model must relearn to distinguish finer positional differences. That is why PI needs a short fine-tune — roughly 1000 steps — after which it comfortably reaches the new length.

The interpolation formula, stated cleanly

It is worth writing PI as a clean transform of the position-to-angle map. Let f(x, m) be RoPE applied to vector x at position m. Position Interpolation defines the extended map as:

scale factor:   s = L' / L          (target length / trained length)

PI angle:       θ'_i(m) = (m / s) · b^(-2i/d)

equivalently:   f'(x, m) = f(x, m · L / L')

every dimension i is slowed by the SAME factor s

The defining trait of PI, and the thing later methods change, is the phrase ‘the same factor s.’ PI is uniform: it divides fast and slow dimensions alike. That uniformity is what makes it simple and robust, but it is also its weakness. The high-frequency dimensions, which were doing perfectly good local work and never approached their trained limit, get needlessly compressed too — blurring the fine positional detail the model relies on for short-range structure. NTK-aware scaling and YaRN both start from the observation that we should not treat every frequency the same.

A worked example: 4K to 16K with PI

Take b = 10000, head dimension d = 128, trained length L = 4096, target L’ = 16384, so s = 4. Look at two dimensions at the extremes of the spectrum and ask what happens to the last token, at position m = 16383.

Fast dim (i = 0):   θ_0 = 1.0 per token,  λ ≈ 6.3 tokens
  naive angle at m=16383:  16383 rad  (wraps thousands of times -- fine, in-domain)
  PI angle:                16383 / 4 = 4095.75 rad  (also fine)

Slow dim (i = 63):  θ_63 ≈ 1.15e-4,  λ ≈ 54,600 tokens
  trained range (m≤4095):   up to 4095 · 1.15e-4 = 0.471 rad
  naive angle at m=16383:      16383 · 1.15e-4 = 1.884 rad  <-- 4x past trained!
  PI angle at m=16383:        (16383/4) · 1.15e-4 = 0.471 rad  <-- back in range

The fast dimension was never in danger — it wraps around the circle constantly, so every angle it ever produces is in-distribution regardless of position. The slow dimension is the one that blew past its trained ceiling of 0.471 rad; PI pulls it exactly back to that ceiling. Notice the collateral damage, though: the fast dimension’s per-token step shrank from 1.0 to 0.25 rad, quartering its local resolution for no benefit. That waste is the opening the next methods exploit.

NTK-aware scaling: change the base, not the positions

NTK-aware scaling (from the researcher ‘bloc97’) reframes the fix. Instead of dividing every position uniformly, it changes the base b. Because the base sits inside the exponent b^(-2i/d), nudging it up affects each dimension by a different amount — a lot for the slow dimensions, almost nothing for the fast ones.

The name comes from Neural Tangent Kernel intuition: networks learn high-frequency features poorly, so it is dangerous to interpolate (blur) the high-frequency dimensions. NTK-aware therefore does the opposite of PI at the fast end — it lets high frequencies extrapolate unchanged (they were never in trouble) while making the low frequencies interpolate (they were). The new base is chosen so that the slowest dimension ends up scaled by roughly the same factor PI would have used:

b' = b · s^(d / (d - 2))       # s = L'/L

then use RoPE unchanged with base b':  θ'_i(m) = m · b'^(-2i/d)

for d=128, s=4:  b' = 10000 · 4^(128/126) = 10000 · 4.09 ≈ 40,900

The striking part: at i = 0, b^0 = 1 no matter the base, so the fastest dimension is completely untouched — full resolution preserved. At i = 63 the effective frequency drops by about 4x, matching PI on the dimension that actually needed it. NTK-aware’s headline feature is that a modest extension often works training-free, though a short fine-tune still helps and larger scale factors leak error into the middle frequencies.

Reading NTK as a per-frequency policy

It pays to see PI and NTK-aware as two points on one axis: how does the scaling vary with frequency? PI applies a flat scale of s to everything. NTK-aware applies a scale that ramps smoothly from about 1 (no change) at the fast end to about s at the slow end. Both keep the slowest dimension inside its trained angular range; they differ entirely in what they do to everything above it.

This reframing explains the empirical results. NTK-aware preserves local positional acuity because it barely touches the fast dimensions, so short-range language modeling stays crisp and the model tolerates the change even without retraining. PI, by uniformly compressing, degrades that local acuity and so needs fine-tuning to recover it — but once fine-tuned, PI is very stable and predictable at large scale factors. Neither is strictly better; they trade training-free convenience against large-scale robustness. The natural question is whether we can get both by being even more deliberate about the middle of the spectrum — which is exactly YaRN’s move.

YaRN, part 1: interpolate by parts

YaRN (Peng et al. 2023) starts from a sharper classification. For each dimension it computes how many full rotations it completes across the original context: r_i = L / λ_i. Three regimes follow. Dimensions that spin many times over the window (high frequency, large r_i) encode purely local offsets and should be left alone — extrapolated, scale 1. Dimensions that don’t even finish one rotation (low frequency, small r_i) encode absolute long-range position and should be fully interpolated — scaled by s, like PI. In between, YaRN uses a linear ramp.

r_i = L / λ_i        # rotations dim i makes over trained length L

  r_i > β  (e.g. 32):   extrapolate  -> keep original θ_i   (local dims)
  r_i < α  (e.g. 1):    interpolate  -> θ_i / s            (global dims)
  α ≤ r_i ≤ β:        linear blend between the two ends

ramp:  γ_i = (r_i - α) / (β - α),  clamped to [0, 1]
        θ'_i = (1 - γ_i)·(θ_i / s) + γ_i·θ_i

This is often called NTK-by-parts: rather than the smooth base-change ramp of NTK-aware, YaRN explicitly picks which dimensions are local, which are global, and which are blended, with thresholds you can tune. The result is that local dimensions keep full resolution, global dimensions are safely interpolated, and only the ambiguous middle band is compromised.

YaRN, part 2: attention temperature

YaRN adds a second, subtler correction that PI and NTK-aware ignore: attention temperature. When you stretch the context, the average magnitude of attention logits drifts, which changes the sharpness (entropy) of the softmax. Longer sequences also mean each query competes against more keys, flattening the distribution. YaRN counteracts this by scaling the logits by a constant 1/t before the softmax:

attention = softmax( (Q K^T) / (t · sqrt(d_k)) ) V

recommended:  sqrt(1/t) = 0.1 · ln(s) + 1      (s = scale factor)

for s = 4:   sqrt(1/t) = 0.1 · 1.386 + 1 = 1.139   ->   t ≈ 0.77

The clever engineering touch is that this temperature is just a constant multiplier on the query and key vectors, so it can be baked into the rotated Q and K at no runtime cost — no change to the attention kernel itself. Empirically the temperature term is a small but real win: it recovers a chunk of the perplexity gap on its own. YaRN’s combination — NTK-by-parts frequencies plus temperature — reaches the target length with far less fine-tuning than PI (often ~400 steps, a fraction of the data), which is why it became the default context-extension recipe for many open models.

Training-free vs fine-tuned: the real trade

The methods line up on a spectrum of how much training they demand. Training-free is seductive — ship a longer window by editing a few lines of the rotary code — and for modest extensions (2x–4x) NTK-aware and dynamic variants genuinely deliver usable quality with zero gradient steps. But training-free quality degrades as the scale factor grows and as you push into the very long tail, because no amount of clever angle-remapping teaches the model to use information 30K tokens back if it never practiced doing so.

Fine-tuned methods (PI always, YaRN with a light touch) spend a few hundred to a few thousand steps on long sequences and in return get robust, large-scale extension — 8x, 16x, and beyond — with perplexity that stays flat across the new window. The honest framing: angle rescaling removes the catastrophic failure (the OOD cliff), and that part can be training-free; fine-tuning removes the residual failure (blurred resolution, unpracticed long-range retrieval), and that part usually cannot. Choose training-free for a quick, moderate bump; budget fine-tuning for a large, dependable window.