NTK-aware Scaling: Extending Context Length in LLMs

Position Interpolation, as we saw in the previous chapter, extends context length by simply scaling down all rotation frequencies. While this works, it comes with a hidden cost: by compressing all frequencies equally, we lose the fine-grained positional distinctions that high-frequency components provide. Consider a model trained with 4,096 tokens. Position Interpolation at 8x scale compresses everything to fit 32,768 positions, but now adjacent tokens that were once easily distinguishable become nearly identical in their high-frequency dimensions.

NTK-aware scaling takes a more surgical approach. Instead of treating all frequencies equally, it recognizes that different frequency bands serve different purposes. High frequencies distinguish nearby tokens, while low frequencies capture long-range structure. By scaling frequencies non-uniformly, NTK-aware methods preserve local precision while still extending the effective context window.

This chapter develops the intuition behind NTK-aware scaling, derives the mathematical formula, and implements both static and dynamic variants. By the end, you'll understand why this approach often outperforms linear interpolation, especially for tasks requiring precise local attention.

To understand why uniform scaling falls short, we need to revisit how RoPE frequencies work. Recall that RoPE assigns each dimension pair $i$ a base frequency that determines how fast the rotation angle changes with position:

where:

• $\theta_i$: the base rotation frequency (in radians per position) for dimension pair $i$
• $b$: the base constant, typically 10000, which controls the overall frequency range
• $i$: the dimension pair index, ranging from 0 to $d/2 - 1$
• $d$: the total embedding dimension (must be even since we work in pairs)
• $2i/d$: the exponent that creates a geometric progression of frequencies across dimensions

The first dimension pair ($i = 0$) has frequency $\theta_0 = 1/10000^0 = 1$, meaning it rotates by 1 radian per position and completes a full rotation every $2\pi \approx 6.28$ positions. The last pair ($i = d/2 - 1$) has a much smaller frequency, completing a rotation only after thousands of positions. This spread creates a multi-scale representation:

• High-frequency pairs ($i$ near 0): Distinguish nearby tokens with precision. Tokens at positions 5 and 6 look very different.
• Low-frequency pairs ($i$ near $d/2$): Capture coarse position information. Tokens far apart have noticeably different rotations.

Position Interpolation addresses context extension by scaling all frequencies uniformly:

where:

• $\theta_i'$: the scaled frequency for dimension pair $i$ after Position Interpolation
• $\theta_i$: the original frequency for dimension pair $i$
• $s$: the context extension factor, computed as $s = L_{\text{target}} / L_{\text{train}}$ (target length divided by training length)

If we want to extend from 4,096 to 32,768 tokens, $s = 32768 / 4096 = 8$. Every frequency gets divided by 8. The high-frequency pair that previously rotated 1 radian per position now rotates only $1/8 = 0.125$ radians. Adjacent tokens, which were once clearly distinguishable, become nearly indistinguishable in these dimensions.

Let's visualize this problem:

The compression is dramatic. With Position Interpolation at 8x scale, adjacent tokens are separated by only 0.125 radians in the highest-frequency dimension, compared to 1 radian originally. This means the model has far less "room" to distinguish nearby positions, potentially harming tasks that require precise local attention.

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.

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.

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$. 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.

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)

Let $b' = b \cdot \alpha$ for some multiplier $\alpha$ we need to determine. 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$)

Substituting the new base into the frequency formula reveals how each dimension is affected:

The effective scaling factor for each dimension 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.

Let's make this concrete with typical values: $d = 64$ dimensions, $s = 8$ extension factor (4k to 32k tokens).

First, compute the NTK exponent:

The modified base becomes:

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.

The visualization confirms our mathematical analysis. The NTK-aware curve (squares) starts at 1 for the highest-frequency pairs and smoothly increases to the full scale factor of 8 for the lowest-frequency pairs. Meanwhile, Position Interpolation (circles) maintains a flat line at 8, treating all dimensions identically.

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 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:

Position Interpolation simply divides all frequencies by the scale factor:

NTK-aware scaling modifies the base according to our derived formula, then computes frequencies using the modified base:

Let's compute frequencies for all three methods and compare them side by side:

The numbers confirm what we derived mathematically. Look at the "PI ratio" column: it's a constant 8.00 across all dimensions, reflecting 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.

A log-scale plot reveals the full picture across all 32 dimension pairs:

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.

To visualize how different scaling methods affect the rotation patterns, let's create heatmaps showing the rotation angles across positions and dimension pairs:

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. 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.

Now that we can compute the frequencies, let's implement the complete RoPE transformation and measure its effect on token distinguishability.

RoPE works by rotating each pair of embedding dimensions by a position-dependent angle. The rotation angle for position $m$ and dimension pair $i$ is simply $m \cdot \theta_i$, where $\theta_i$ is the frequency we computed above:

The ultimate test of our scaling method is whether adjacent tokens remain distinguishable after rotation. If tokens at positions $m$ and $m+1$ become too similar, the model loses its ability to attend precisely to nearby positions.

We'll measure this using cosine similarity between adjacent token embeddings. Lower similarity means more distinguishable tokens:

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.

Static NTK-aware scaling uses a fixed extension factor $s$. But what if the actual sequence length varies? A model configured for 32k tokens shouldn't apply aggressive scaling when processing only 2k tokens. Ideally, the scaling should adapt to the current context: no modification for short sequences, progressive scaling as sequences grow longer.

Dynamic NTK scaling adapts the base in real-time based on the current sequence length:

where:

