YaRN › Summary

YaRN provides a principled and practically effective approach to extending context length in RoPE-based language models. By combining selective frequency interpolation via the ramp function with attention temperature scaling, it addresses the two distinct problems that simpler methods leave unsolved: frequency distortion and entropy shift.

The key insight driving YaRN is that different dimension pairs in RoPE have fundamentally different needs. High-frequency pairs (short wavelengths) have learned the full rotation cycle during training and can handle extended positions without modification. Low-frequency pairs (long wavelengths) have only seen a small arc of their rotation cycle and need interpolation to avoid extrapolation. The ramp function encodes this insight directly, applying a dimension-specific decision based on each pair's wavelength ratio $r_i = \lambda_i / L$.

The temperature correction addresses a separate problem: even with perfect frequency adjustments, interpolation compresses the geometry of the rotated embedding space. This makes all keys look more similar to every query and flattening the attention distribution. The formula $t = 0.1 \ln(s) + 1$ provides a calibrated correction that grows logarithmically with the extension factor. The logarithmic form matches the empirical observation that entropy shift saturates rather than growing linearly.

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 increases attention entropy. YaRN compensates with a temperature scaling factor $\sqrt{t}$ where $t = 0.1 \ln(s) + 1$, applied to pre-softmax scores.

• 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 that addresses entropy.

• Practical deployment: YaRN has been integrated into many open-source models and inference frameworks. This shows its practical viability for real-world context extension at scales from 4x to 32x.

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.