LongRoPE2: Near-Lossless LLM Context Window Scaling

TL;DR

LongRoPE2 tackles the challenge of extending LLM context windows by diagnosing RoPE OOD as resulting from undertraining in higher rotation dimensions. It introduces a needle-driven, evolutionary search to identify real per-dimension rescaling factors and a mixed-context training regime to preserve short-context performance while enabling long-context understanding. Empirical results on Phi3-mini-3.8B and LLaMA3-8B show near-lossless 128k context extension with minimal loss to short-context tasks, outperforming YaRN, NTK, and Meta’s larger counterparts using far fewer training tokens. This approach offers a practical pathway to substantially longer context processing in existing LLMs and lays groundwork for scalable, near-lossless context expansion.

Abstract

LongRoPE2 is a novel approach that extends the effective context window of pre-trained large language models (LLMs) to the target length, while preserving the performance on the original shorter context window. This is achieved by three contributions: (1) a hypothesis that insufficient training in higher RoPE dimensions contributes to the persistent out-of-distribution (OOD) issues observed in existing methods; (2) an effective RoPE rescaling algorithm that adopts evolutionary search guided by "needle-driven" perplexity to address the insufficient training problem; (3) a mixed context window training approach that fine-tunes model weights to adopt rescaled RoPE for long-context sequences while preserving the short-context performance with the original RoPE. Extensive experiments on LLaMA3-8B and Phi3-mini-3.8B across various benchmarks validate the hypothesis and demonstrate the effectiveness of LongRoPE2. Remarkably, LongRoPE2 extends LLaMA3-8B to achieve a 128K effective context length while retaining over 98.5% of short-context performance, using only 10B tokens -- 80x fewer than Meta's approach, which fails to reach the target effective context length. Code will be available at https://github.com/microsoft/LongRoPE.

Figures (10)

• Figure 1: LongRoPE2-extended LLaMA3-8B achieves the best performance at a 128k context length among $\sim$10B models.
• Figure 2: (a) RoPE OOD (red area) when extending context length from 2k to 4k. (b) Per-dimensional RoPE rescaling factor from different approaches for extending Phi3-mini from 2k to 128k, all aligning with RoPE OOD theory. (c) Performance of Phi3-mini-128k after fine-tuning. Existing methods fail to achieve an effective 128k context length and show noticeable short-context performance drop.
• Figure 3: Sequence length required to span the theoretical period during Phi3-mini pre-training for different RoPE dimensions. Insufficient training in higher RoPE dimensions leads to shorter effective RoPE ranges and longer actual periods.
• Figure 4: Scale factors across different RoPE rescaling approaches.
• Figure 5: Mixed context window training to improve both short and long context capabilities.