NTK-aware Scaling › From Intuition to Formula: The NTK-aware Approach

Now that we've seen the problem with uniform scaling, let's develop a solution from first principles. The journey from intuition to a working formula requires answering three questions: What property do we want? How can we achieve it mathematically? And does the result match our expectations?

The key realization is that not all frequencies are equal in importance. Think of RoPE's multi-frequency design like a ruler with both centimeter and millimeter markings. The millimeter marks (high frequencies) give you precision for small measurements, while the centimeter marks (low frequencies) help you measure longer distances quickly. Position Interpolation is like shrinking the entire ruler by 8x, making even the millimeter marks too small to read. What we really want is to keep the millimeter precision while adjusting the centimeter scale.

Another useful analogy: think of the frequency dimensions as instruments in an orchestra. The high-frequency dimensions are like the piccolo and violins, playing rapid, precise notes that define the immediate rhythmic detail of the music. The low-frequency dimensions are like the bass and cello. This provides the slow harmonic foundation that gives the piece its overall structure. Position Interpolation is like telling every instrument to play at 1/8 tempo, which ruins the fast parts even though the slow parts could tolerate it. A better approach would be to slow only the bass section while keeping the violins at full tempo.

Translating this intuition into concrete requirements:

1. Preserve high frequencies: The fastest-rotating dimension pairs (small $i$) should remain unchanged. These are our "millimeter marks" for local position distinctions.

2. Scale low frequencies: The slowest-rotating dimension pairs (large $i$) can be compressed by the full factor $s$. These are our "centimeter marks" that need adjustment for longer contexts.

3. Smooth transition: Intermediate dimensions should scale gradually between these extremes, not abruptly jump.

These three requirements are not arbitrary design choices. They follow from thinking about what the model learned during training. For a model trained on sequences up to $L_{\text{train}}$ tokens, every positional relationship it has seen involves positions within that range. High-frequency dimensions complete many full rotations within that range and therefore carry meaningful, learned patterns. Low-frequency dimensions may only complete a fraction of a rotation, so extending the context just means continuing the pattern they were already set up to handle.

How can we achieve dimension-dependent scaling? The original RoPE frequency formula is $\theta_i = 1/b^{2i/d}$. Position Interpolation modifies the output by dividing: $\theta_i' = \theta_i / s$. This gives uniform scaling because every frequency is divided by the same constant.

Instead, we'll modify the input: the base $b$. This is the key creative step in the derivation. By operating on the base rather than the output, we can exploit the mathematical structure of the exponential to get dimension-dependent effects automatically.

If we replace $b$ with a larger base $b'$, all frequencies decrease (slower rotation). The key insight is that the exponent $2i/d$ causes this decrease to affect different dimensions differently:

• When $i = 0$: $\theta_0' = 1/(b')^0 = 1$. No matter what $b'$ is, the highest frequency remains 1.
• When $i$ is large: $\theta_i' = 1/(b')^{2i/d}$ shrinks more because the exponent is larger.

This is exactly the property we want. By choosing the right $b'$, we can leave high frequencies untouched while scaling low frequencies by whatever factor we need. The mathematical structure of the exponential gives us the smooth, dimension-dependent transition between these extremes as a natural consequence, without any ad hoc weighting or interpolation scheme.

Notice the elegance here: we do not need to specify how each individual dimension should be scaled. We just need to find one number, $b'$, and the exponential structure of the formula takes care of distributing the scaling correctly across all dimensions. This is a beautiful example of how choosing the right mathematical representation of a problem can make the solution fall out almost automatically.

This geometric view makes the difference concrete. In the high-frequency dimension (left), original RoPE spreads positions around the circle, allowing the model to distinguish them easily. Position Interpolation bunches them together near the starting point, reducing distinguishability. NTK-aware scaling preserves the original spread. In the low-frequency dimension (right), all methods behave similarly since both Position Interpolation and NTK-aware apply full compression to these slow-rotating dimensions.

Let's work backward from our requirements to find the exact formula for $b'$. We want a scaling function where:

• $\text{scale}_0 = 1$ (no scaling at highest frequency)
• $\text{scale}_{d/2-1} = s$ (full scaling at lowest frequency)

We write the modified base as $b' = b \cdot \alpha$ for some multiplier $\alpha$ to be determined. This parameterization is useful because it lets us express the new base as a proportional increase over the old one, and we can reason about the multiplier $\alpha$ independently. Substituting into the frequency formula:

Using the exponent rule $(ab)^n = a^n \cdot b^n$, we can separate this:

The effective scaling factor for dimension $i$ becomes:

Now we apply our constraints. At $i = 0$, we need $\text{scale}_0 = 1$:

This is automatically satisfied for any $\alpha > 0$, since any number raised to zero equals one. Our first requirement is met regardless of what $\alpha$ we choose.

At $i = d/2 - 1$ (the last dimension pair), we need $\text{scale}_{d/2-1} = s$:

Simplifying the exponent step by step:

To solve for $\alpha$, we raise both sides to the power $d/(d-2)$:

This gives us the complete NTK-aware formula:

where:

• $b'$: the modified base for NTK-aware scaling
• $b$: the original base (typically 10000)
• $s$: the context extension factor (e.g., $s = 8$ to extend from 4k to 32k tokens)
• $d$: the embedding dimension
• $d/(d-2)$: an exponent slightly greater than 1 (for $d = 64$, this equals $64/62 \approx 1.032$)

The exponent $d/(d-2)$ is worth examining closely. For large $d$, this converges toward 1, meaning the formula approaches $b' = b \cdot s$. For small $d$, it is significantly larger than 1. In practice, the head dimension $d$ used in modern transformers is typically 64 or 128, so the exponent is either $\approx 1.032$ or $\approx 1.016$. These are small but meaningful corrections that ensure the boundary condition at $i = d/2 - 1$ is met exactly. If you were to mistakenly use $b' = b \cdot s$ instead of the correct formula, the lowest-frequency dimension would be scaled by slightly less than $s$, producing a small but systematic mismatch at the low-frequency end of the spectrum.

Substituting the new base into the frequency formula reveals how each dimension is affected. Starting from the frequency definition, replacing $b$ with $b \cdot s^{d/(d-2)}$, and simplifying:

The effective scaling factor for each dimension, defined as the ratio of the original frequency to the new frequency, is:

Let's verify this matches our design goals:

• When $i = 0$ (highest frequency): $\text{scale}_0 = s^0 = 1$. The highest frequencies are not scaled at all. ✓
• When $i = d/2 - 1$ (lowest frequency): $\text{scale}_{d/2-1} = s^{(d-2)/(d-2)} = s$. The lowest frequencies are scaled by the full factor. ✓
• Intermediate dimensions: Scaling increases smoothly from 1 to $s$ as $i$ increases. ✓

This is precisely what we wanted: high frequencies preserved, low frequencies scaled, smooth transition between.

Notice also that the scaling exponent $2i/(d-2)$ grows roughly linearly with dimension index $i$. This means the scaling factor is exponential in the dimension index, which mirrors the original exponential spacing of the frequencies. In other words, NTK-aware scaling respects the logarithmic structure of RoPE. It does not arbitrarily compress some frequencies more than others; it applies compression in a way that is proportional to the frequency's original "slowness."

Let's make this concrete with typical values: $d = 64$ dimensions, $s = 8$ extension factor (4k to 32k tokens). This is a realistic scenario corresponding to extending LLaMA-2, which was trained with a 4,096-token context, to a 32,768-token context window.

First, compute the NTK exponent:

The exponent is only slightly above 1, which means the modified base is only slightly larger than what you would get by simply multiplying $b$ by $s$. This small correction is what makes the last dimension pair satisfy the boundary condition exactly.

The modified base becomes:

In practice, this means the model will behave as if it was trained with a base of 87,200 rather than 10,000. You can think of this as a kind of "stretching" of the frequency spectrum: by raising the base, we make all frequencies slightly smaller (slower rotation), but the effect is much stronger at the low-frequency end because those dimensions have a large exponent in the denominator.

Now we can compute scaling factors for specific dimensions:

The gradient is smooth and continuous. High-frequency pairs (small $i$) stay close to their original values, while low-frequency pairs (large $i$) are compressed to fit the extended context. Notice that the "moderate scaling" region, around dimension pairs 12 to 20, corresponds to the intermediate frequencies that provide mid-range positional encoding, useful for tracking structure at the paragraph level rather than the word or document level. NTK-aware scaling applies intermediate compression here, which is intuitively the right thing to do: these dimensions can tolerate some compression, but not the full factor.