• $b'(L)$: the dynamically computed base, now a function of sequence length $L$
• $b$: the original base (typically 10000)
• $L$: the current sequence length being processed
• $L_{\text{train}}$: the context length the model was trained with
• $L/L_{\text{train}}$: the effective extension factor, computed on-the-fly
• $d/(d-2)$: the same exponent as in static NTK-aware scaling

The formula is identical to static NTK scaling, except that $s = L/L_{\text{train}}$ is computed dynamically rather than fixed in advance.

When $L \leq L_{\text{train}}$, the ratio $L/L_{\text{train}} \leq 1$, and we clamp it to 1 (no scaling). When $L$ exceeds the training length, scaling kicks in proportionally. A sequence of 8k tokens (2x the training length of 4k) gets $s = 2$; a sequence of 32k tokens (8x) gets $s = 8$.

The effective base scales smoothly with sequence length. At 32k tokens (8x the training length), the base increases to approximately 87k, matching our earlier static calculation.

To understand the practical impact, let's examine how attention scores behave under different scaling methods. We'll create synthetic queries and keys at various distances and measure the attention patterns:

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.

Let's directly compare how well each method preserves the relative position encoding properties:

The deviation plot (right) makes the advantage clear: NTK-aware scaling maintains smaller deviations from the original behavior, particularly for small relative positions where local attention patterns matter most.

An alternative perspective on NTK-aware scaling is to analyze the interpolation factor for each frequency. We can express NTK-aware scaling as a blend between "no interpolation" and "full interpolation":

This interpolation perspective provides an intuitive understanding: NTK-aware scaling smoothly transitions from "preserve high frequencies" (interpolation factor near 0) to "fully scale low frequencies" (interpolation factor near 1).

NTK-aware scaling works better than Position Interpolation, but it comes with its own trade-offs and limitations.

The core tension in context extension is that you cannot perfectly preserve all the properties of the original position encoding while also extending the context window. NTK-aware scaling prioritizes high-frequency preservation at the cost of more aggressive low-frequency compression. For tasks that rely heavily on precise long-range position information, this trade-off may not be ideal. Document retrieval, where finding a specific passage requires accurate global positioning, might suffer compared to summarization tasks where local coherence matters more.

Additionally, NTK-aware scaling, like Position Interpolation, typically requires fine-tuning to achieve optimal performance. While models may work "out of the box" with NTK scaling applied at inference time, the attention patterns learned during training assumed different frequency relationships. Fine-tuning on longer sequences helps the model adapt to the new frequency landscape. The amount of fine-tuning needed is generally comparable to Position Interpolation: a few hundred to a few thousand steps on representative long-context data often suffices.

The choice between static and dynamic NTK scaling depends on deployment constraints. Static scaling is simpler to implement and has deterministic behavior, but requires knowing the maximum context length in advance. Dynamic scaling adapts gracefully to varying sequence lengths but adds computational overhead, as frequencies must be recomputed based on current sequence length. For most applications where sequence lengths are predictable, static scaling is sufficient.

Finally, NTK-aware scaling is specific to RoPE and doesn't transfer to other position encoding schemes. Models using absolute position embeddings, ALiBi, or other mechanisms require different extension strategies. This limits the generality of the approach, though the dominance of RoPE in modern architectures means NTK-aware scaling is broadly applicable.

When implementing NTK-aware scaling, several parameters control the behavior of the context extension:

• base (default: 10000): The original RoPE base constant. This value comes from the pretrained model and should match what was used during training. Common values are 10000 (LLaMA, Mistral) or 500000 (some newer models with extended context).

• scale (s): The context extension factor, computed as target length divided by training length. For example, extending from 4k to 32k tokens gives $s = 8$. Larger values enable longer contexts but increase the distortion from original frequencies.

• d (embedding dimension): The model's embedding dimension, used in computing the NTK exponent $d/(d-2)$. This is fixed by the model architecture and typically ranges from 64 to 128 for the head dimension in modern transformers.

• train_length (for dynamic scaling): The context length the model was originally trained with. This serves as the threshold below which no scaling is applied. Using the correct value is critical for dynamic NTK to work properly.

When choosing between static and dynamic scaling, consider your deployment scenario. Static scaling with a fixed $s$ is simpler and more predictable, suitable when you know the maximum sequence length in advance. Dynamic scaling adapts to varying inputs but adds computational overhead for recomputing frequencies.

NTK-aware scaling provides a principled approach to extending context length in RoPE-based models by recognizing that different frequency bands serve different purposes:

• Key insight: High-frequency components distinguish nearby tokens; low-frequency components capture long-range structure. Uniform scaling damages local precision unnecessarily.

• The NTK formula: Replace base $b$ with $b' = b \cdot s^{d/(d-2)}$, where $s$ is the context extension factor. This preserves high frequencies while scaling low frequencies.

• Frequency-dependent scaling: The effective scaling factor increases from 1 (no change) for the highest frequencies to $s$ (full scaling) for the lowest frequencies.

• Dynamic variant: Adapt the base in real-time based on current sequence length: $b'(L) = b \cdot (L/L_{\text{train}})^{d/(d-2)}$.

• Practical benefits: Better preservation of local attention patterns compared to Position Interpolation, especially for tasks requiring precise nearby-token relationships.

NTK-aware scaling improved context extension, but researchers continued seeking even better solutions. In the next chapter, we'll explore YaRN (Yet another RoPE extension method), which builds on NTK-aware principles while adding attention scaling and temperature adjustments for further improvements.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about NTK-aware scaling and context extension in transformer models.