NTK-aware Scaling › Applying RoPE with NTK-aware Scaling
Language AI HandbookAll chapters NTK-aware Scaling › Applying RoPE with NTK-aware ScalingOne section · Sections in this chapter5 / 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
Now that we can compute the frequencies, let's implement the complete RoPE transformation and measure its effect on token distinguishability. This section is where theory meets practice: we apply the frequencies to actual vectors and observe whether NTK-aware scaling delivers on its promise of better local discriminability.
The RoPE Transformation
RoPE works by rotating each pair of embedding dimensions by a position-dependent angle. The rotation angle for position and dimension pair is simply , where is the frequency we computed above. This is a standard 2D rotation matrix applied to consecutive pairs of dimensions. Importantly, the rotation is applied to both the query and key vectors at each position, so that the inner product between a query at position and a key at position depends only on the relative offset through the rotation difference:
def apply_rope(x, frequencies):
"""Apply RoPE to a sequence of vectors.
Args:
x: Input vectors of shape (seq_len, d)
frequencies: Base frequencies for each dimension pair, shape (d/2,)
Returns:
Rotated vectors of same shape as x
"""
seq_len, d = x.shape
positions = np.arange(seq_len)
# Compute rotation angles: (seq_len, d/2)
# Each position m gets angle m * theta_i for each dimension pair i
angles = np.outer(positions, frequencies)
# Split input into pairs for 2D rotation
x_pairs = x.reshape(seq_len, -1, 2) # (seq_len, d/2, 2)
# Precompute sin and cos for efficiency
cos_angles = np.cos(angles) # (seq_len, d/2)
sin_angles = np.sin(angles)
# Apply 2D rotation: [x, y] -> [x*cos - y*sin, x*sin + y*cos]
x_rotated = np.zeros_like(x_pairs)
x_rotated[:, :, 0] = (
x_pairs[:, :, 0] * cos_angles - x_pairs[:, :, 1] * sin_angles
)
x_rotated[:, :, 1] = (
x_pairs[:, :, 0] * sin_angles + x_pairs[:, :, 1] * cos_angles
)
return x_rotated.reshape(seq_len, d)Measuring Token Distinguishability
The ultimate test of our scaling method is whether adjacent tokens remain distinguishable after rotation. If tokens at positions and become too similar, the model loses its ability to attend precisely to nearby positions.
Think of it this way: the attention mechanism computes scores by taking inner products between queries and keys. If all nearby keys produce nearly the same score for a given query, the softmax distribution becomes flat over those positions. The model cannot selectively attend to "the word right before me" versus "the word two steps back," because both produce the same attention score. This is the practical consequence of high-frequency compression, and it is what we want to measure.
We'll measure this using cosine similarity between adjacent token embeddings. Lower similarity means more distinguishable tokens. We use random embeddings here as a proxy for actual model outputs, which is valid because the distinguishability depends primarily on the rotations applied, not on the specific content of the vectors:
# Create one deterministic sample after the shared setup seed.
seq_len = 10
d = 64
embeddings = np.random.randn(seq_len, d) * 0.1
# Apply RoPE with different scaling methods
rotated_original = apply_rope(embeddings, freq_original)
rotated_pi = apply_rope(embeddings, freq_pi)
rotated_ntk = apply_rope(embeddings, freq_ntk)# Compute cosine similarity between adjacent tokens
def cosine_similarity(a, b):
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
# Calculate similarities for each method
similarities = {"original": [], "pi": [], "ntk": []}
for i in range(seq_len - 1):
similarities["original"].append(
cosine_similarity(rotated_original[i], rotated_original[i + 1])
)
similarities["pi"].append(
cosine_similarity(rotated_pi[i], rotated_pi[i + 1])
)
similarities["ntk"].append(
cosine_similarity(rotated_ntk[i], rotated_ntk[i + 1])
)Cosine similarity between adjacent token embeddings:
Positions Original PI NTK
--------------------------------------------
0 & 1 -0.0311 0.0152 -0.0237
1 & 2 -0.2305 -0.2127 -0.2279
2 & 3 -0.0370 -0.0254 -0.0345
3 & 4 0.2245 0.2120 0.2135
4 & 5 -0.0194 0.0068 -0.0210
5 & 6 -0.0112 0.0196 -0.0044
6 & 7 -0.1081 -0.0881 -0.1076
7 & 8 0.1659 0.1631 0.1591
8 & 9 0.0982 0.1182 0.1049
--------------------------------------------
Average 0.0057 0.0232 0.0065
The visualization makes the pattern clear: Position Interpolation significantly increases similarity between adjacent tokens (making them harder to distinguish), while NTK-aware scaling keeps the similarity close to the original. This preservation of token distinguishability is the practical benefit of frequency-dependent scaling.
In practice, the gap between NTK-aware and Position Interpolation is most pronounced in tasks where local syntactic structure matters. Question answering, where the model must identify a precise span of text, benefits from the preserved local discriminability. Summarization, where the model reads long text and produces a concise summary, is somewhat less affected because the model mostly needs to understand what is there rather than to precisely track where it is. This is an important calibration point when choosing a context extension method for your specific use case.