YaRN: Extending Context Length with Selective Interpolation and Temperature Scaling

Extending context length in large language models has become a central challenge. Models trained on sequences of 2,048 or 4,096 tokens struggle when asked to process documents spanning tens of thousands of positions. Position Interpolation offered one solution: scale down position indices to fit within the trained range. NTK-aware scaling improved upon this by preserving high-frequency components that capture local relationships. But both methods still leave performance on the table, particularly as the extension factor grows large.

YaRN, which stands for "Yet another RoPE extensioN," addresses limitations in both approaches. The key insight is that extending context length doesn't just require fixing position encodings. It also requires compensating for a subtle but significant change in attention score distributions. When we stretch positions across longer sequences, the entropy of attention distributions shifts in ways that degrade model behavior. YaRN tackles this through a combination of targeted frequency interpolation and an attention temperature correction.

This chapter develops YaRN from its motivation through its complete formulation. We'll see why existing methods fall short at large extension factors, derive the attention scaling mechanism, and implement YaRN step by step. By the end, you'll understand how this technique enables models to maintain quality at context lengths 16x or 32x beyond their training distribution.

Before diving into YaRN, let's understand why Position Interpolation and NTK-aware scaling leave room for improvement. Both methods successfully enable RoPE-based models to process longer sequences, but they introduce subtle distortions that accumulate as the extension factor increases.

Position Interpolation works by scaling all position indices by a factor $s$: a position $m$ becomes $m/s$. This keeps all positions within the trained range, but it compresses the entire frequency spectrum uniformly. High-frequency dimensions that previously distinguished positions 1 apart now need to distinguish positions $s$ apart after compression. The model loses fine-grained position discrimination.

NTK-aware scaling addresses this by adjusting the base frequency rather than scaling positions directly. This preserves high-frequency components while stretching low-frequency ones. The approach works well for moderate extensions (2x to 4x), but at larger factors, even NTK-aware scaling begins to degrade.

Let's visualize how these methods transform the frequency spectrum:

The visualization reveals a fundamental trade-off. Position Interpolation maintains the relative spacing between frequencies (the lines are parallel on a log scale) but shifts everything down. NTK-aware scaling bends the curve, keeping high frequencies close to original while aggressively stretching low frequencies. Neither approach is optimal across all dimension pairs.

Beyond frequency adjustments, there's a more subtle issue that neither PI nor NTK-aware scaling addresses: attention entropy shifts. When we modify RoPE frequencies, we change the distribution of attention scores in ways that affect model behavior.

Recall that attention scores are computed as the scaled dot product between query and key vectors:

where:

• $\text{score}(m, n)$: the attention score between a query at position $m$ and a key at position $n$
• $\mathbf{q}_m$: the RoPE-rotated query vector at sequence position $m$
• $\mathbf{k}_n$: the RoPE-rotated key vector at sequence position $n$
• $d_k$: the dimension of the key vectors, used for scaling to prevent dot products from growing too large
• $\mathbf{q}_m^\top \mathbf{k}_n$: the dot product between query and key, measuring their similarity

When we interpolate positions, the effective distances between tokens change. Positions that were far apart now appear closer in the rotated space. This compresses the range of attention scores, making them more uniform. Higher uniformity means higher entropy: the model attends more diffusely rather than focusing on specific positions.

Let's quantify this effect by computing attention entropy under different interpolation schemes:

The entropy values reveal the effect of position interpolation on attention patterns. While the differences might seem small in absolute terms, they compound across layers and affect the model's ability to retrieve and aggregate information from specific positions. YaRN addresses this by introducing an attention temperature correction.

YaRN combines two complementary mechanisms to enable high-quality context extension:

1. Ramp-based frequency interpolation: Rather than treating all dimension pairs equally (like PI) or smoothly transitioning across all pairs (like NTK), YaRN uses a ramp function that applies no interpolation to high-frequency pairs, full interpolation to low-frequency pairs, and a smooth transition in between.

2. Attention temperature scaling: YaRN introduces a scaling factor $\sqrt{t}$ to attention scores, where $t \geq 1$ is derived to compensate for the entropy increase caused by interpolation. When $s = 1$ (no extension), $t = 1$ and no scaling is applied.

The combined approach modifies both the rotation frequencies and the attention computation. For frequencies, YaRN applies a dimension-specific adjustment:

where:

