NTK-aware Scaling › Dynamic NTK Scaling

Michael BrenndoerferUpdated July 1, 20255 min read
Language AI HandbookAll chapters
Part 18 of 67Long Context
Chapter 128NTK-aware Scaling › Dynamic NTK Scaling
Reading section 6 of 14 · Dynamic NTK Scaling · 5 minSection 6 of 14 · 5 min

Book contents

Section 6 of 14 · 5 min read

NTK-aware Scaling › Dynamic NTK ScalingOne section · Sections in this chapter6 / 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

Static NTK-aware scaling uses a fixed extension factor ss. 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.

A model serving a chatbot receives requests of widely varying lengths. Short questions, paragraph-length responses, and multi-page code reviews all pass through the same model. If you apply aggressive 8x static scaling to a short 200-token prompt, you are compressing frequencies that do not need to be compressed at all. The model's fine-tuned patterns for short-context reasoning may degrade unnecessarily.

Dynamic NTK scaling adapts the base in real-time based on the current sequence length. The formula is identical to static scaling, but the scale factor is computed from the actual sequence length at inference time rather than being fixed in advance:

b′(L)=b⋅(LLtrain)dd−2b'(L) = b \cdot \left(\frac{L}{L_{\text{train}}}\right)^{\frac{d}{d-2}}

where:

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

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

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

The clamping behavior is important. Without it, short sequences would have a scale factor less than 1, which would make the base smaller than the original and effectively increase frequencies. This would make adjacent tokens rotate faster than they did during training, which is a different kind of distortion. The clamp ensures dynamic NTK is always either "no change" or "compress," never "expand."

In[19]:
Code
def dynamic_ntk_frequencies(d, base, train_length, current_length):
    """Compute dynamically scaled NTK-aware frequencies.

    Args:
        d: Embedding dimension
        base: Original RoPE base (typically 10000)
        train_length: Training context length
        current_length: Current sequence length

    Returns:
        Frequencies for each dimension pair
    """
    if current_length <= train_length:
        # No scaling needed
        return compute_rope_frequencies(d, base)

    # Dynamic scaling factor
    scale = current_length / train_length
    ntk_base = base * (scale ** (d / (d - 2)))
    return compute_rope_frequencies(d, ntk_base)
In[20]:
Code
# Demonstrate dynamic scaling at different sequence lengths
d = 64
base = 10000
train_length = 4096
test_lengths = [2048, 4096, 8192, 16384, 32768]

freq_by_length = {
    L: dynamic_ntk_frequencies(d, base, train_length, L) for L in test_lengths
}
Out[21]:
Console
Dynamic NTK scaling: effective base at different sequence lengths
Training length: 4096

  Seq Length    Scale   Effective Base
--------------------------------------
        2048     1.00            10000
        4096     1.00            10000
        8192     2.00            20452
       16384     4.00            41829
       32768     8.00            85550

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. At 4,096 tokens or below (within the training distribution), the base remains at 10,000, meaning the model behaves exactly as it did during training. This is a clean and desirable property: dynamic NTK is a strict generalization of the original model, falling back to the original behavior whenever the sequence is short enough.

Notice that the effective base does not jump discontinuously at the training-length boundary. There is a smooth transition that begins exactly at LtrainL_{\text{train}}, which avoids any sudden behavioral change as sequences cross this threshold. This is important in applications like streaming generation, where the context window grows one token at a time; the model's behavior changes gradually rather than abruptly as the context crosses the training boundary.

Out[22]:
Visualization
Line plot showing how RoPE frequencies change with sequence length under dynamic NTK scaling, with different curves for different sequence lengths.
Dynamic NTK-aware frequencies adapt based on sequence length. Short sequences (at or below training length) use original frequencies, while longer sequences progressively modify frequencies to extend context.