NTK-Aware Scaling

The Real Problem: PI Crushes Local Resolution

Pretraining a 70B model on a 32K context costs the same as pretraining it on a 4K context — about eight times more. The attention matrix grows as $O(L^2)$, the activation memory grows linearly, and the rare extremely-long documents that actually teach the model long-range structure are the bottleneck. Nobody wants to redo this. So almost every modern LLM is trained on a short context and extended afterwards — Llama 2 went 4K → 32K, Llama 3 went 8K → 128K, DeepSeek-V3 went 4K → 128K.

The obvious extension method — Position Interpolation (PI) by Chen et al. (2023) — is breathtakingly simple. RoPE turns position $m$ into a rotation by angle $m \theta_i$ on each dimension pair $i$. To extend from $L_{\text{train}}$ to $L_{\text{new}} = s \cdot L_{\text{train}}$, PI says: feed $m/s$ in instead of $m$. The model now sees the same angles it always saw, just at finer spacing. Reasonable, right?

Wrong — and the reason is the only thing that matters in this section. RoPE encodes position across a spectrum of frequencies. The first dimension pair rotates almost a full revolution between adjacent tokens ($\theta_0 = 1$). The last dimension pair barely moves over the entire trained context ($\theta_{d/2-1} \approx 10^{-4}$). These two ends of the spectrum do different jobs. High-frequency pairs distinguish neighbours — they let attention say "this token is the next token, not the one five tokens away." Low-frequency pairs encode long-range position — they let attention say "this token came from paragraph 3, not paragraph 17."

NTK-Aware Scaling, proposed in mid-2023 by the pseudonymous Reddit user bloc97 and rapidly adopted across the LLaMA ecosystem, fixes this with a one-line change. Same goal as PI. Same zero training cost. But it interpolates only the dimensions that needed it.

Intuition: Position Is a Frequency Spectrum

Imagine a piano with $d/2$ keys, one per dimension pair. Striking key $i$ at position $m$ rings a single tone at frequency $\theta_i$; the chord of all keys at position $m$ is the position embedding. The leftmost keys are the tweeters — they vibrate fast and distinguish tokens that are just one or two apart. The rightmost keys are the woofers — they vibrate slowly and distinguish chapters from chapters.

Now you record the piano at a sample rate that covers exactly $L_{\text{train}}$ seconds (the trained context). You want to play it back for $L_{\text{new}} = s \cdot L_{\text{train}}$ seconds without re-recording. You have two options:

1. Slow the whole tape by $s$ (PI). The bass survives, but every tweeter note is now an indistinct hum — you can no longer hear adjacent-token detail. The high frequencies have been driven below the model's ability to discriminate them.
2. Keep the tweeters at their original speed; slow only the woofers (NTK-Aware). The low frequencies were the only ones that ever needed to be stretched, because they were the only ones whose period was longer than $L_{\text{train}}$ in the first place. The tweeters were already cycling through full revolutions every few tokens — there is nothing to stretch.

The name "NTK-Aware" comes from a Neural Tangent Kernel argument: in the NTK limit, the contribution of each frequency component to the kernel decays geometrically with frequency, so the high-frequency components are the ones that carry fine-grained positional information through the network and must not be distorted. The piano analogy is the engineer's version of the same claim.

The Mathematical Idea

RoPE assigns dimension-pair $i \in [0, d/2)$ the angular frequency

$\theta_i = \text{base}^{-2i/d}$ with $\text{base} = 10000$.

At position $m$, the rotation angle is $m \theta_i$. The highest frequency is $\theta_0 = 1$; the lowest is $\theta_{d/2-1} = \text{base}^{-(d-2)/d}$ which is tiny (≈ $1.16 \cdot 10^{-4}$ for $d=64$).

Position Interpolation (PI)

Replace $m$ with $m / s$. Equivalently, the new effective frequencies are

$\theta_i^{\text{PI}} = \theta_i / s$ for all $i$.

NTK-Aware

Keep the position $m$ integer-valued. Replace the base with a new constant $\text{base}^{\prime}$ chosen so that the lowest frequency gets exactly the same compression as PI:

$(\text{base}^{\prime})^{-2(d/2-1)/d} = \text{base}^{-2(d/2-1)/d} \cdot \frac{1}{s}$

Solve for $\text{base}^{\prime}$:

$\text{base}^{\prime} = \text{base} \cdot s^{\,d/(d-2)}$

The new frequencies are then

$\theta_i^{\text{NTK}} = (\text{base}^{\prime})^{-2i/d}$.

