RoPE Scaling

RoPE rotates dimension pair $i$ by $m\theta_i$ at position $m$, with $\theta_i = b^{-2i/d}$. Nothing in that formula breaks at $m = 5000$ for a model trained to 4096. What breaks is that the fast pairs have wrapped around into phase relationships the attention weights were never fit against, and the scores they produce are not merely inaccurate but unrelated to distance.

Every extension method rewrites the rotation so that a longer sequence lands inside the phase range the model already understands.

Position interpolation

The blunt version. Divide every position by the scale factor $s$, so the longer sequence occupies the same angular span as the original.

No phase is out of range, because none exceeds what training covered. The price is resolution: adjacent tokens are now $\theta_i / s$ apart instead of $\theta_i$, and at $s = 8$ the fastest pair can no longer cleanly separate neighbours. Chen et al. recover this with about 1000 steps of fine-tuning — cheap, but not free, and the model gets slightly worse at short contexts.

NTK-aware scaling

Instead of scaling positions, scale the base.

Because $\theta_i = b^{-2i/d}$, raising $b$ barely moves the high-frequency pairs ($i$ small) and moves the low-frequency pairs a great deal. Local resolution survives; the slow dimensions that encode long-range position absorb the stretch. This works with no fine-tuning at modest $s$, which is why it spread through the open-weights community before either paper was published.

YaRN

YaRN makes the frequency dependence explicit rather than implicit. Compare each pair’s wavelength $\lambda_i = 2\pi/\theta_i$ to the original context $L$:

• $\lambda_i \ll L$ — the pair completes many rotations inside the training window, so the model has seen its full phase range. Leave it alone.
• $\lambda_i > L$ — the pair never completes a rotation, so its phase is an absolute coordinate. Interpolate it fully by $1/s$.
• In between — ramp linearly between the two.

YaRN adds one more piece that is easy to miss and does real work: attention temperature. Longer contexts mean more terms in the softmax and therefore flatter distributions, so the logits are scaled by $t$ with $\sqrt{1/t} = 0.1\ln s + 1$ — a fixed multiplier folded into the rotation at no runtime cost. With both, YaRN reaches the same perplexity as position interpolation on roughly a tenth of the fine-tuning tokens.

import math

import torch
from torch import Tensor


def yarn_frequencies(
    dim: int,
    base: float = 10_000.0,
    scale: float = 8.0,
    original_context: int = 8192,
    alpha: float = 1.0,     # below this many rotations: interpolate fully
    beta: float = 32.0,     # above this many rotations: leave untouched
) -> Tensor:
    """Per-pair inverse frequencies, ramped between interpolated and original."""
    i = torch.arange(0, dim, 2).float()
    inv_freq = 1.0 / (base ** (i / dim))                      # θ_i

    # Rotations completed inside the original context window.
    rotations = original_context * inv_freq / (2 * math.pi)
    gamma = ((rotations - alpha) / (beta - alpha)).clamp(0, 1)

    return (1 - gamma) * (inv_freq / scale) + gamma * inv_freq


def attention_temperature(scale: float) -> float:
    """Softmax sharpening that compensates for the longer sum. YaRN §3.4."""
    return 0.1 * math.log(scale) + 1.0

Llama 3.1’s configuration is this shape with different names — low_freq_factor: 1, high_freq_factor: 4, factor: 8 over a base already raised to 500 000 during pre-training. The large base is the tell: the model was built to be extended, with the slow dimensions given room in advance.

Scaling makes long positions representable. It does not make the model use them. A model extended to 128k typically retains strong retrieval at the beginning and end of its window and much weaker retrieval in the middle, and no amount of frequency arithmetic addresses that — it is a property of what the attention heads learned to do, not of how position is encoded.