YaRN: Yet Another RoPE Extensionn Method

The Problem YaRN Solves

Both linear interpolation and NTK scaling are global — they apply the same transformation to all RoPE frequency dimensions. But different dimensions encode different kinds of positional information:

• High-frequency dims (small wavelength): encode fine-grained local position. They should not be interpolated — compressing their cycles destroys local structure.
• Low-frequency dims (large wavelength): encode long-range position. They can be linearly interpolated without harm.
• Mid-frequency dims: need something in between.

YaRN handles each group differently.

Visual: How YaRN Treats Different Frequency Bands

The Three Zones

YaRN divides the \(d/2\) frequency dimensions into three groups based on their wavelength \(\lambda_i = 2\pi/\theta_i\) relative to the training length \(L\) and target length \(L'\):

Default hyperparameters: \(\alpha = 1\), \(\beta = 32\) (tuned empirically). The ramp function smoothly interpolates between the two strategies across the mid-frequency range.

The Ramp Function

For each dimension \(i\), YaRN defines a blending factor \(r(i)\):

r(i) = 0   if high-frequency (no change)
r(i) = 1   if low-frequency (full interpolation)
r(i) = smooth ramp  otherwise

The effective frequency for dimension \(i\) becomes:

Where \(s = L'/L\) is the scale factor. When \(r(i) = 0\): \(\theta_i\) unchanged (high-freq). When \(r(i) = 1\): \(\theta_i / s\) (full interpolation). In between: a blend.

This gives each dimension group the treatment it needs, rather than applying a single global rule.

The Attention Temperature Fix

A subtlety that NTK scaling ignores: when you change RoPE frequencies, the distribution of attention logits shifts. Longer contexts naturally produce larger dot products, and the softmax temperature becomes miscalibrated.

YaRN addresses this with a learned attention temperature correction:

Where \(t = 0.1 \cdot \ln(s) + 1\) (with \(s = L'/L\)). For \(s = 4\) (4× context extension), \(t \approx 1.138\).

This dampens attention logits slightly, keeping the softmax distribution well-calibrated at longer contexts. Without this correction, models tend to “spread” attention too uniformly at long range — a well-known failure mode.

Worked Example: YaRN on LLaMA-2 (4k → 32k)

Model: LLaMA-2 7B, \(L = 4096\), \(L' = 32768\), scale \(s = 8\). Head dim \(d = 128\), so \(d/2 = 64\) frequency pairs. Default \(\alpha = 1\), \(\beta = 32\).

Wavelength for dimension \(i\):

Classify each dimension:

• \(\lambda_i < L \cdot \alpha = 4096\) → high-freq, no change (\(r = 0\))
• \(\lambda_i > L' \cdot \beta = 32768 \times 32 = 1{,}048{,}576\) → low-freq, full interpolation (\(r = 1\))
• Otherwise → mid-freq ramp

For \(i = 60\) (near highest freq): \(\lambda_{60} = 2\pi \cdot 10000^{60/64} \approx 42{,}000\). Since \(42{,}000 > 4096\) but \(< 1{,}048{,}576\) → mid-freq ramp.

For \(i = 10\) (low-freq): \(\lambda_{10} = 2\pi \cdot 10000^{10/64} \approx 590\). Since \(590 < 4096\) → high-freq, unchanged.

For \(i = 0\) (lowest freq): \(\lambda_0 = 2\pi \cdot 10000^{0} = 6.28\). Since \(6.28 < 4096\) → high-freq, unchanged.

Temperature correction for \(s = 8\):

All attention logits are divided by an extra factor of \(1.208\) on top of \(\sqrt{d_k}\), keeping softmax sharper at 32k context than without the correction.

Why YaRN Became So Popular

It hits a practical sweet spot: much better than naive interpolation, easier to deploy than fully model-specific search, and cheap enough that many RoPE-based LLMs adopted it as the default long-context extension recipe.

Training Recipe

YaRN requires minimal fine-tuning:

1. Modify RoPE with the three-zone frequency scheme
2. Apply attention temperature correction
3. Fine-tune for ~400 steps on long-context data (compared to thousands for full context extension training)

This makes YaRN practical: you can take an existing model (e.g., LLaMA-2 7B trained at 4096 tokens) and extend it to 128k context with a short fine-tuning run.

Results vs Other Methods

YaRN consistently outperforms both methods on long-context benchmarks (SCROLLS, LongBench) at the same scale, with less fine-tuning than linear interpolation.

Models Using YaRN

• Mistral 7B v0.2 (context extension from 8k to 32k)
• Qwen2 series (various context lengths)
• LLaMA-2 fine-tuned variants (community-produced 32k/64k/128k models)

YaRN is the standard method for context extension in the open-source community.

Comparison of Context Extension Methods

Summary

YaRN improves on earlier RoPE extension methods by:

1. Treating different frequency bands differently (local, transitional, long-range)
2. Correcting attention temperature to maintain focus at long context
3. Requiring minimal fine-tuning (~400 steps)

It is the current community standard for extending the context of open-weight models, used in Mistral and many LLaMA derivatives.

References

• Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. ICLR 2024 (YaRN: combines NTK-Aware Scaling with attention temperature correction and frequency-interpolation to extend context with minimal fine-tuning).
• Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv 2021 (RoPE: the base positional encoding that YaRN extends for long-context models).
• Chen, S., Wong, S., Luo, L., & Tian, Y. (2023). Extending Context Window of Large Language Models via Positional Interpolation. arXiv 2023 (Position Interpolation: the predecessor approach that YaRN improves upon by addressing the frequency-dimension asymmetry).