Look at the two endpoints. For $i = 0$: $\theta_0^{\text{NTK}} = 1$ (unchanged from original). For $i = d/2 - 1$: $\theta_{d/2-1}^{\text{NTK}} = \theta_{d/2-1} / s$ (matches PI). In between, the ratio $\theta_i^{\text{NTK}} / \theta_i$ smoothly interpolates from 1 (no compression at the high end) down to $1/s$ (full PI compression at the low end). The in-between bowed curve is what saves local resolution.

Why this exponent? The denominator $d-2$ appears because RoPE's lowest pair is at index $d/2-1$, so the exponent on $\text{base}$ for that pair is $-2(d/2-1)/d = -(d-2)/d$. Inverting that exponent to solve for $\text{base}^{\prime}$ gives $s^{d/(d-2)}$. For $d \to \infty$ the exponent approaches $s^1 = s$; for $d = 64$ it is $s^{1.032}$ — barely larger than $s$.

Manual Numerical Walkthrough

One worked example with $d = 8$, $s = 4$, $L_{\text{train}} = 1024$, $L_{\text{new}} = 4096$. Small enough to do every step by hand; structurally identical to the real $d = 128$ case.

Visualizing the Frequency Stretch

Move the sliders. The x-axis is dimension-pair index $i$ from $0$ (tweeters, left) to $d/2-1$ (woofers, right). The y-axis is the rotation angle each pair sees at the new boundary position $L_{\text{new}}$, plotted on log scale. The dashed amber line is what the model saw at training; if a curve sits above it, that pair is being asked to extrapolate.

Three things are visible at every (d, s):

• The grey original curve sits above the amber training line everywhere — that's the problem in one picture.
• The red PI curve sits exactly on the amber line for every pair. Safe — but no headroom for the high-frequency end, which is where local information lives.
• The cyan NTK-Aware curve touches the amber line at the woofer end and rises smoothly toward the original at the tweeter end. The bowed shape is the entire point.

Plain Python: Three Ways to Stretch RoPE

No tensors yet. Just three pure-Python functions that compute the per-pair frequencies for the original RoPE, for Position Interpolation, and for NTK-Aware Scaling. The whole comparison fits in twenty lines.

