NLP Foundations
RoPE Scaling: NTK-aware and YaRN
How a 4k-token base checkpoint becomes a 128k-token release without pretraining from scratch, by rescaling RoPE's rotation frequencies rather than its raw positions.
advanced · 10 min read · Premium
A model pretrained with 4096-token sequences, deployed naively at 32000 tokens, falls apart within a few hundred tokens past its training length: RoPE's fast-rotating dimension pairs alias into angles the network never trained on, and quality collapses. Nearly every long-context release since 2023, from open-weight 32k and 128k variants to commercial context-window expansions, gets there not by pretraining at the target length but by rescaling RoPE's rotation frequencies on a short base checkpoint and lightly fine-tuning. The specific choice of rescaling matters enormously.
Position interpolation: the blunt fix
The simplest approach, position interpolation, linearly rescales position indices before computing rotation angles. If the model trained at L_train = 4096 and you need L_target = 32768, every position is divided by L_target / L_train = 8 before the rotation angle is computed, so the entire target sequence maps back into the 0..4096 range of angles the model actually trained on. It works, and needs comparatively little fine-tuning, on the order of a thousand steps, to adapt.
The problem is that it rescales every frequency band uniformly. The slow, long-range dimension pairs benefit from being compressed back into familiar territory, but so do the fast, local dimension pairs, whose fine-grained resolution for nearby tokens gets blurred even though those tokens didn't need rescaling at all. Local syntax and short-range dependencies pay a tax that only long-range structure needed to pay.
NTK-aware scaling: treat frequency bands differently
NTK-aware scaling, first circulated informally in the open-source community in 2023 and later formalised and analysed in the YaRN paper, takes the opposite approach at the level of what it changes: instead of rescaling positions, it rescales the RoPE base frequency itself, using roughly base_new = base_old * (L_target / L_train)^(d / (d - 2)). This stretches the slow, long-range dimension pairs far more aggressively than the fast, local ones, which are left nearly untouched. The name draws an analogy to neural tangent kernel arguments about networks having a harder time learning high-frequency functions from low-frequency-biased training data; the practical upshot is that local resolution survives far better than under naive interpolation, and some variants ("dynamic NTK") produce usable results with no fine-tuning at all.
YaRN: banded interpolation plus an attention-temperature fix
YaRN (Yet another RoPE extensioN) refines this further by treating the d/2 frequency pairs as three explicit bands, based on how their wavelength compares to the target context length:
fast dims (wavelength << target length): left mostly untouched, extrapolate on their own
slow dims (wavelength >> target length): interpolated, same idea as position interpolation
mid dims: a ramp function blends the two regimes, avoiding a hard boundary between them
YaRN also adds a correction unrelated to frequency: it scales the pre-softmax attention logits by a length-dependent constant, an "attention temperature" adjustment, to counteract the fact that longer contexts spread softmax mass over more keys and naturally flatten attention distributions relative to what the model saw in training. Without this, even a perfectly rescaled positional signal can leave attention entropy higher than the model was trained to expect.
Keep reading with Pro.
You're reading the preview. Unlock the full concept plus the library, study plans, the AI mentor, and daily emails.