• $\theta'_i$: the adjusted frequency for dimension pair $i$ after YaRN modification
• $\theta_i$: the original RoPE frequency for dimension pair $i$, computed as $1/10000^{2i/d}$
• $\gamma(r_i)$: the interpolation factor, a value between $1/s$ and $1$ that depends on the wavelength ratio
• $r_i = \lambda_i / L$: the wavelength-to-context ratio for dimension pair $i$
• $\lambda_i$: the wavelength for dimension pair $i$ (positions per full rotation)
• $L$: the original training context length

For attention scores, YaRN introduces a temperature correction:

where:

• $\text{score}'(m, n)$: the temperature-adjusted attention score
• $\sqrt{t}$: the temperature scaling factor, where $t > 1$ increases score magnitude
• $t$: the temperature parameter, computed from the extension factor $s$

Let's develop each component in detail.

To understand why some dimension pairs need interpolation while others don't, we need to think about rotation in terms of wavelength rather than frequency. While frequency tells us how fast something rotates (radians per position), wavelength tells us how far we must travel for one complete cycle (positions per full rotation). Wavelength provides a more intuitive picture because we can directly compare it to the context length.

Think of it this way: if a dimension pair completes 100 full rotations during the training context, it has learned to use those rotations to distinguish positions at a fine-grained level. Extending the context by 4x means it now completes 400 rotations. No problem. The pair can still distinguish positions just as well as before. But if a dimension pair completes only 0.1 rotations during training (barely moving at all), extending by 4x means it now needs to cover 0.4 rotations. That's still not a complete cycle, and the model has never seen rotation patterns beyond what it learned during training.

The wavelength $\lambda_i$ of dimension pair $i$ is the inverse of the frequency, scaled by $2\pi$:

where:

• $\lambda_i$: the wavelength (in positions) for dimension pair $i$, representing how many sequence positions correspond to one complete rotation
• $\theta_i = 1/10000^{2i/d}$: the base frequency for dimension pair $i$, measured in radians per position
• $d$: the total embedding dimension (must be even)
• $i$: the dimension pair index (0, 1, ..., $d/2 - 1$)
• $2\pi$: the number of radians in a complete rotation (one full cycle)
• $10000$: the RoPE base constant, which controls the frequency range

The second equality follows by substituting the definition of $\theta_i$ and simplifying: $2\pi / (1/10000^{2i/d}) = 2\pi \cdot 10000^{2i/d}$.

The critical question becomes: how does each wavelength compare to the training context length $L$? If $\lambda_i < L$, the pair completes at least one full rotation during training, meaning the model has seen all possible rotation states. If $\lambda_i > L$, the pair completes less than one rotation, meaning the model has only seen a fraction of the rotation cycle. This ratio $r_i = \lambda_i / L$ will be central to YaRN's design.

The "vs Original L" column shows the wavelength ratio $r_i = \lambda_i / L$. This ratio is the key to understanding which dimension pairs need interpolation:

• When $r_i < 1$, the wavelength is shorter than the context, meaning the pair completes multiple full rotations during training. These pairs have learned the full rotation cycle.
• When $r_i > 1$, the wavelength exceeds the context, meaning the pair doesn't complete even one rotation during training. These pairs operate in a limited portion of the rotation cycle.

Let's trace through the table to build intuition. Dimension pair 0 has a wavelength of about 6 positions and $r_0 \approx 0.003$. It completes roughly 340 full rotations within the training context, so it has thoroughly learned how to use rotation for position encoding. When we extend to 8,192 positions, pair 0 still completes over 1,300 rotations. No problem here.

At the other extreme, dimension pair 31 has a wavelength of around 60,000 positions and $r_{31} \approx 29$. During training, it completes only about 3% of a single rotation. The model has learned to encode position using just this small arc of the rotation cycle. If we extend the context by 4x without interpolation, we'd ask pair 31 to cover 12% of its rotation cycle, using rotation states it has never encountered. This is extrapolation, and it degrades model quality.

The pattern is clear: pairs with small $r_i$ don't need interpolation; pairs with large $r_i$ do. The question is where to draw the line, and whether the transition should be abrupt or smooth. YaRN's ramp function provides the answer.

Now we can design a function that decides how much interpolation each dimension pair receives. We want three behaviors:

1. For pairs with small $r_i$ (short wavelengths, many rotations during training): apply no interpolation, leaving them at their original frequencies.
2. For pairs with large $r_i$ (long wavelengths, few rotations during training): apply full interpolation, scaling them by $1/s$ just like Position Interpolation would.
3. For pairs with intermediate $r_i$: apply partial interpolation, blending smoothly between the two extremes.

This is exactly what a ramp function achieves. YaRN defines $\gamma(r)$ as a piecewise function:

where:

• $r = \lambda_i / L$: the wavelength-to-context ratio for dimension pair $i$, comparing the rotation period to the training context
• $\lambda_i$: the wavelength for dimension pair $i$
• $L$: the original training context length (e.g., 2048 or 4096 tokens)
• $s$: the extension scale factor (e.g., 4 for extending from 2048 to 8192)
• $\alpha$: the lower threshold (typically 1.0), below which no interpolation is applied
• $\beta$: the upper threshold (typically 32.0), above which full interpolation is applied
• $w = (r - \alpha) / (\beta - \alpha)$: the interpolation weight in the ramp region, ranging from 0 to 1
• $\gamma(r)$: the output interpolation factor, ranging from 1 (no interpolation) to $1/s$ (full interpolation)

Let's unpack the middle case, which is the most interesting. The weight $w$ measures where $r$ falls within the transition region $[\alpha, \beta]$. At the left edge ($r = \alpha$), we have $w = 0$, so the expression becomes $(1 - 0) + 0 \cdot (1/s) = 1$. At the right edge ($r = \beta$), we have $w = 1$, giving $(1 - 1) + 1 \cdot (1/s) = 1/s$. For values between the edges, we get a linear blend. This creates a smooth ramp rather than an abrupt step, which helps the model adapt more gracefully.

To summarize, the ramp function creates three distinct regions:

1. Short wavelengths ($r < \alpha$): These pairs already rotate many times within the original context. They don't need interpolation because they can naturally handle the extended positions.

2. Long wavelengths ($r > \beta$): These pairs rotate slowly and need full Position Interpolation treatment to avoid extrapolation.

3. Middle wavelengths ($\alpha \leq r \leq \beta$): These pairs receive partial interpolation, smoothly blending between the two extremes.

Let's visualize the ramp function and its effect on frequencies:

The ramp function creates a piecewise linear transition on the log-wavelength scale. The left panel shows how $\gamma$ drops from 1 to $1/s$ as wavelength increases. The right panel shows the resulting frequency spectrum: YaRN preserves high frequencies (matching the original) while smoothly transitioning to interpolated low frequencies.

We've now addressed the frequency problem with the ramp function. But recall that we identified two problems at the start: frequency distortion and entropy shift. Even with perfect frequency adjustments, interpolation changes something fundamental about attention patterns.

To understand why, think about what interpolation does geometrically. When we slow down rotations (by multiplying frequencies by $\gamma < 1$), positions that were previously "far apart" in the rotated embedding space now appear "closer." Consider two tokens at positions 0 and 100. In the original RoPE, pair 0 might rotate these to be 100 radians apart (after wrapping). With interpolation, the same pair rotates them to only 25 radians apart (for 4x extension). This compression happens across all interpolated dimension pairs.

This compression has a direct consequence for attention scores. The dot product between query and key vectors measures their similarity. When positions appear closer together in the rotated space, the dot products become more similar to each other. The range of attention scores shrinks.

Now recall how softmax works. Given a vector of scores $\mathbf{z}$:

When input scores are spread out (high variance), exponentiating amplifies the differences, producing a peaked distribution that focuses on the highest scores. When input scores are compressed (low variance), the exponentials are more similar, producing a flatter, more uniform distribution. This increased uniformity is what we measure as higher entropy.

The solution is conceptually simple: if scores are too compressed, stretch them back out. We do this by multiplying the scores by a factor greater than 1 before applying softmax. In temperature scaling terminology, we reduce the "temperature" (making the distribution sharper) by dividing by a temperature $T < 1$, or equivalently, multiplying by $1/T$. YaRN frames this as multiplying by $\sqrt{t}$ where $t > 1$.

The temperature factor is computed as:

where:

• $t$: the attention temperature scaling factor (always $\geq 1$)
• $s$: the context extension scale factor (e.g., 4 for 4x extension)
• $\ln(s)$: the natural logarithm of the extension factor
• $0.1$: an empirically determined coefficient that controls the rate of temperature increase
• $1$: the base value ensuring $t = 1$ when $s = 1$ (no extension)

