NTK-aware Scaling › Attention Score Analysis
Language AI HandbookAll chapters NTK-aware Scaling › Attention Score AnalysisOne section · Sections in this chapter7 / 14 ↓
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
To understand the practical impact, let's examine how attention scores behave under different scaling methods. Frequency comparisons tell us about the mechanics, but attention scores tell us about the downstream effect on the model's ability to focus on relevant tokens. We'll create synthetic queries and keys at various distances and measure how the attention weight between them changes with distance under each method.
The expected behavior is a gradual decay: tokens closer to the query should receive higher attention than tokens far away, all else being equal. This is because RoPE encodes relative position directly into the inner product, and vectors at small relative offsets have more similar rotations than vectors at large relative offsets. If context extension distorts this decay pattern, the model may inappropriately up-weight distant tokens or fail to distinguish near from far.
def compute_attention_scores(q, k):
"""Compute scaled dot-product attention scores."""
d = q.shape[-1]
scores = q @ k.T / np.sqrt(d)
return scores
def create_position_aware_vectors(positions, d, frequencies):
"""Create vectors with RoPE applied at specified positions."""
n = len(positions)
# The shared setup seed keeps this example deterministic across variants.
base_vectors = np.random.randn(n, d) * 0.1
# Create rotation matrices for each position
rotated = np.zeros_like(base_vectors)
for idx, pos in enumerate(positions):
angles = pos * frequencies
cos_a = np.cos(angles)
sin_a = np.sin(angles)
for i in range(d // 2):
x, y = base_vectors[idx, 2 * i], base_vectors[idx, 2 * i + 1]
rotated[idx, 2 * i] = x * cos_a[i] - y * sin_a[i]
rotated[idx, 2 * i + 1] = x * sin_a[i] + y * cos_a[i]
return rotated# Examine attention between a query and keys at varying distances
d = 64
base = 10000
train_length = 4096
extended_length = 32768
scale = extended_length / train_length
# Different scaling approaches
freq_orig = compute_rope_frequencies(d, base)
freq_pi = position_interpolation_frequencies(d, base, scale)
freq_ntk = ntk_aware_frequencies(d, base, scale)
# Query at position 1000, keys at various distances
query_pos = 1000
key_distances = [1, 2, 5, 10, 50, 100, 500, 1000]
key_positions = [query_pos + dist for dist in key_distances]
The attention decay pattern reveals a subtle but important distinction. While Position Interpolation flattens the attention curve (making nearby and distant tokens more similar in score), NTK-aware scaling preserves more of the original decay structure, especially at short distances.
The most significant differences appear at very short distances, within a few positions of the query. This is where the high-frequency dimensions dominate the inner product, because they are the ones that change fastest with relative position. Since NTK-aware scaling preserves these high-frequency dimensions, the attention behavior at short distances stays closer to the original. At long distances, both methods produce similar behavior because the long-range structure is dominated by low-frequency dimensions, which both methods compress to a similar degree.