NTK-aware Scaling › Implementation
Language AI HandbookAll chapters NTK-aware Scaling › ImplementationOne section · Sections in this chapter4 / 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
With the formula derived and verified, let's translate it into code. We'll build the implementation in layers, starting with the core frequency computation and progressively adding the full RoPE transformation. The implementation is remarkably simple, which is one of the key practical advantages of NTK-aware scaling over more complex methods. The entire context extension reduces to a single change in one parameter.
In practice, you would integrate this into a model's RoPE layer by patching the frequency computation before inference. In frameworks like Hugging Face Transformers, this means replacing the inv_freq buffer inside the rotary embedding module. Libraries like llama.cpp and vLLM implement NTK-aware scaling as a configurable option that can be toggled with a single flag.
Computing Frequencies
The foundation is a function that computes RoPE frequencies for any base value. This same function serves both the original frequencies and NTK-aware frequencies, since the only difference is the base we pass in. By keeping the function generic, we avoid code duplication and make it easy to swap between methods:
def compute_rope_frequencies(d, base=10000):
"""Compute original RoPE frequencies for d/2 dimension pairs."""
dim_pairs = np.arange(d // 2)
frequencies = 1.0 / (base ** (2 * dim_pairs / d))
return frequenciesPosition Interpolation simply divides all frequencies by the scale factor:
def position_interpolation_frequencies(d, base=10000, scale=1.0):
"""Compute Position Interpolation frequencies (uniform scaling)."""
original = compute_rope_frequencies(d, base)
return original / scalePosition Interpolation is the simplest possible extension: just divide every frequency by the scale factor. This is a useful baseline that is easy to understand and easy to implement, but as we have seen, it does not handle all frequency bands well.
NTK-aware scaling modifies the base according to our derived formula, then computes frequencies using the modified base:
def ntk_aware_frequencies(d, base=10000, scale=1.0):
"""Compute NTK-aware scaled frequencies (non-uniform scaling)."""
# Modify the base according to NTK formula: b' = b * s^(d/(d-2))
ntk_base = base * (scale ** (d / (d - 2)))
return compute_rope_frequencies(d, ntk_base)The elegance of this implementation is worth appreciating: the entire NTK-aware scaling logic is a one-liner. Computing the modified base and passing it to the same frequency function as before handles everything. This simplicity is a major practical advantage: you can retrofit NTK-aware scaling onto any existing RoPE implementation with a single line change.
Comparing the Approaches
Let's compute frequencies for all three methods and compare them side by side. Looking at the raw numbers helps build intuition before we see the visualizations:
# Compare the two approaches
d = 64
base = 10000
scale = 8 # Extend context by 8x
freq_original = compute_rope_frequencies(d, base)
freq_pi = position_interpolation_frequencies(d, base, scale)
freq_ntk = ntk_aware_frequencies(d, base, scale)Context extension factor: 8x Embedding dimension: 64 Frequency comparison (first 5 dimension pairs): Dim Original PI NTK PI ratio NTK ratio -------------------------------------------------------------- 0 1.000000 0.125000 1.000000 8.00 1.00 1 0.749894 0.093737 0.701242 8.00 1.07 2 0.562341 0.070293 0.491741 8.00 1.14 3 0.421697 0.052712 0.344829 8.00 1.22 4 0.316228 0.039528 0.241809 8.00 1.31 Frequency comparison (last 5 dimension pairs): Dim Original PI NTK PI ratio NTK ratio -------------------------------------------------------------- 27 0.000422 0.000053 0.000069 8.00 6.12 28 0.000316 0.000040 0.000048 8.00 6.54 29 0.000237 0.000030 0.000034 8.00 7.00 30 0.000178 0.000022 0.000024 8.00 7.48 31 0.000133 0.000017 0.000017 8.00 8.00
The numbers confirm what we derived mathematically. Look at the "PI ratio" column: it's a constant 8.00 across all dimensions. This reflects Position Interpolation's uniform scaling. Now look at "NTK ratio": it starts near 1.00 for dimension 0 and gradually increases toward 8.00 for the highest dimensions. This is the dimension-dependent scaling in action.
The NTK ratio column tells you something important about what the model will experience at inference time. For dimension 0, the frequencies are essentially unchanged, so the model's learned patterns for detecting adjacent tokens are fully intact. For the highest dimensions, the frequencies are reduced by the full factor of 8, which is the necessary price of extending the context: you need the slow-rotating dimensions to cover 8x more positions, so they must rotate 8x more slowly. The difference from Position Interpolation is that this cost is paid entirely by the dimensions that are designed to handle long-range structure, not by the dimensions that handle local patterns.
Visualizing the Frequency Spectrum
A log-scale plot reveals the full picture across all 32 dimension pairs. The logarithmic scale is important here because the frequencies themselves span many orders of magnitude:
On the log scale, observe how NTK-aware frequencies (triangles) track the original (circles) closely for low dimension indices, where high frequencies live. As we move right toward higher dimension indices (lower frequencies), the NTK curve gradually diverges to match Position Interpolation (squares). This is exactly the "preserve high frequencies, scale low frequencies" behavior we designed.
The crossover point, where NTK-aware frequencies visibly diverge from the original, corresponds roughly to the middle of the dimension range. In a model with head dimension, this is around dimension pair 15 to 20. Dimensions below this threshold receive essentially no scaling and behave as if the context window were not extended. Dimensions above this threshold receive progressively more scaling and are responsible for encoding the longer-range positional structure needed for the extended context.
Rotation Angle Heatmaps
To visualize how different scaling methods affect the rotation patterns, let's create heatmaps showing the rotation angles across positions and dimension pairs. In these heatmaps, each row corresponds to a dimension pair and each column to a position. The color represents the rotation angle modulo , so rapid color cycling indicates a high-frequency dimension while slow color change indicates a low-frequency one:
The heatmaps reveal the key difference between methods. In the original (left), the bottom rows (high-frequency dimensions) show rapid color cycling as position increases, while top rows (low-frequency dimensions) change slowly. Position Interpolation (center) uniformly slows down all cycling, making all rows look more like the slow top rows. NTK-aware scaling (right) preserves the rapid cycling in the bottom rows while slowing only the top rows. This visual confirms that NTK-aware scaling selectively modifies frequencies based on dimension.
Looking at the NTK heatmap versus the original, a key observation is that the bottom half of the rows looks nearly identical in both. This is the visual representation of "high frequencies preserved." The top half of the rows is where you see the difference, and even there the NTK-aware version looks more similar to the original than the Position Interpolation version does. The original's top rows complete a fraction of a full rotation over the 20 positions shown, and the NTK-aware version matches this better than Position Interpolation, which further compresses them.