LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

The limit of formula-based frequency scaling

YaRN and NTK-aware scaling showed that non-uniform treatment of RoPE frequency dimensions is better than the uniform compression linear interpolation applies. But both methods still derive their rescaling factors from a single closed-form formula applied to each dimension based on its wavelength. The optimal rescaling factor for each dimension is not actually determined by wavelength alone — it depends on the specific frequency, the training data distribution, and how the model's attention heads have specialized during pretraining. A formula that groups dimensions by wavelength thresholds is a reasonable heuristic, but it is not the optimal solution.

LongRoPE makes this precise: the rescaling factors $\{\lambda_d\}_{d=1}^{D/2}$, one per RoPE frequency pair, are free parameters to be optimized rather than values to be computed from a formula. The question becomes: can we find the optimal set of $D/2$ factors efficiently?

Evolutionary search over rescaling factors

LongRoPE frames the search as a black-box optimization problem. The objective is to minimize perplexity on a validation set at the target context length $L_{\text{target}}$, evaluated by running a forward pass with a candidate set of rescaling factors applied to the RoPE frequencies. No gradients through the optimization loop are needed — only repeated forward passes, each of which is cheap relative to training.

The search algorithm is a population-based evolutionary strategy. Starting from a population of candidate factor sets (initialized around known good baselines like NTK scaling), the algorithm evaluates perplexity for each candidate, selects the better-performing ones, applies mutation (small perturbations to individual $\lambda_d$ values), and iterates. The search space is $\mathbb{R}^{D/2}$, which for a typical model has on the order of 64 dimensions — small enough for evolutionary search to converge without exhaustive grid traversal.

The result is a set of per-dimension factors that is strictly more expressive than any formula-derived set: it can capture the fact that dimension 12 in a particular model benefits from aggressive interpolation while dimension 13 does not, a pattern no wavelength-based heuristic would identify.

Two-stage extension

Jumping directly from training length $L_{\text{train}}$ to a very long target context $L_{\text{target}}$ in a single fine-tuning stage causes catastrophic degradation — the model's weights are too far from the regime the new positional encodings require. LongRoPE addresses this with a two-stage approach:

1. Stage 1: search for optimal rescaling factors for an intermediate length $L_{\text{mid}}$ (e.g., 128K when targeting 2M), fine-tune for ~1000 steps at $L_{\text{mid}}$ using those factors.
2. Stage 2: search again for optimal factors for $L_{\text{target}}$, fine-tune for another ~1000 steps.

Each search uses the current model checkpoint as its base, so the factors found in stage 2 are appropriate for a model that has already been adapted to $L_{\text{mid}}$. The total fine-tuning budget is roughly 2000 steps — comparable to YaRN's budget for a much shorter extension — yet yields stable perplexity at 2 million tokens.

Short-context recovery

A persistent failure mode of long-context extension methods is regression on short sequences. After adapting a model to very long contexts, its performance on the standard benchmarks it was originally strong on (MMLU, GSM8K, etc.) can degrade measurably. The rescaling factors optimized for $L_{\text{target}}$ are not the same ones that were optimal at the original training length, and the fine-tuning distribution is dominated by long sequences.

LongRoPE fixes this with a brief re-fine-tuning step at the original training length $L_{\text{train}}$, using the original (unmodified) RoPE frequencies. This step takes roughly 400 additional steps on a short-context mix. At inference time, the model uses two configurations: the searched long-context factors for inputs longer than some threshold, and the original factors for shorter inputs. The short-context recovery step largely restores the original benchmark numbers without undoing the long-context capability.

Results

Applied to Phi-3-mini, LongRoPE achieves a 2 million token context window with stable perplexity across the full range — from a few hundred tokens to 2M. At intermediate lengths (128K, 256K), LongRoPE outperforms YaRN in perplexity at a similar or smaller fine-tuning budget. The per-dimension factors found by evolutionary search are measurably better than the formula-based alternatives, confirming that the added search cost is worthwhile.

The core contribution is the separation of concerns: finding the right rescaling factors is a search problem, not a formula problem, and treating it as such yields strictly better results. The two-stage extension and short-context recovery are engineering additions that make the approach practical, and together they define the recipe for long-context adaptation that subsequent work builds on.