YaRN › Attention Entropy and the Temperature Problem
Language AI HandbookAll chapters YaRN › Attention Entropy and the Temperature ProblemOne section · Sections in this chapter3 / 16 ↓
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
Beyond frequency adjustments, there is a more subtle issue that neither PI nor NTK-aware scaling addresses: attention entropy shifts. When we modify RoPE frequencies, we change the distribution of attention scores in ways that affect model behavior in a systematic and problematic direction. Understanding this problem requires thinking carefully about what attention score distributions represent and why their entropy matters.
Recall that attention scores are computed as the scaled dot product between query and key vectors. These scores are then passed through a softmax to produce attention weights that sum to one. The resulting distribution tells the model how much to focus on each key position. A low-entropy distribution is sharply peaked, meaning the model strongly attends to one or a few positions. A high-entropy distribution is flat, meaning the model spreads its attention diffusely. During training, models learn attention patterns calibrated to the natural entropy of the position encoding scheme they were trained with. If you change that entropy, you change the effective behavior of every attention head in every layer.
Recall that attention scores are computed as the scaled dot product between query and key vectors:
where:
- : the attention score between a query at position and a key at position
- : the RoPE-rotated query vector at sequence position
- : the RoPE-rotated key vector at sequence position
- : the dimension of the key vectors, used for scaling to prevent dot products from growing too large in magnitude
- : the dot product between query and key, measuring their geometric similarity after rotation
When we interpolate positions, the effective distances between tokens change in the rotated embedding space. Positions that were far apart now appear closer in that space. This compression happens because we have slowed down the rotation rates, so two tokens that used to be separated by many radians of rotation are now separated by fewer. This reduction in angular separation increases the cosine similarity between their rotated representations, which increases the dot product, which compresses the range of attention scores.
Higher score compression means higher entropy: the model attends more diffusely rather than focusing on specific positions. The softmax function is particularly sensitive to the scale of its inputs. When inputs are spread out, the largest values dominate strongly and the output is peaked. When inputs are compressed, the softmax output becomes more uniform. This is the same phenomenon that temperature scaling in language model generation exploits, and it is exactly what YaRN needs to counteract.
The entropy of an attention distribution measures how spread out the attention weights are across sequence positions. Low entropy means the model focuses sharply on a few positions, which is typical of attention heads specialized for specific patterns (e.g., a head that always attends to the previous token, or one that retrieves the most semantically similar prior context). High entropy means attention is distributed more evenly across many positions, which is useful for aggregating broad context but loses the ability to pinpoint specific information. Changes to position encoding can inadvertently shift this entropy, degrading the model's ability to both focus (for retrieval) and aggregate (for synthesis).
The key insight is that this entropy problem is not fixed by improving the frequency adjustment alone. Even a perfect frequency adjustment that places every dimension pair's rotation frequency exactly where we would want it will still change the geometry of the rotated embedding space. The act of enabling longer contexts changes the statistics of attention, and YaRN compensates with a direct correction to attention scores.
Let's quantify this effect by computing attention entropy under different interpolation schemes:
def softmax(x, axis=-1):
"""Numerically stable softmax."""
x_max = np.max(x, axis=axis, keepdims=True)
exp_x = np.exp(x - x_max)
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
def compute_entropy(attention_weights):
"""Compute entropy of attention distribution."""
# Avoid log(0) by adding small epsilon
eps = 1e-10
log_weights = np.log(attention_weights + eps)
entropy = -np.sum(attention_weights * log_weights, axis=-1)
return entropy
def apply_rope(x, position, frequencies):
"""Apply RoPE to a vector at given position."""
d = len(x)
x_pairs = x.reshape(-1, 2)
angles = position * frequencies
cos_angles = np.cos(angles)
sin_angles = np.sin(angles)
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.flatten()
def compute_attention_scores(seq_len, d_model, frequencies, seed=42):
"""Compute attention score matrix for random Q, K vectors."""
rng = np.random.default_rng(seed)
# Generate random query and key vectors
Q = rng.standard_normal((seq_len, d_model))
K = rng.standard_normal((seq_len, d_model))
# Apply RoPE
Q_rotated = np.array(
[apply_rope(Q[i], i, frequencies) for i in range(seq_len)]
)
K_rotated = np.array(
[apply_rope(K[i], i, frequencies) for i in range(seq_len)]
)
# Compute scaled dot-product scores
scores = Q_rotated @ K_rotated.T / np.sqrt(d_model)
return scores# Compare entropy across methods
seq_len = 64
d_model = 64
scale = 4.0
# Original context (positions 0 to seq_len-1)
original_scores = compute_attention_scores(seq_len, d_model, original_freqs)
original_weights = softmax(original_scores)
original_entropy = compute_entropy(original_weights)
# Position interpolation (pretend we're at 4x length, but scale down)
pi_scores = compute_attention_scores(seq_len, d_model, pi_freqs)
pi_weights = softmax(pi_scores)
pi_entropy = compute_entropy(pi_weights)
# NTK-aware
ntk_scores = compute_attention_scores(seq_len, d_model, ntk_freqs)
ntk_weights = softmax(ntk_scores)
ntk_entropy = compute_entropy(ntk_weights)Mean attention entropy across query positions: Method Mean Entropy Std Dev ------------------------------------------------------- Original RoPE 3.6643 0.1826 Position Interpolation 3.6601 0.1594 NTK-aware 3.6838 0.1464
The entropy values reveal the effect of position interpolation on attention patterns. Both PI and NTK-aware scaling increase entropy compared to the original RoPE. While the differences might seem modest in absolute terms, they compound across layers and affect the model's ability to retrieve and aggregate information from specific positions. A model that learned to perform sharp, focused retrieval with original RoPE will now perform softer, more diffuse retrieval after extension, degrading quality on tasks that require pinpoint attention. YaRN addresses this by introducing an attention temperature correction that pushes the entropy back toward the original distribution.