YaRN: Yet another RoPE extensioN for context window extension.

YaRN modifies RoPE frequency bands to handle longer sequences than the model was originally trained on. It achieves this by scaling different frequency components based on their wavelength relative to the original context length.

Key Insight

RoPE encodes position via rotations at different frequencies. For context extension, high-frequency (local position) info should be preserved while low-frequency (global position) info needs scaling:

High-frequency bands (short wavelength):
  - Capture local positional relationships
  - Left unchanged (factor = 1)

Low-frequency bands (long wavelength):
  - Capture global position in sequence
  - Scaled down by 1/scale factor

Middle bands:
  - Linear interpolation between the two regimes

Formula

For each dimension i, compute a scaling factor:

• If wavelength < high_freq_threshold: factor = 1 (unchanged)
• If wavelength > low_freq_threshold: factor = 1/scale (full scaling)
• Otherwise: linear interpolation between 1 and 1/scale

The thresholds are derived from:

• high_freq_wavelen = original_max_position / high_freq_factor
• low_freq_wavelen = original_max_position / low_freq_factor

Usage

# Build a YaRN-modified RoPE layer
model = YARN.build(
  embed_dim: 64,
  scale: 8,
  original_max_position: 2048
)

# Get the frequency table directly
freqs = YARN.yarn_freqs(64,
  scale: 8,
  original_max_position: 2048
)

# Apply YaRN to query/key tensors
{q_rotated, k_rotated} = YARN.apply_yarn(q, k,
  scale: 8,
  original_max_position: 2048
)

References

• "YaRN: Efficient Context Window Extension of Large Language Models" (Peng et al., 2023) — https://arxiv.org/abs/2309.00071