YaRN vs NTK-Aware RoPE Scaling: Why Long Context Breaks

Take a Llama-class model trained at 8k context, set max_position_embeddings to 128k, and feed it a 40k-token document. It will not error. It will produce fluent, confident, completely wrong output — often degenerating into repetition somewhere past the original training length. Nothing crashed. The rotary embeddings just handed the attention layer angles it has never seen.

RoPE scaling is the family of fixes for that, and the three main variants — Position Interpolation, NTK-aware, and YaRN — are not interchangeable. Picking the wrong one silently degrades either your long-context recall or your short-context quality, and the failure is invisible unless you eval both.

TL;DR

• RoPE encodes position as a rotation angle m·θ_i per dimension pair. Low-frequency pairs have wavelengths longer than the training context, so the model has only ever seen a partial arc of them. Extrapolating past training length puts those dims in unvisited angular territory.
• Position Interpolation (PI) divides every position by a scale factor s. It works, but it compresses high-frequency dims too — the ones encoding local token ordering — so short-range performance drops.
• NTK-aware scaling instead raises the RoPE base (10000 → 10000·s^(d/(d-2))), which interpolates low frequencies heavily and high frequencies barely. No fine-tuning needed, but it under-scales the lowest frequencies and still extrapolates at the edges.
• YaRN = NTK-by-parts + attention temperature. It partitions dimensions by wavelength (extrapolate short, interpolate long, ramp the middle) and sharpens attention logits by 1/t where √(1/t) = 0.1·ln(s) + 1. This is what Llama 3.1's rope_type: "llama3" config implements a variant of.
• Biggest production trap: applying a static scale at short sequence lengths. Use dynamic scaling, and never mix KV cache entries computed under different RoPE configs.

Why does RoPE break past the training context length?

Because the low-frequency dimension pairs never complete a full rotation during training, so the model learns their behavior only over a fraction of the unit circle.

RoPE splits each attention head's d-dim vector into d/2 pairs and rotates pair i by angle m·θ_i, where m is the token position and:

θ_i = base^(-2i/d)        base = 10000, d = 128 (typical head_dim)
λ_i = 2π / θ_i            # wavelength in tokens

Run the numbers for d=128, 64 pairs:

pair i	wavelength λ (tokens)
0	6.3
32	628
50	~8,200
63	~54,400

At an 8192-token training context, roughly the first 50 pairs complete at least one full rotation. The top ~14 never do. Pair 63 sees only about 15% of a single cycle across the entire training corpus.

That asymmetry is the whole story. High-frequency dims are well-trained periodic features — the model knows what angle 5.0 rad means because it saw it thousands of times per sequence. Low-frequency dims are effectively a monotonic ramp the model uses as a coarse "how far back is this" signal, and it has only ever seen the first sliver of that ramp.

Push position to 40,000 and pair 55's angle lands in a region with zero training support. Attention logits, which depend on cos(θ_i·(m−n)) terms, go out of distribution. In practice you see logit magnitudes blow up and the attention distribution collapse onto a few tokens.

Why does Position Interpolation hurt short-context quality?

Because PI is uniform: it squeezes the well-trained high-frequency dims by the same factor as the untrained low-frequency ones.

PI (Chen et al.) rescales the position itself:

m' = m / s     where s = target_len / train_len

Every angle now lands inside the trained range. Clean, one line of code, and it genuinely works with a short fine-tune. But consider pair 0 with λ = 6.3 tokens. At s = 8, adjacent tokens now differ by 0.16 rad instead of 1.0 rad. The dimension that distinguished "the token immediately before" from "three tokens before" has had its resolution cut 8x. You have taken your local-syntax machinery and blurred it to save your long-range machinery.

Measured effect: PI models fine-tuned for long context reliably regress on short-sequence perplexity relative to the base model. If your workload is 95% sub-8k requests with occasional 100k documents, uniform PI is a bad trade.

How does NTK-aware scaling fix the frequency imbalance?

By scaling the RoPE base instead of the position, which produces an interpolation factor that varies smoothly with frequency — near-zero for high frequencies, near-full for low ones.

base_new = base * s ** (d / (d - 2))   # d = head_dim

For d=128, s=8, that's 10000 · 8^(64/63) ≈ 81,700. Pair 0's wavelength moves from 6.28 to ~6.28 (unchanged — θ_0 = base^0 = 1 regardless of base). Pair 63's wavelength grows by roughly s. Exactly the behavior you want.

The catch: the top pairs get an effective scale slightly below s, so the highest positions still extrapolate a bit. That's why practitioners running NTK-aware without fine-tuning use s' = s · 1.5 or more — you overshoot to buy margin. It's a hack that works and it's why "just bump rope_theta" is folklore that half-works.

What does YaRN add on top of NTK-aware?

Two things: a hard partition of dimensions by wavelength (NTK-by-parts) rather than a smooth curve, and a temperature correction on the attention logits.

Part 1 — the ramp. Define r_i = L / λ_i, the number of full rotations pair i completes within the original context L. Then:

• r_i > β (default 32): the pair has cycled many times. It is fully trained. Do not interpolate.
• r_i < α (default 1): the pair never completed a cycle. Fully interpolate (PI).
• In between: linear ramp.

