Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

YaRN › Attention Temperature Scaling

Michael BrenndoerferUpdated June 28, 20257 min read

Language AI HandbookAll chapters

Part 18 of 67Long Context

Chapter 129YaRN › Attention Temperature Scaling

Reading section 7 of 16 · Attention Temperature Scaling · 7 minSection 7 of 16 · 7 min

Section 7 of 16 · 7 min read

YaRN › Attention Temperature ScalingOne section · Sections in this chapter7 / 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

We've now addressed the frequency problem with the ramp function. But recall that we identified two problems at the start: frequency distortion and entropy shift. Even with perfect frequency adjustments, interpolation changes something fundamental about attention patterns, and the ramp function alone cannot fix it. This is where temperature scaling enters the picture.

To understand why, think about what interpolation does geometrically to the embedding space. When we slow down rotations (by multiplying frequencies by $\gamma < 1$ ), positions that were previously "far apart" in the rotated embedding space now appear "closer." Consider two tokens at positions 0 and 100. In the original RoPE, dimension pair 0 might rotate these two positions to vectors that are 100 radians apart in angle (after wrapping around the circle many times). With 4x interpolation applied to that pair, the same two positions are rotated to vectors only 25 radians apart. This compression happens across all interpolated dimension pairs simultaneously, and its effect accumulates in the dot product between query and key.

This compression has a direct consequence for attention scores. The dot product between query and key vectors measures their similarity. When positions appear closer together in the rotated space, the dot products become more similar across all pairs of positions. The range of attention scores shrinks: scores that would have spanned, say, $[-3, 3]$ before interpolation might now span only $[-1, 1]$ . When this compressed range is passed through softmax, the output distribution becomes flatter.

Recall how softmax works. Given a vector of scores $\mathbf{z}$ :

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_j e^{z_j}}

where $z_i$ is the $i$ -th attention score and the sum runs over all key positions. When input scores are spread out (high variance), the exponential function amplifies the differences, producing a distribution that peaks strongly at the highest score and falls off quickly. When input scores are compressed (low variance), the exponentials are more similar to each other, producing a flatter distribution. This increased uniformity is precisely what we measure as higher entropy, and it means the model is attending to many positions at once rather than focusing on the most relevant ones.

The solution is conceptually simple: if scores are too compressed, stretch them back out before applying softmax. We do this by multiplying the scores by a factor greater than 1. This is equivalent to reducing the "temperature" in the softmax, which sharpens the output distribution. YaRN frames this as multiplying by $\sqrt{t}$ where $t > 1$ . The choice of $\sqrt{t}$ rather than $t$ is deliberate, as we will explain shortly.

The temperature factor $t$ is computed from the extension scale $s$ via a formula derived from empirical fitting across many models and extension factors:

t = 0.1 \cdot \ln(s) + 1

where:

$t$ : the attention temperature scaling factor, always $\geq 1$ for any $s \geq 1$
$s$ : the context extension scale factor (e.g., 4 for 4x extension)
$\ln(s)$ : the natural logarithm of the extension factor
$0.1$ : an empirically determined coefficient that controls how aggressively temperature increases with extension factor
$1$ : the base value ensuring $t = 1$ when $s = 1$ , so no correction is applied when there is no extension

Why does this formula use $\ln(s)$ rather than $s$ directly? This reflects an empirical observation about how entropy changes with extension factor. The entropy shift does not grow linearly as we extend further. Going from 4x to 8x extension doubles the context but adds far less than twice the entropy shift that going from 1x to 4x introduced. The logarithm captures this diminishing-returns relationship: doubling the extension factor adds only $\ln(2) \approx 0.693$ to $\ln(s)$ , contributing just $0.069$ more to $t$ regardless of where we start. This makes the temperature correction well-calibrated across a wide range of extension factors.

Why $\sqrt{t}$ rather than $t$ directly as the score multiplier? This relates to how variance scales under linear transformation. If we multiply a set of scores by a constant $c$ , the variance of the scores is multiplied by $c^2$ . To increase variance by a factor of $t$ (undoing the compression from interpolation), we need to multiply scores by $\sqrt{t}$ , because $(\sqrt{t})^2 = t$ . The choice of $\sqrt{t}$ ensures that the temperature correction acts on the variance of attention scores, not their absolute magnitude, which is the right quantity to adjust for matching entropy.

In[14]:

Code

def compute_yarn_temperature(scale):
    """Compute YaRN attention temperature scaling factor."""
    return 0.1 * np.log(scale) + 1.0


