YaRN and NTK-Aware Scaling: How Llama 3 Got to 128K

YaRN ("Yet another RoPE extensioN", Peng et al. 2023) is the most successful technique in production for extending RoPE-based models to long context. Both Llama 3 and many community fine-tunes use YaRN or YaRN-shaped recipes.

The insight YaRN exploits

RoPE encodes position via sinusoids at many frequencies. The high-frequency dimensions (small θ) wrap around quickly — they're the ones that distinguish nearby positions. The low-frequency dimensions wrap slowly — they're the ones that distinguish far-apart positions, and therefore the ones that go off-distribution when you extend context.

Naive Position Interpolation rescales all frequencies uniformly, which works but compresses the local patterns the model already learned. NTK-aware scaling rescales only the low-frequency dimensions, leaving high-frequency dimensions alone. YaRN refines this further with a temperature-scaled attention modification that improves the result.

The receipts

YaRN demonstrated that LLaMA models could be extended from 4K → 128K context with about 10× fewer fine-tuning tokens than full retraining and 2.5× fewer training steps than prior interpolation methods. Llama 3 used a YaRN-style schedule to extend its 8K-trained base to a 128K context window in the released model card.

Code

NTK-aware scaling — sketch·python

# Standard RoPE inverse frequencies:
# inv_freq[i] = 1 / (base^(2i / d))

# NTK-aware: scale base proportionally to context extension
def ntk_aware_inv_freq(d_head, scaling_factor, base=10000.0):
    # If we want to extend context by `scaling_factor`,
    # adjust base so that the lowest-frequency dimension
    # reaches the new range while leaving high-frequency dims alone.
    new_base = base * (scaling_factor ** (d_head / (d_head - 2)))
    inv_freq = 1.0 / (new_base ** (torch.arange(0, d_head, 2).float() / d_head))
    return inv_freq

# YaRN adds:
#  - per-dimension interpolation (the "ramp")
#  - attention temperature: scale attention scores by t = 0.1 ln(s) + 1
#    where s is the scaling factor, to maintain entropy.

Exercise

Find an open-weight model with a documented context extension recipe (e.g., NousResearch/Yarn-Llama-2-7b-128k). Read its README/config. Identify which YaRN parameters they used and why. Then run perplexity on a long document at training-window length and at the extended length. How much does perplexity change?