Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

YaRN › The Complete YaRN Formula

Michael BrenndoerferUpdated June 28, 20255 min read

Language AI HandbookAll chapters

Part 18 of 67Long Context

Chapter 129YaRN › The Complete YaRN Formula

Reading section 8 of 16 · The Complete YaRN Formula · 5 minSection 8 of 16 · 5 min

Section 8 of 16 · 5 min read

YaRN › The Complete YaRN FormulaOne section · Sections in this chapter8 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

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 coordinated steps. Each step builds on the previous one, and together they constitute the full YaRN algorithm.

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 wavelength ratio $r_i = \lambda_i / L$ and the adjusted frequency:

\theta'_i = \theta_i \cdot \gamma\left(\frac{\lambda_i}{L}\right)

where:

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

This step is performed once at initialization and the adjusted frequencies are stored. During inference, every query and key vector is rotated using $\theta'_i$ rather than $\theta_i$ . The computational cost is identical to standard RoPE; only the frequency values change.

Step 2: Apply RoPE with adjusted frequencies

For query or key vector $\mathbf{x}$ at position $m$ , apply the block-diagonal rotation matrix using the adjusted frequencies. Each pair of dimensions $(2i, 2i+1)$ is rotated by angle $\theta'_i \cdot m$ :

\text{YaRN-RoPE}(\mathbf{x}, m) = \begin{pmatrix} R(\theta'_0 \cdot m) & & \\ & \ddots & \\ & & R(\theta'_{d/2-1} \cdot m) \end{pmatrix} \mathbf{x}

where:

$\text{YaRN-RoPE}(\mathbf{x}, m)$ : the rotated vector, 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 at position $m$ , increasing linearly with position at the adjusted rate $\theta'_i$
The block-diagonal structure means each dimension pair is rotated independently of all others

The rotation matrix $R(\phi)$ for angle $\phi$ is:

R(\phi) = \begin{pmatrix} \cos\phi & -\sin\phi \\ \sin\phi & \cos\phi \end{pmatrix}

This is the standard 2D rotation matrix. Applying it to a pair $(x_{2i}, x_{2i+1})$ gives:

\begin{pmatrix} x'_{2i} \\ x'_{2i+1} \end{pmatrix} = \begin{pmatrix} x_{2i}\cos(\theta'_i m) - x_{2i+1}\sin(\theta'_i m) \\ x_{2i}\sin(\theta'_i m) + x_{2i+1}\cos(\theta'_i m) \end{pmatrix}

Because rotation is an isometry (it preserves vector lengths), YaRN-RoPE does not change the magnitude of query or key vectors. It only changes their direction, which is all that matters for computing dot products.

Step 3: Apply temperature scaling to attention scores

After computing the raw attention scores from the rotated query and key vectors, apply the temperature correction before softmax:

\text{score}(m, n) = \sqrt{t} \cdot \frac{\mathbf{q}_m^\top \mathbf{k}_n}{\sqrt{d_k}}

where:

$\text{score}(m, n)$ : the temperature-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 (the standard $\sqrt{d_k}$ prevents dot products from growing too large)
$t = 0.1 \ln(s) + 1$ : the temperature factor, computed once from the extension scale
$\sqrt{t}$ : the multiplier applied to raw scores, increasing the spread of scores to counteract entropy inflation
This scaling occurs before the softmax normalization, not after

In[19]:

Code

class YaRNRoPE:
    """YaRN-adjusted Rotary Position Embedding."""

    def __init__(
        self,
        d_model,
        original_context=2048,
        scale=4.0,
        base=10000,
        alpha=1.0,
        beta=32.0,
    ):
        """Initialize YaRN RoPE.

        Args:
            d_model: Embedding dimension (must be even)
            original_context: Original training context length
            scale: Context extension factor
            base: RoPE frequency base
            alpha: Lower wavelength threshold
            beta: Upper wavelength threshold
        """
        self.d_model = d_model
        self.original_context = original_context
        self.scale = scale
        self.alpha = alpha
        self.beta = beta

        # Compute adjusted frequencies
        self.original_freqs = compute_rope_frequencies(d_model, base)
        self.wavelengths = compute_wavelengths(d_model, base)

        self.gamma_values = np.array(
            [
                yarn_gamma(w, original_context, scale, alpha, beta)
                for w in self.wavelengths
            ]
        )

        self.adjusted_freqs = self.original_freqs * self.gamma_values

        # Compute temperature factor
        self.temperature = compute_yarn_temperature(scale)

    def apply(self, x, position):
        """Apply YaRN-adjusted RoPE to a vector.

        Args:
            x: Input vector of shape (d_model,)
            position: Position index

        Returns:
            Rotated vector of shape (d_model,)
        """
        return apply_rope(x, position, self.adjusted_freqs)

    def apply_batch(self, x):
        """Apply YaRN-adjusted RoPE to a batch of vectors.

        Args:
            x: Input tensor of shape (seq_len, d_model)

        Returns:
            Rotated tensor of shape (seq_len, d_model)
        """
        seq_len, d = x.shape
        positions = np.arange(seq_len)

        # Compute all rotation angles
        angles = np.outer(positions, self.adjusted_freqs)

        cos_angles = np.cos(angles)
        sin_angles = np.sin(angles)

        x_pairs = x.reshape(seq_len, -1, 2)

        x_rotated = np.stack(
            [
                x_pairs[:, :, 0] * cos_angles - x_pairs[:, :, 1] * sin_angles,
                x_pairs[:, :, 0] * sin_angles + x_pairs[:, :, 1] * cos_angles,
            ],
            axis=-1,
        )

        return x_rotated.reshape(seq_len, d)

    def scale_attention(self, scores):
        """Apply temperature scaling to attention scores.

        Args:
            scores: Attention scores of shape (seq_len, seq_len)

        Returns:
            Scaled scores
        """
        return scores * np.sqrt(self.temperature)

Let's test the complete YaRN implementation:

In[20]:

Code

# Create YaRN module
yarn = YaRNRoPE(d_model=64, original_context=2048, scale=4.0)

# Test with random input
seq_len = 32
test_input = np.random.randn(seq_len, 64)

# Apply YaRN RoPE
rotated = yarn.apply_batch(test_input)

Out[21]:

Console

YaRN Configuration:
  Original context: 2048
  Extension scale: 4.0x
  Extended context: 8192
  Temperature: 1.1386
  Attention scale: sqrt(t) = 1.0671

Input shape: (32, 64)
Output shape: (32, 64)
Magnitude preserved: True

The implementation confirms that YaRN preserves vector magnitudes (rotation is an isometry) while applying the adjusted frequencies. The extended context (8,192 positions) is 4x larger than the training context (2,048 positions), and the temperature correction of approximately 1.067 will be applied to attention scores before softmax.