Why $\sqrt{t}$ rather than $t$ directly? This relates to how variance scales under multiplication. If we multiply scores by a constant $c$, the variance of the scores is multiplied by $c^2$. So to increase variance by a factor of $t$, we need to multiply the scores by $\sqrt{t}$.

The logarithmic form of $t$ captures an empirical observation: the entropy shift doesn't grow linearly with the extension factor. Larger extensions need progressively less aggressive correction. Doubling the extension factor from 4x to 8x doesn't require doubling the temperature adjustment. Instead, it adds a constant amount: $\ln(8) - \ln(4) = \ln(2) \approx 0.69$, contributing an additional $0.069$ to $t$.

The $\sqrt{t}$ column shows the actual multiplier applied to attention scores. Even for very large extension factors (128x), the multiplier stays below 1.2, indicating that temperature correction is a subtle adjustment rather than a dramatic rescaling.

The temperature factor grows slowly with extension factor. For a 4x extension, we multiply attention scores by $\sqrt{1.139} \approx 1.067$. For 32x extension, the multiplier is $\sqrt{1.347} \approx 1.161$. These modest corrections help maintain attention sharpness.

Let's verify the effect on entropy:

Temperature scaling helps bring the entropy closer to the original distribution. While the match isn't perfect (which is expected given the fundamental changes to the attention geometry), the correction prevents the excessive entropy increase that would otherwise occur.

Now we can put together the complete YaRN transformation. Given a model trained on context length $L$ that we want to extend by factor $s$, the process involves three steps.

Step 1: Compute adjusted frequencies

For each dimension pair $i$ with base frequency $\theta_i$ and wavelength $\lambda_i = 2\pi/\theta_i$, compute the adjusted frequency:

where:

• $\theta'_i$: the adjusted frequency for dimension pair $i$
• $\theta_i$: the original RoPE frequency
• $\gamma(\cdot)$: the ramp function defined earlier
• $\lambda_i / L$: the wavelength ratio that determines the interpolation amount

Step 2: Apply RoPE with adjusted frequencies

For query or key vector $\mathbf{x}$ at position $m$, apply the block-diagonal rotation matrix:

where:

