YaRN: Efficient Context Window Extension of Large Language Models

The problem with extending RoPE

Rotary Position Embedding encodes position as a rotation angle. Each coordinate pair $d$ in a query or key vector is rotated by $m \cdot \theta_d$, where $m$ is the token position and $\theta_d = \text{base}^{-2d/D}$ is a dimension-specific base frequency. At training time, every angle seen by the model lives in a bounded range determined by the maximum sequence length. Push context beyond that training length and some rotation angles become out-of-distribution — the model has never learned to interpret them.

The naive fix is linear position interpolation (Chen et al., 2023): divide position indices by a scale factor $s = L_{\text{target}} / L_{\text{train}}$ so that a context of length $L_{\text{target}}$ occupies the same angular range as $L_{\text{train}}$ did. This keeps all angles in-distribution, but at a cost: every dimension is compressed uniformly. High-frequency dimensions — those already rotating fast and completing many full cycles within the training length — are over-compressed. They lose the fine-grained local signal they carried, and the model needs thousands of fine-tuning steps to re-learn positional discriminability.

YaRN's non-uniform frequency scaling

YaRN's key insight is that different RoPE frequency groups need different treatment. Given a scale factor $s$, the minimum training length for dimension $d$ to see at least one full rotation cycle is $\lambda_d = 2\pi / \theta_d$. Comparing $\lambda_d$ to $L_{\text{train}}$ and $L_{\text{target}}$ partitions the dimensions into three groups:

• High-frequency dimensions ($\lambda_d \ll L_{\text{train}}$): these dimensions complete many cycles even within the original training window. They already generalize to new positions — the model has seen all their angles repeatedly. No interpolation needed; use the original frequencies unchanged.
• Low-frequency dimensions ($\lambda_d \gg L_{\text{target}}$): these complete less than one cycle even within the extended context. They are not really positional encoders in the usual sense — they act more like a global position tag. YaRN extrapolates these using NTK-aware scaling, which spreads the base frequency out to accommodate longer ranges.
• Medium-frequency dimensions: these fall between the two extremes. YaRN applies standard linear interpolation here, scaling down frequencies to bring the extended positions into the training range.

This tripartite treatment is the practical mechanism behind the NTK-by-parts approach named in the paper. Rather than a single global formula, each dimension gets the interpolation policy appropriate to its wavelength.

Attention temperature correction

Extending context length changes the statistical regime of attention beyond just position angles. As the number of keys grows, the softmax distribution over them flattens — with more keys, each receives a smaller share of probability mass on average, and the effective entropy of the distribution increases. A model trained at length $L_{\text{train}}$ has learned to operate with a certain distribution sharpness. At $s \times L_{\text{train}}$, that sharpness is gone: attention heads become diffuse and lose their ability to retrieve specific tokens precisely.

YaRN corrects this by multiplying attention logits by a temperature inverse $1/t$ before the softmax, where

$$t = 0.1 \ln(s) + 1.$$

The scale factor $s$ controls how much longer the context is, and $\ln(s)$ grows slowly, so the correction is modest and smooth. At $s = 1$ (no extension), $t = 1$ and nothing changes. At $s = 32$ (extending from 4K to 128K), $t \approx 1.35$ and logits are multiplied by $\approx 0.74$, sharpening the attention distribution back toward the entropy the model was trained on. This correction requires no additional parameters and no changes to the model architecture.

Results and adoption

Applied to LLaMA 2 7B, YaRN extends the context window from 4K to 128K tokens using roughly 400 fine-tuning steps — compared to thousands required by linear interpolation to achieve comparable perplexity. At 128K context, perplexity on long-document benchmarks is competitive with models trained natively at that length. The combination of fewer fine-tuning steps and better preservation of short-context performance made YaRN the practical route for open-weight long-context adaptation.

YaRN's influence extended well beyond the paper: Mistral 7B v0.2, Mixtral, and many derivative community models use YaRN or its NTK-by-parts frequency grouping as their context extension method. It established non-uniform frequency scaling as the correct framing for the problem — a framing that subsequent work, including LongRoPE, extended further by dropping the hand-designed group boundaries in favor of learned per-dimension factors.