YaRN + NTK Scaling — Extending RoPE Context

RoPE, in the one form YaRN needs

RoPE splits a head of dimension d into d/2 coordinate pairs and rotates pair i at position m by angle m · θ_i, with

θ_i = b^(-2i/d),   i = 0 .. d/2 - 1,   b = 10000 (typically)
λ_i = 2π / θ_i = 2π · b^(2i/d)      # wavelength, measured in tokens

Because the attention logit between positions m and n depends only on (m - n) · θ_i, RoPE is relative: each pair is a clock, and the phase difference between two clocks is the encoded distance. The pairs differ enormously in speed. With d = 128, pair 0 has λ_0 = 2π ≈ 6.3 tokens while pair 63 has λ_63 ≈ 5.5×10^4 tokens. That four-order-of-magnitude spread is the whole reason one global scale factor is the wrong tool, and it is the quantity YaRN keys off.

Why linear interpolation hurts

Position Interpolation (PI) extends context L → L' by rescaling the position itself: m → m / s with s = L' / L. Every angle shrinks by the same factor, so no dimension sees a phase it did not see in pretraining — which is why PI is stable and cheap to fine-tune.

The damage is concentrated at the fast end. With s = 16, pair 0’s effective wavelength goes from ~6.3 tokens to ~100. Two adjacent tokens, once separated by a full radian of phase in that dimension, now differ by about 0.063 rad. The high-frequency dimensions carry exactly the local information — ‘this token comes immediately after that one’ — and PI compresses their signal into the noise floor of the dot product.

The symptom is a model that still summarizes a 64k document acceptably but gets sloppier at short range: worse perplexity on ordinary passages, muddier syntax, degraded copying. You bought reach with resolution.

NTK-aware scaling: move the base, not the position

The NTK-aware fix reframes the problem: rather than rescaling m, raise the base b so the slowest pair stretches by s while the fastest is barely touched. Solve for it by demanding the last pair (i = d/2 - 1) interpolate fully:

