Gen AI

Long-Context Extension

Long-context extension rescales a model's positional signal so sequences longer than its trained length still land back in the range it actually learned.

A position beyond the trained context length gets its RoPE angle rescaled back down into the range the model actually trained on, rather than extrapolated past it

Why Does This Exist?

Every model ships with a context window fixed at training time, and rotary position embeddings are, in most current models, the mechanism computing the angles behind that limit. RoPE gives attention a clean relative-distance signal within trained range, but it doesn't erase the underlying wall: a model trained with positions up to 4,096 has only ever seen rotation angles up to whatever $4{,}096\theta$ works out to. Ask it to process position 10,000, and the rotation formula still produces a valid number, but it's an angle the model's attention patterns were never shaped around. Quality degrades past the trained length, the exact same failure additive positional encoding had, just showing up at a different layer of the mechanism.

Retraining from scratch at a longer length fixes it, but it's expensive, and often isn't necessary. Long-context extension is a family of techniques — position interpolation, NTK-aware scaling, YaRN — that instead take a model already trained at length $L$ and get it working at some longer length $L'$ , by changing how positions map to RoPE angles rather than retraining the whole model on new data from zero.

Think of It Like This

Compressing a longer ruler back onto a shorter one

Imagine a ruler marked from 0 to 100, and a model that only ever learned to read the marks in that range accurately — ask it to read mark 250 on some hypothetical extended ruler, and it's reading territory it's never calibrated against. One fix: instead of handing it the raw 250, rescale it first — squeeze the 0-to-250 range down so it fits inside 0-to-100 before reading it. Mark 250 might get mapped down to something like mark 100, mark 125 down to something like mark 50, and so on. The model reads marks it already knows how to read; the rescaling did the work of fitting a longer range into familiar territory.

That's what position interpolation does to RoPE angles. Instead of computing the angle for position 10,000 directly and feeding the model something it never trained on, it rescales position 10,000 down into the 0-to-4,096 range the model actually trained on, computes the angle for that rescaled value, and only then runs attention. The model is still reading angles from its comfort zone — the input position was compressed to get it there.

How It Actually Works

Position interpolation: linear rescaling

The simplest approach, position interpolation, scales every position down by a fixed ratio before computing its RoPE angle. Extending a model from trained length $L$ to target length $L'$ means every position $p$ gets replaced with $p \cdot (L / L')$ before the rotation angle is computed — position $L'$ maps down to exactly $L$ , position $L'/2$ maps down to $L/2$ , and the whole extended range folds evenly back into the trained range. The diagram above shows this: a position past the trained boundary gets pulled back inside it, rather than extrapolated past it.

NTK-aware scaling: not every frequency needs the same treatment

Plain linear interpolation has a cost: it compresses every frequency pair by the same ratio, including the highest-frequency pairs that were already distinguishing very nearby positions from each other well within trained range. Squeezing those down too can blur fine-grained relative distances the model relied on for local structure. NTK-aware scaling instead scales different frequency pairs by different amounts — leaving high frequencies closer to untouched, since they're responsible for fine nearby distinctions, and stretching the scaling further for low frequencies, which is where the actual long-range extrapolation problem lives.

YaRN: combining both, with a temperature adjustment

YaRN builds on the NTK-aware idea, using a smooth ramp between "leave this frequency mostly alone" and "rescale this frequency substantially," rather than picking one scaling rule for every frequency. It also adjusts the softmax temperature slightly to counteract a side effect of the rescaling, since spreading attention over a longer effective range tends to flatten the score distribution somewhat. Because it targets the specific frequencies that actually need adjusting rather than rescaling everything uniformly, YaRN typically extends usable context length further, and with less fine-tuning, than plain position interpolation.

Show Me the Code

Linear position interpolation for RoPE angles, showing how a position past the trained length gets rescaled back inside it before the angle is computed.

1import numpy as np2
3
4def rescaled_angle(position: int, theta: float, trained_len: int, target_len: int) -> float:5    scale = trained_len / target_len              # < 1 when extending beyond trained_len6    return theta * (position * scale)7
8
9theta = 0.0110trained_len, target_len = 4096, 8192               # doubling the usable length11
12angle_at_boundary = rescaled_angle(trained_len, theta, trained_len, target_len)13print(round(angle_at_boundary, 4))  # -> 20.48 -- position 4096 scaled by 0.5 first14
15angle_past_boundary = rescaled_angle(target_len, theta, trained_len, target_len)16print(round(angle_past_boundary, 4))  # -> 40.96 -- position 8192 maps back to trained_len=4096 pre-scale17print(angle_past_boundary == theta * trained_len)  # -> True -- lands exactly on the trained boundary's angle

Position 8192, the very edge of the extended window, produces the identical angle the model saw at its own trained boundary of 4,096 — exactly the folding-back the analogy describes.

Watch Out For

Assuming extension is free quality

Every extension technique trades something: linear position interpolation compresses fine-grained nearby distinctions along with the distant ones it was meant to fix, and even NTK-aware and YaRN approaches typically need at least some fine-tuning at the target length to perform well, rather than working perfectly zero-shot. "The model now accepts more tokens" and "the model uses those extra tokens as effectively as its original range" are different claims — the same distinction context windows draws between a hard limit and uniform quality within it.

Applying a RoPE-based extension technique to an ALiBi model

Position interpolation, NTK-aware scaling, and YaRN all work by rescaling RoPE's rotation angles specifically — they have nothing to rescale on a model using ALiBi, which never computes a position-dependent angle in the first place. ALiBi's own length generalization comes from its linear penalty staying well-defined at any distance, not from a rescaling step; conflating the two families means applying a fix aimed at the wrong mechanism.

The Quick Version

Long-context extension lets a model trained at length $L$ handle sequences longer than $L$ , without retraining from scratch.
Position interpolation rescales positions linearly, folding the extended range back inside the trained range before RoPE's angle is computed.
NTK-aware scaling rescales frequencies unevenly, protecting high-frequency pairs that handle fine local distinctions.
YaRN combines uneven frequency scaling with a softmax temperature adjustment, and generally extends further with less fine-tuning.
All of these techniques rescale RoPE angles specifically; they don't apply to ALiBi-based models, which generalize by a different mechanism.