NTK-aware Scaling › Key Parameters

When implementing NTK-aware scaling, several parameters control the behavior of the context extension. Getting these right is important for correct behavior and good performance.

• 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). Using the wrong base will produce frequencies that are scaled from the wrong starting point, degrading performance in unpredictable ways. Always check the model card or configuration file for the original base.

• 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. In practice, $s$ between 2 and 8 works well with minimal fine-tuning, while $s$ above 16 typically requires significant adaptation training.

• 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. Note that this is the head dimension, not the total model dimension; if the model has a total embedding dimension of 4096 with 32 attention heads, each head has dimension $4096 / 32 = 128$, and you should use $d = 128$ in the formula.

• train_length (for dynamic scaling): The context length the model was originally trained with. This is the threshold below which no scaling is applied. Using the correct value is important for dynamic NTK to work properly. If you set this too low, the model will apply scaling even for short sequences that do not need it. If you set it too high, the model will fail to scale up for sequences that exceed the training context.

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. More importantly, dynamic scaling has implications for KV-cache reuse, as described in the Limitations section. For most production deployments, static scaling is recommended unless your application has highly variable context lengths that would make aggressive static scaling wasteful for short inputs.