require  θ'_(d/2-1) = θ_(d/2-1) / s
         b'^(-(d-2)/d) = b^(-(d-2)/d) / s
         (b' / b)^((d-2)/d) = s
  ⇒      b' = b · s^( d / (d-2) )

d = 128, s = 16:  b' = 10000 · 16^(64/63) ≈ 10000 · 16.7 ≈ 1.67×10^5

This degrades short-range behavior far less than PI and works passably with no fine-tuning at all. But it is still one smooth curve applied to every dimension. The slow pairs end up slightly under-interpolated, so they emit phases beyond anything seen in training — which is why practitioners find they must set a scale factor larger than s to reach a target length. YaRN’s diagnosis: the curve is continuous where the underlying phenomenon is not.

The quantity that matters: rotations per context

YaRN asks a sharper question of each pair: how many full revolutions did it complete inside the original training window?

r_i = L / λ_i     # rotations of pair i across the original context L

The answer partitions the head into two genuinely different regimes. If r_i is large — hundreds of rotations — the model saw that clock at every phase during pretraining. Positions up to L' ask for nothing new. Extrapolation is safe, and interpolation is actively harmful because it destroys local resolution.

If r_i < 1, the clock never completed a revolution in training. The model has seen only a fraction of that dimension’s cycle, so any position beyond L pushes it into phase territory it has never observed. Here extrapolation is catastrophic and interpolation is the only safe move — and it costs almost nothing, because a dimension that slow carried no fine-grained signal anyway.

NTK-by-parts: the ramp

YaRN turns that dichotomy into a per-dimension blend. Pick thresholds α and β in units of rotations (the paper uses α = 1, β = 32 for LLaMA models), then mix the fully-interpolated frequency with the untouched one via a clamped ramp:

γ(r) = 0                     if r < α     (slow  → interpolate fully)
       1                     if r > β     (fast  → leave alone)
       (r - α) / (β - α)     otherwise    (blend)

θ'_i = (1 - γ(r_i)) · (θ_i / s)  +  γ(r_i) · θ_i

Read the endpoints: at γ = 0 the pair behaves exactly like PI; at γ = 1 it is untouched, i.e. pure extrapolation. Note what is not scaled — the position m stays raw. YaRN moves frequencies, not positions.

Worked example: 4k → 64k on a 128-dim head

Take d = 128, b = 10000, L = 4096, L' = 65536 (so s = 16), α = 1, β = 32, and evaluate r_i = L / (2π · 10000^(i/64)):

i = 0 :  λ = 6.28      r = 652     γ = 1     untouched
i = 21:  λ = 129       r = 31.7    γ ≈ 1   ramp begins
i = 32:  λ = 464       r = 8.8     γ = 0.25  mostly interpolated
i = 45:  λ = 4080      r = 1.00    γ = 0     full PI from here
i = 63:  λ = 5.5×10^4  r = 0.074   γ = 0     full PI

So of 64 pairs, roughly 21 are left completely alone, 24 sit on the ramp, and 19 are interpolated exactly as PI would. A third of the head keeps full short-range resolution — adjacent tokens still differ by a radian in pair 0 — while the dimensions that would have emitted unseen phases stay pinned inside the trained range. That asymmetry, not any clever function, is where YaRN’s quality comes from.

The second half: attention temperature

Interpolation has a side effect independent of any single dimension: compressing angular differences shrinks the spread of q · k logits, so the softmax flattens and attention entropy rises. YaRN corrects this with a temperature:

attn = softmax( q·k^T / (t · sqrt(d_k)) )

empirical fit (LLaMA):   sqrt(1/t) = 0.1 · ln(s) + 1

s = 16 :  sqrt(1/t) = 1.277  →  logits × 1.63
s = 32 :  sqrt(1/t) = 1.347  →  logits × 1.81

The elegant part is that it is free. Rotation commutes with scalar multiplication, so scaling both q and k by sqrt(1/t) multiplies their dot product by 1/t — and you get that by multiplying the precomputed cos and sin tables by sqrt(1/t). No kernel change, no extra FLOP, unmodified FlashAttention. In Hugging Face configs this is attention_factor.

The complete transform

Both halves live entirely in table construction, run once at load time:

inputs: d, b, L, L', α=1, β=32;   s = L' / L
af = 0.1 · ln(s) + 1                      # attention factor = sqrt(1/t)

for i in 0 .. d/2 - 1:
    θ_i  = b^(-2i/d)
    r_i  = L · θ_i / (2π)                 # rotations over original context
    γ_i  = clamp( (r_i - α) / (β - α), 0, 1 )
    θ'_i = (1 - γ_i) · θ_i / s + γ_i · θ_i

for m in 0 .. L' - 1:
    cos[m, i] = cos(m · θ'_i) · af        # tables: [L', d/2]
    sin[m, i] = sin(m · θ'_i) · af

That is the whole method: d/2 frequencies recomputed, two tables scaled by a constant. Nothing in the attention kernel, the weights, or the KV-cache layout changes — which is also why a wrong line here produces no crash, only a quietly worse model.

Fine-tuning budget and dynamic YaRN

YaRN degrades gracefully with zero fine-tuning, which makes it the default choice for a quick context bump. Its headline result, though, is efficiency under fine-tuning: it reaches target context with roughly 0.1% of the original pretraining tokens — on the order of 400 optimizer steps — where PI-style recipes need about ten times the tokens. It also generalizes somewhat past the length it was tuned on.

Dynamic YaRN recomputes s = max(1, l / L) from the current sequence length l, so short prompts run at s = 1 with no interpolation, removing the short-context regression a static scale imposes. The catch is autoregressive: when s changes, cached keys were rotated with stale frequencies. You accept the inconsistency, recompute the cache at boundaries, or step s in coarse jumps — the last is what most serving stacks do.

What it costs on a CPU-bound small model

YaRN’s runtime cost is zero: the per-pair frequencies are d/2 scalars and the temperature is a constant multiply on tables you already build. The only new resource is table memory: cos and sin at [65536, 64] in fp32 is 2 × 65536 × 64 × 4 B ≈ 33 MB — fine on a laptop, annoying on an embedded target, avoidable by rotating on the fly.

The real wall is elsewhere. Consider a 1B-parameter SLM with 24 layers, 8 KV heads and head dim 64: each token caches 2 × 8 × 64 × 24 = 24576 values, about 49 KB in fp16. At 64k tokens that is ~3.2 GB of KV cache, and prefill attention is O(N^2) in a bandwidth-starved setting. YaRN makes long context coherent, not affordable. Pair it with GQA, KV quantization, and chunked prefill.

Pitfalls that silently break it

Every common YaRN bug is silent: the model runs, the loss is finite, the output is just worse.

Wrong L. The scale factor is target length over the original trained context, not whatever the tokenizer config advertises. For a model already shipped with an enlarged base (Llama 3’s b = 5×10^5), YaRN computed against the wrong L double-scales the slow dimensions.

Scaling positions too. YaRN modifies θ only; a leftover PI-style m / s applies the squeeze twice.

Porting α and β blindly. They are thresholds on rotations, so they depend on b, d, and L. Change any of those and the crossover pairs move.

Dropping the temperature. The ramp is the famous half, but attention_factor is worth measurable perplexity at s ≥ 16. Validate with passkey retrieval at several depths, not average perplexity, which hides positional blind spots.