YaRN › Worked Example: Tracing a Dimension Pair

Let's trace YaRN's computation for a specific dimension pair to make the mathematics concrete. We'll work through a 4x extension from context length 2,048 to 8,192 with a 64-dimensional model.

Consider dimension pair $i = 10$ (the 11th pair out of 32 total). We'll compute everything step by step.

Step 1: Compute the original frequency.

The base frequency formula gives:

$$\theta_{10} = \frac{1}{10000^{2 \times 10 / 64}} = \frac{1}{10000^{0.3125}} = \frac{1}{10000^{5/16}}$$

Computing numerically: $10000^{0.3125} = 10^{4 \times 0.3125} = 10^{1.25} \approx 17.78$, so $\theta_{10} \approx 0.0563$.

Step 2: Compute the wavelength.

The wavelength is:

$$\lambda_{10} = \frac{2\pi}{\theta_{10}} = \frac{2\pi}{0.0563} \approx 111.6 \text{ positions}$$

So one complete rotation of dimension pair 10 takes about 112 positions.

Step 3: Compute the wavelength ratio.

With original context $L = 2048$:

$$r_{10} = \frac{\lambda_{10}}{L} = \frac{111.6}{2048} \approx 0.054$$

Since $r_{10} = 0.054 < \alpha = 1.0$, dimension pair 10 falls in the "no interpolation" region. The model completes about 18 full rotations ($2048 / 111.6 \approx 18.4$) during the training context, so it has learned all possible rotation states.

Step 4: Apply the ramp function.

Since $r_{10} < \alpha$, we have $\gamma(r_{10}) = 1$. The adjusted frequency is:

$$\theta'_{10} = \theta_{10} \cdot \gamma(r_{10}) = 0.0563 \cdot 1 = 0.0563$$

Dimension pair 10 is unchanged: it rotates at its original speed.

Now let's examine dimension pair $i = 28$ (near the low-frequency end).

Step 1: Compute the original frequency.

$$\theta_{28} = \frac{1}{10000^{2 \times 28 / 64}} = \frac{1}{10000^{0.875}} \approx \frac{1}{5623} \approx 0.0001778$$

Step 2: Compute the wavelength.

$$\lambda_{28} = \frac{2\pi}{0.0001778} \approx 35,325 \text{ positions}$$

One full rotation of dimension pair 28 takes over 35,000 positions.

Step 3: Compute the wavelength ratio.

$$r_{28} = \frac{35325}{2048} \approx 17.25$$

Since $1.0 < r_{28} = 17.25 < 32.0 = \beta$, dimension pair 28 falls in the ramp (transition) region.

Step 4: Apply the ramp function.

The linear weight is:

$$w = \frac{r - \alpha}{\beta - \alpha} = \frac{17.25 - 1.0}{32.0 - 1.0} = \frac{16.25}{31.0} \approx 0.524$$

The interpolation factor is:

$$\gamma(r_{28}) = (1 - w) + w \cdot \frac{1}{s} = (1 - 0.524) + 0.524 \cdot \frac{1}{4} = 0.476 + 0.131 = 0.607$$

The adjusted frequency is:

$$\theta'_{28} = \theta_{28} \cdot \gamma(r_{28}) = 0.0001778 \cdot 0.607 \approx 0.0001079$$

Dimension pair 28 is slowed down to about 61% of its original speed. With the extended context of 8,192 positions, a token at position 8,192 would produce a rotation angle of $\theta'_{28} \times 8192 \approx 0.884$ radians, which corresponds to about 14% of a full rotation. The model has learned rotation states up to about $\theta_{28} \times 2048 = 0.1778 \times 2048 = 0.364$ radians in the original training, so 0.884 radians is outside the training distribution. The gamma correction reduces this to $\theta'_{28} \times 8192 \approx 0.884$ radians... wait, let's check with the corrected frequency: $0.0001079 \times 8192 \approx 0.884$ radians. Pair 28 still extrapolates somewhat (since we applied only partial interpolation), which is intentional. Full interpolation would be too conservative; partial interpolation strikes a balance between accessing new context and staying close to the training distribution.

Step 5: Temperature correction.

For $s = 4$:

$$t = 0.1 \cdot \ln(4) + 1 = 0.1 \cdot 1.386 + 1 = 1.139$$

The attention score multiplier is $\sqrt{1.139} \approx 1.067$. Every pre-softmax attention score is multiplied by this factor before the softmax is applied.

This worked example shows the key principle: YaRN applies a per-dimension decision based on physical intuition (how much of the rotation cycle has been trained), rather than a global formula that treats all dimensions identically.