import math

def yarn_ramp(head_dim, base, orig_ctx, scale, alpha=1.0, beta=32.0):
    """Returns per-pair inv_freq for YaRN (NTK-by-parts)."""
    idx = [2 * i for i in range(head_dim // 2)]
    inv_freq_extrap = [1.0 / (base ** (j / head_dim)) for j in idx]   # no scaling
    inv_freq_interp = [f / scale for f in inv_freq_extrap]            # pure PI

    out = []
    for f_e, f_i in zip(inv_freq_extrap, inv_freq_interp):
        wavelen = 2 * math.pi / f_e
        r = orig_ctx / wavelen                       # rotations within training ctx
        gamma = (r - alpha) / (beta - alpha)         # 0 -> interp, 1 -> extrap
        gamma = min(1.0, max(0.0, gamma))
        out.append(f_i * (1 - gamma) + f_e * gamma)
    return out

# attention temperature: fold into cos/sin so no kernel change
def yarn_mscale(scale):
    return 0.1 * math.log(scale) + 1.0   # = sqrt(1/t)

For d=128, base=10000, L=8192, β=32: pairs 0–25 (λ < 256 tokens) are left untouched, pairs ~50–63 are fully interpolated, and the ~24 pairs in between ramp. That's a concrete, inspectable partition — you can print it and sanity-check it against your model config.

Part 2 — attention temperature. Longer contexts mean more keys competing in the softmax, which raises attention entropy: the distribution flattens and the model attends to everything a little. YaRN counteracts this by scaling logits by 1/t, with √(1/t) = 0.1·ln(s) + 1.

At s = 8 that's √(1/t) ≈ 1.208, so logits are sharpened by 1.208² ≈ 1.46 — a 46% boost before softmax. The elegant part: you multiply the precomputed cos and sin tables by 1.208. Both q and k get rotated by those tables, so the dot product picks up the square automatically. Zero changes to your attention kernel, zero runtime cost.

YaRN's headline claim is that it reaches comparable long-context quality with far less continued pretraining than PI — the paper reports roughly 10x fewer tokens and 2.5x fewer training steps. Treat the exact ratios as method-and-model dependent, but the qualitative gap is large and reproducible.

What does this look like in a real config?

Llama 3.1 ships a variant of NTK-by-parts directly in config.json:

{
  "max_position_embeddings": 131072,
  "rope_theta": 500000.0,
  "rope_scaling": {
    "rope_type": "llama3",
    "factor": 8.0,
    "low_freq_factor": 1.0,
    "high_freq_factor": 4.0,
    "original_max_position_embeddings": 8192
  }
}

Note rope_theta: 500000 — the base was already raised during pretraining, before any scaling. Then factor: 8.0 extends 8192 → 65536 of "native" range, and low/high_freq_factor define the same wavelength partition as YaRN's α/β, expressed as ratios instead of rotation counts.

The equivalent for explicit YaRN in vLLM:

vllm serve <model> \
  --max-model-len 65536 \
  --rope-scaling '{"rope_type":"yarn","factor":8.0,
                   "original_max_position_embeddings":8192}'

What are the failure modes in production?

Four, in rough order of how often they bite.

1. Static scaling applied to short requests. If you serve with a fixed s=8, a 500-token prompt is being interpolated by 8x for no reason, and quality drops on the majority of your traffic. Use dynamic scaling: compute s = max(1, seq_len / orig_ctx) per request. HuggingFace exposes this as rope_type: "dynamic".

2. KV cache poisoning across configs. RoPE is applied to keys before they enter the cache. Change rope_scaling and every cached key from the old config is silently misaligned. Prefix caches, session caches, and any cross-request reuse must be keyed on the RoPE config hash. This one produces the most confusing bug reports because the model works fine on a cold cache.

3. Passkey retrieval as your only eval. Needle-in-a-haystack passes at 128k on models that cannot summarize a 30k-token document coherently. Retrieval of a verbatim string is a low bar — the low-frequency dims only need to be monotonic, not accurate. Eval on multi-hop QA or long-document summarization where the model must weigh many distant spans.

4. Forgetting the temperature. People port the YaRN inv_freq computation and drop mscale because it looks like an unrelated hack. Without it you get YaRN's frequency handling with PI-era attention entropy, and long-context accuracy sits well below what the method should deliver.

So why does long context break, and which scaling should you use?

RoPE-based long context breaks because low-frequency rotary dimensions never complete a full rotation during training, so any position past the training length rotates them into angles the model has no learned response for. Uniform Position Interpolation fixes that by squeezing all frequencies, which sacrifices the high-frequency dims that encode local token order. NTK-aware scaling fixes the imbalance by raising the RoPE base so interpolation strength varies with frequency, but under-scales at the low end. YaRN is the version to use in production: partition dimensions by how many rotations they complete within the original context, extrapolate the fast ones, interpolate the slow ones, ramp the middle, and multiply your cos/sin tables by 0.1·ln(s) + 1 to restore attention sharpness. Serve it dynamically so short prompts pay nothing, key your KV cache on the RoPE config, and eval on real long-document reasoning rather than passkey retrieval.