1import math
2
3BASE = 10000          # original RoPE base frequency
4D    = 64             # head dimension (must be even)
5L_TRAIN = 4096        # context length the model was trained on
6L_NEW   = 32768       # target context length after extension
7s = L_NEW / L_TRAIN   # extension factor (8x here)
8
9# 1) Original RoPE frequencies for one head, no extension.
10def rope_freqs(d, base):
11    # pair index i = 0..d/2-1; theta_i = base^(-2i/d)
12    return [base ** (-2 * i / d) for i in range(d // 2)]
13
14# 2) Position Interpolation (PI). Compress positions: m' = m / s.
15#    Equivalently, divide every frequency by s.
16def pi_freqs(d, base, s):
17    return [theta / s for theta in rope_freqs(d, base)]
18
19# 3) NTK-Aware Scaling. Stretch the base instead of the positions.
20#    Solve: (new_base)^(-2 (d/2 - 1) / d) = base^(-2 (d/2 - 1) / d) / s
21#    => new_base = base * s^(d / (d - 2))
22def ntk_freqs(d, base, s):
23    new_base = base * (s ** (d / (d - 2)))
24    return rope_freqs(d, new_base)
25
26orig = rope_freqs(D, BASE)
27pi   = pi_freqs(D, BASE, s)
28ntk  = ntk_freqs(D, BASE, s)
29
30print(f"highest-freq pair (i=0):")
31print(f"  orig theta = {orig[0]:.6f}, PI = {pi[0]:.6f}, NTK = {ntk[0]:.6f}")
32print(f"lowest-freq pair  (i=d/2-1 = {D//2-1}):")
33print(f"  orig theta = {orig[-1]:.3e}, PI = {pi[-1]:.3e}, NTK = {ntk[-1]:.3e}")
34
35# What rotation angle does each pair see at the new boundary position?
36def max_angle(thetas, m):
37    return [m * t for t in thetas]
38
39orig_at_train = max_angle(orig, L_TRAIN)
40pi_at_new     = max_angle(pi,   L_NEW)
41ntk_at_new    = max_angle(ntk,  L_NEW)
42
43print()
44print("Boundary-position rotation angles:")
45print(f"  orig at L_train, highest pair:   {orig_at_train[0]:8.1f} rad")
46print(f"  PI   at L_new,   highest pair:   {pi_at_new[0]:8.1f} rad")
47print(f"  NTK  at L_new,   highest pair:   {ntk_at_new[0]:8.1f} rad")
48print(f"  orig at L_train, lowest pair:    {orig_at_train[-1]:.4f} rad")
49print(f"  PI   at L_new,   lowest pair:    {pi_at_new[-1]:.4f} rad")
50print(f"  NTK  at L_new,   lowest pair:    {ntk_at_new[-1]:.4f} rad")

Running this prints exactly the numbers from the walkthrough table above (scaled up to d = 64): the highest-frequency pair is left alone by NTK and crushed by 8× by PI; the lowest-frequency pair is landed on the same value by both methods.

PyTorch: Drop-In NTK-Aware RoPE

Production RoPE precomputes two tables — a cosine table and a sine table, both of shape $(L_{\text{max}}, d/2)$ — and reads from them during attention. Extension means rebuilding those tables. PI changes the positions column; NTK-Aware changes the base constant. Everything else — the rotation kernel, the attention math, the KV cache — stays identical. This is why both methods are described as training-free: zero changes to weights, zero changes to the rest of the forward pass.

shape (max_pos, 32) — one angle per (m, i)

(32768, 32) — same as PI

(2, 8, 32768, 64) — batch, heads, seq, dim

(2, 8, 32768, 64) — same shape

1import torch
2
3def build_rope_cos_sin(d, max_pos, base):
4    """Precompute cos/sin tables for RoPE — exactly as in Llama/DeepSeek."""
5    half = d // 2
6    pair_idx  = torch.arange(half, dtype=torch.float32)   # [0, 1, ..., half-1]
7    inv_freq  = 1.0 / (base ** (2 * pair_idx / d))        # theta_i
8    positions = torch.arange(max_pos, dtype=torch.float32)
9    # Outer product: angles[m, i] = m * theta_i
10    angles = positions[:, None] * inv_freq[None, :]       # (max_pos, half)
11    return angles.cos(), angles.sin()                     # both (max_pos, half)
12
13def rotate_half(x):
14    # Split last dim in half, rotate the pair: (x1, x2) -> (-x2, x1)
15    x1, x2 = x[..., ::2], x[..., 1::2]
16    return torch.stack((-x2, x1), dim=-1).flatten(-2)
17
18def apply_rope(q, cos, sin):
19    # q: (..., seq, d). cos/sin: (seq, d/2) -> broadcast over last dim by repeating.
20    cos_full = cos.repeat_interleave(2, dim=-1)            # (seq, d)
21    sin_full = sin.repeat_interleave(2, dim=-1)            # (seq, d)
22    return q * cos_full + rotate_half(q) * sin_full
23
24# ---- The two extension methods ----
25D, L_TRAIN, L_NEW = 64, 4096, 32768
26s = L_NEW / L_TRAIN
27
28# (a) Position Interpolation: same base, but feed positions m / s into the table.
29def pi_cos_sin(d, max_pos, base, s):
30    cos, sin = build_rope_cos_sin(d, max_pos, base)
31    # Resample so that the m-th entry corresponds to true position m/s.
32    # In practice: build for max_pos new positions with positions = arange / s.
33    half = d // 2
34    inv_freq  = 1.0 / (base ** (2 * torch.arange(half) / d))
35    positions = torch.arange(max_pos) / s                  # KEY: divide
36    angles    = positions[:, None] * inv_freq[None, :]
37    return angles.cos(), angles.sin()
38
39# (b) NTK-Aware: new base; positions stay integer-valued.
40def ntk_cos_sin(d, max_pos, base, s):
41    new_base = base * (s ** (d / (d - 2)))                 # KEY: stretch base
42    return build_rope_cos_sin(d, max_pos, new_base)
43
44# Plug into a forward pass exactly like before:
45q = torch.randn(2, 8, L_NEW, D)                            # (batch, heads, seq, d)
46cos_ntk, sin_ntk = ntk_cos_sin(D, L_NEW, base=10000, s=s)
47q_rope = apply_rope(q, cos_ntk, sin_ntk)
48
49print("cos/sin table:", cos_ntk.shape)                     # (32768, 32)
50print("q after RoPE:", q_rope.shape)                       # (2, 8, 32768, 64)
51print("new base used:", 10000 * (s ** (D / (D - 2))))

In Hugging Face Transformers, this exact substitution is selected by setting \text{rope\_scaling} = \{ \text{type}: \text{"dynamic_ntk"$, \text{factor}: s \}$ in the model config. In vLLM and SGLang it is --rope-scaling-type dynamic-ntk --rope-scaling-factor 8. Same math, same shapes, three minutes to enable on any RoPE-based LLM.

At Massive Scale: Why NTK Alone Is Not Enough

NTK-Aware was a breakthrough because for the first time we could push Llama-1 7B to 8K and 16K context with almost no perplexity regression — and zero training cost. But at extreme extension ratios ($s \geq 16$) and on modern $d = 128$ heads, two problems appear that bare NTK-Aware cannot fix.

Problem 1: The middle band still gets distorted

The bowed curve in the visualizer is smooth, but it does not match the truth of where extrapolation is safe and where it is not. Frequencies whose period is just barely longer than $L_{\text{train}}$ need PI-style compression (they have never seen a full revolution). Frequencies whose period is much shorter than $L_{\text{train}}$ need no compression at all (they cycled through every phase many times). NTK-Aware interpolates between these two regimes continuously, which means the mid-band pairs — the ones at the crossover — get the wrong treatment.

Problem 2: Attention logit scale changes

Stretching the base widens the inner product distribution of the rotated query and key. Softmax temperature interacts with this shift; at large $s$ the attention distribution becomes more uniform (over-smoothed) than what the model was trained on. This shows up as degraded retrieval quality at very long context even when perplexity looks healthy.

The standard remedy: YaRN

YaRN, covered next section, addresses both:

1. It thresholds the frequencies into three bands — fully extrapolate the high-freq band (NTK-style), fully interpolate the low-freq band (PI-style), and linearly ramp between them in the middle. This is a piecewise version of NTK-Aware that does not over-stretch the mids.
2. It rescales the attention logits by a small constant to compensate for the softmax temperature shift — typically $\sqrt{1 + 0.1 \ln s}$.

DeepSeek-V3's 128K context is built with YaRN on top of an original 4K-trained MLA RoPE, with $s = 32$. The middle-band correction matters enormously at that ratio — bare NTK-Aware at $s = 32$ would lose retrieval performance on the "needle in a haystack" tests that 128K context exists to pass.

Engineering Reality and Gotchas

1. Pre-compute once, share across the batch. The cos/sin tables depend only on (d, max_pos, base). Build them once at model load, store as buffers, and broadcast against any (batch, heads) shape. Re-building them per forward pass is a common performance bug.
2. Dynamic NTK at inference. A useful variant — "Dynamic NTK" — adjusts $s$ per forward pass based on the current sequence length: $s_t = \max(1, L_t / L_{\text{train}})$. For short prompts ($L_t \leq L_{\text{train}}$) you do nothing; for long prompts you stretch only as much as needed. The trade-off is that the cos/sin table must be rebuilt whenever the sequence crosses $L_{\text{train}}$ in a streaming setting — a one-time cost.
3. Combines with quantisation cleanly. Because the cos/sin tables are just data, they quantise to FP8/INT8 with no extra accuracy loss beyond what RoPE already had. NTK-Aware extension on a quantised model is essentially free.
4. The PI–NTK boundary depends on d. For $d = 16$ the exponent $d/(d-2) = 1.14$ is meaningfully larger than 1, so the new base inflates more aggressively. For $d = 128$ it is 1.016 — almost identical to a pure $s$ stretch. NTK-Aware behaves more like a re-derived PI on big heads and more like an aggressive tweeter-preserver on tiny heads.
5. Always test on retrieval, not just perplexity. A well-known failure mode is that bare NTK-Aware at $s = 16$ can improve perplexity on long-form text (because the low frequencies are now correctly spaced) while degrading needle-in-haystack accuracy (because the mid-band pairs are subtly miscalibrated). Run BABILong, RULER, or NIH before declaring victory.
6. Short fine-tune helps a lot, even though it's "training-free". The original NTK-Aware post billed it as zero-training, and it does work zero-shot. But the standard production recipe today is: apply NTK-Aware or YaRN, then fine-tune for $\sim 1$B tokens at the new context length. This recovers the small perplexity gap that any frequency-domain rescaling introduces. DeepSeek-V3 explicitly does this in the YaRN stage.

NTK-Aware Scaling is, on the surface, a one-line code change: replace the RoPE base. Under the surface, it is the first technique that treated position as a frequency spectrum and asked which frequencies actually need adjustment. Every long-context method that followed — YaRN, LongRoPE, the dynamic variants — is a refinement of the same idea. If you understand which dimension pairs are tweeters and which are woofers, you understand modern long-context extension.