• $\text{YaRN-RoPE}(\mathbf{x}, m)$: the rotated vector, now incorporating YaRN's frequency adjustments
• $R(\theta'_i \cdot m)$: a 2$\times$2 rotation matrix for angle $\theta'_i \cdot m$, applied to dimension pair $i$
• $\theta'_i \cdot m$: the rotation angle, which increases linearly with position $m$ at the adjusted rate
• The block-diagonal structure means each dimension pair is rotated independently

Step 3: Apply temperature scaling to attention scores

After computing the raw attention scores, apply the temperature correction:

where:

• $\text{score}(m, n)$: the scaled attention score between positions $m$ and $n$
• $\mathbf{q}_m$: the YaRN-rotated query vector at position $m$
• $\mathbf{k}_n$: the YaRN-rotated key vector at position $n$
• $d_k$: the dimension of key vectors
• $t = 0.1 \ln(s) + 1$: the temperature factor, computed from the extension scale
• $\sqrt{t}$: the scaling multiplier applied to raw scores (increases score magnitude to sharpen attention)
• The scaling occurs before softmax normalization

Let's test the complete YaRN implementation:

The implementation confirms that YaRN preserves vector magnitudes (rotation is an isometry) while applying the adjusted frequencies.

Let's visualize how YaRN affects the attention patterns compared to other methods:

The attention heatmaps reveal the differences between methods. YaRN produces patterns that more closely resemble the original RoPE attention, with sharper focus on relevant positions rather than the diffuse patterns seen with Position Interpolation.

An important practical consideration is whether YaRN requires fine-tuning. The answer depends on the extension factor:

For modest extensions (2x to 4x): YaRN can often be applied without any fine-tuning, leveraging the model's existing weights. The frequency adjustments and temperature scaling are designed to minimize the distribution shift.

For larger extensions (8x to 32x): Brief fine-tuning significantly improves quality. The YaRN paper recommends 200 to 400 training steps on long-context data, much less than training from scratch.

For extreme extensions (64x+): Longer fine-tuning becomes necessary, though still much shorter than pre-training from scratch.

The key advantage of YaRN over naive approaches is the efficiency of this fine-tuning. Because the method preserves the essential structure of the position encoding while making targeted adjustments, the model needs to learn only minor adaptations rather than fundamentally new position representations.

Let's summarize how YaRN compares to other context extension methods:

The main advantages of YaRN are:

1. Selective preservation: High-frequency dimensions that don't need interpolation are left unchanged, maintaining fine-grained position discrimination.

2. Entropy control: Temperature scaling prevents the attention distribution from becoming too diffuse, preserving the model's ability to focus.

3. Efficient adaptation: The targeted adjustments require less fine-tuning to adapt to than more disruptive methods.

4. Predictable behavior: The ramp function provides intuitive control over which dimensions are interpolated and by how much.

The "Mean Entropy" column shows how diffuse the attention distribution is, with lower values indicating sharper focus. The "Rel. Pos. Error" measures whether the same relative distance produces the same attention score at different absolute positions. Values near zero indicate perfect relative position encoding.

The comparison shows that YaRN achieves entropy closer to the original while maintaining the relative position property (low position error). The combination of targeted interpolation and temperature scaling addresses both the frequency and entropy aspects of context extension.

Despite its effectiveness, YaRN has limitations worth understanding.

Hyperparameter sensitivity: The $\alpha$ and $\beta$ parameters control the ramp function's transition region. While the defaults (1 and 32) work well for many models, some architectures may benefit from tuning. Models with different RoPE bases or dimension sizes may require adjustment.

Temperature approximation: The temperature formula $t = 0.1 \ln(s) + 1$ is empirically derived rather than theoretically optimal. It works well in practice but may not perfectly compensate for entropy shifts in all scenarios.

Integration complexity: Unlike Position Interpolation (which only modifies position indices) or NTK-aware scaling (which only modifies the base), YaRN requires two separate modifications. Both the frequency adjustment and the attention scaling must be implemented correctly.

Interaction with other optimizations: When combined with techniques like FlashAttention or grouped-query attention, care must be taken to apply the temperature scaling at the correct point in the computation. The scaling should occur before the softmax, not after.

These limitations are manageable in practice. YaRN has been successfully integrated into many open-source models, and the community has developed reference implementations that handle the integration details correctly.

When implementing YaRN, several parameters control its behavior:

• scale: The context extension factor (e.g., 4.0 for extending from 2048 to 8192 tokens). Larger values enable longer contexts but may require more fine-tuning.

• alpha: The lower wavelength threshold (default: 1.0). Dimension pairs with wavelength-to-context ratio below this value receive no interpolation. Decreasing $\alpha$ applies interpolation to more dimension pairs.

• beta: The upper wavelength threshold (default: 32.0). Dimension pairs with wavelength-to-context ratio above this value receive full interpolation. Increasing $\beta$ delays full interpolation to higher wavelength pairs.

• base: The RoPE frequency base (default: 10000). This should match the base used in the original model's RoPE implementation.

• original_context: The context length the model was originally trained on (e.g., 2048 or 4096). This determines the wavelength ratios used in the ramp function.

For most applications, the default values of $\alpha = 1$ and $\beta = 32$ work well. The primary parameter to adjust is scale, which should be set to the desired extension factor. If the model uses a non-standard RoPE base, ensure base matches the original configuration.

YaRN provides a principled approach to extending context length in RoPE-based language models. By combining selective frequency interpolation with attention temperature scaling, it addresses limitations in both Position Interpolation and NTK-aware methods.

Key takeaways:

• Wavelength-based interpolation: YaRN uses a ramp function to apply no interpolation to high-frequency dimension pairs, full interpolation to low-frequency pairs, and a smooth transition in between. This preserves fine-grained position information where it matters.

• Attention temperature correction: Context extension changes attention score distributions, increasing entropy. YaRN compensates with a temperature scaling factor $\sqrt{t}$ where $t = 0.1 \ln(s) + 1$.

• Efficient fine-tuning: The targeted adjustments minimize distribution shift, enabling effective context extension with as few as 200 to 400 fine-tuning steps for moderate extension factors.

• Complementary to other methods: YaRN can be seen as a refinement that combines insights from Position Interpolation (full interpolation for long wavelengths) and NTK-aware scaling (preservation of short wavelengths), while adding the novel temperature correction.

• Practical deployment: YaRN has been integrated into many open-source models and inference frameworks, demonstrating its practical viability for real-world context extension.

The next chapter explores attention sinks, a phenomenon where transformer models allocate disproportionate attention to initial tokens regardless of their semantic relevance, and how this affects long-context processing.