# Compute temperature for various extension factors
extension_factors = [2, 4, 8, 16, 32, 64, 128]
temperatures = [compute_yarn_temperature(s) for s in extension_factors]

Out[15]:

Console

YaRN temperature scaling factors:
Extension Factor     Temperature (t)      sqrt(t)             
------------------------------------------------------------
2                    1.0693               1.0341
4                    1.1386               1.0671
8                    1.2079               1.0991
16                   1.2773               1.1302
32                   1.3466               1.1604
64                   1.4159               1.1899
128                  1.4852               1.2187

The $\sqrt{t}$ column shows the actual multiplier applied to attention scores before softmax. Even for very large extension factors (128x), the multiplier stays below 1.2, indicating that temperature correction is a subtle adjustment rather than a dramatic rescaling. For a 4x extension, attention scores are scaled by approximately 1.067, a barely perceptible change in absolute terms but a meaningful correction for the entropy shift introduced by interpolation.

The temperature factor grows slowly with extension factor. For a 4x extension, we multiply attention scores by $\sqrt{1.139} \approx 1.067$ . For 32x extension, the multiplier is $\sqrt{1.347} \approx 1.161$ . These modest corrections help maintain attention sharpness across a wide range of extension scenarios. The key insight is that even small corrections to the pre-softmax scores can have significant effects on the output distribution, because the exponential in softmax amplifies differences.

Let's verify the effect of temperature scaling on entropy:

In[16]:

Code

def apply_yarn_rope(
    x, position, frequencies, scale, context_length, alpha=1.0, beta=32.0
):
    """Apply YaRN-adjusted RoPE."""
    d = len(x)
    wavelengths = compute_wavelengths(d)

    # Compute gamma for each dimension pair
    gamma_vals = np.array(
        [yarn_gamma(w, context_length, scale, alpha, beta) for w in wavelengths]
    )

    # Adjust frequencies
    adjusted_freqs = frequencies * gamma_vals

    return apply_rope(x, position, adjusted_freqs)


def compute_yarn_attention_scores(
    seq_len, d_model, scale, context_length, apply_temperature=True, seed=42
):
    """Compute attention scores with YaRN adjustments."""
    rng = np.random.default_rng(seed)

    original_freqs = compute_rope_frequencies(d_model)

    Q = rng.standard_normal((seq_len, d_model))
    K = rng.standard_normal((seq_len, d_model))

    Q_rotated = np.array(
        [
            apply_yarn_rope(Q[i], i, original_freqs, scale, context_length)
            for i in range(seq_len)
        ]
    )
    K_rotated = np.array(
        [
            apply_yarn_rope(K[i], i, original_freqs, scale, context_length)
            for i in range(seq_len)
        ]
    )

    scores = Q_rotated @ K_rotated.T / np.sqrt(d_model)

    if apply_temperature:
        t = compute_yarn_temperature(scale)
        scores = scores * np.sqrt(t)

    return scores

In[17]:

Code

# Compare entropy with and without temperature scaling
seq_len = 64
d_model = 64
scale = 4.0
context_length = 2048

# YaRN without temperature
yarn_scores_no_temp = compute_yarn_attention_scores(
    seq_len, d_model, scale, context_length, apply_temperature=False
)
yarn_weights_no_temp = softmax(yarn_scores_no_temp)
yarn_entropy_no_temp = compute_entropy(yarn_weights_no_temp)

# YaRN with temperature
yarn_scores_with_temp = compute_yarn_attention_scores(
    seq_len, d_model, scale, context_length, apply_temperature=True
)
yarn_weights_with_temp = softmax(yarn_scores_with_temp)
yarn_entropy_with_temp = compute_entropy(yarn_weights_with_temp)

Out[18]:

Console

Effect of YaRN temperature scaling on attention entropy:
Method                         Mean Entropy    Std Dev        
------------------------------------------------------------
Original RoPE                  3.6643           0.1826
YaRN (no temperature)          3.6643           0.1825
YaRN (with temperature)        3.5992           0.2102

Temperature scaling moves the entropy closer to the original distribution. While the match is not perfect (which is expected given the fundamental geometric changes to the attention space that interpolation introduces), the correction prevents the excessive entropy increase that would otherwise occur. In practice, the remaining difference is small enough that fine-tuning on long-context data can close the gap entirely, which is why YaRN requires so few fine-tuning steps compared to training from scratch.