LongRoPE2: Near-Lossless LLM Context Window Scaling

📄 arXiv: 2502.20082v1 📥 PDF

作者: Ning Shang, Li Lyna Zhang, Siyuan Wang, Gaokai Zhang, Gilsinia Lopez, Fan Yang, Weizhu Chen, Mao Yang

分类: cs.CL

发布日期: 2025-02-27

🔗 代码/项目: GITHUB

💡 一句话要点

LongRoPE2：通过进化搜索和混合训练实现LLM近乎无损的上下文窗口扩展

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文学习 RoPE缩放 进化搜索 大型语言模型 上下文窗口扩展

📋 核心要点

  1. 现有方法在扩展LLM上下文窗口时，常在高RoPE维度上训练不足，导致分布外问题和性能下降。
  2. LongRoPE2通过进化搜索优化RoPE重缩放，并结合混合上下文窗口训练，解决高维度训练不足问题。
  3. 实验表明，LongRoPE2能将LLaMA3-8B扩展到128K上下文长度，同时保持98.5%的短上下文性能，且训练成本远低于Meta的方法。

📝 摘要（中文）

LongRoPE2是一种新颖的方法，它将预训练的大型语言模型（LLM）的有效上下文窗口扩展到目标长度，同时保持在原始较短上下文窗口上的性能。这通过三个贡献实现：（1）一个假设，即在高RoPE维度中训练不足会导致现有方法中观察到的持续的分布外（OOD）问题；（2）一种有效的RoPE重缩放算法，该算法采用由“needle-driven”困惑度引导的进化搜索来解决训练不足的问题；（3）一种混合上下文窗口训练方法，该方法微调模型权重以采用重新缩放的RoPE用于长上下文序列，同时使用原始RoPE保持短上下文性能。在LLaMA3-8B和Phi3-mini-3.8B上进行的大量实验验证了该假设并证明了LongRoPE2的有效性。值得注意的是，LongRoPE2扩展了LLaMA3-8B以实现128K的有效上下文长度，同时保留了超过98.5%的短上下文性能，仅使用了10B个token——比Meta的方法少80倍，而Meta的方法未能达到目标有效上下文长度。

🔬 方法详解

问题定义：现有方法在扩展LLM的上下文窗口时，通常会遇到性能下降的问题，尤其是在原始训练长度之外。一个关键的痛点是，模型在高维度的RoPE（旋转位置编码）上的训练不足，导致在长序列上出现分布外（OOD）的问题，从而影响模型的泛化能力和性能。

核心思路：LongRoPE2的核心思路是通过更有效地重新缩放RoPE，并结合混合上下文窗口训练，来解决高维度RoPE训练不足的问题。通过进化搜索找到最佳的RoPE缩放策略，使得模型在长序列上也能保持良好的性能。同时，通过混合训练，模型既能适应新的长上下文RoPE，又能保持在原始短上下文上的性能。

技术框架：LongRoPE2的整体框架包含三个主要阶段：1) RoPE重缩放：使用进化搜索算法，基于“needle-driven”困惑度指标，寻找最优的RoPE缩放策略。2) 混合上下文窗口训练：使用包含短上下文和长上下文序列的混合数据集，对模型进行微调，使其适应新的RoPE缩放策略。3) 评估：在各种基准测试上评估模型在长上下文和短上下文上的性能。

关键创新：LongRoPE2的关键创新在于其RoPE重缩放算法和混合上下文窗口训练方法。RoPE重缩放算法通过进化搜索，能够更有效地利用RoPE的不同维度，从而缓解高维度训练不足的问题。混合上下文窗口训练则保证了模型在扩展上下文窗口的同时，不会牺牲在原始短上下文上的性能。与现有方法相比，LongRoPE2能够以更低的训练成本实现更大的上下文窗口扩展。

关键设计：RoPE重缩放算法使用进化搜索，目标是最小化“needle-driven”困惑度。该困惑度通过在长序列中插入一个“needle”信息，并评估模型能否正确检索到该信息来衡量。混合上下文窗口训练使用原始RoPE和重缩放后的RoPE，分别对应短上下文和长上下文序列。训练损失函数是标准语言建模损失，优化器使用AdamW。具体参数设置（如进化搜索的种群大小、迭代次数，混合训练中短上下文和长上下文序列的比例）未知，需要在代码公开后进一步分析。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LongRoPE2在LLaMA3-8B上实现了128K的有效上下文长度，同时保留了超过98.5%的短上下文性能。与Meta的方法相比，LongRoPE2仅使用了10B个token进行训练，训练成本降低了80倍，且Meta的方法未能达到相同的有效上下文长度。这些实验结果表明，LongRoPE2是一种高效且有效的LLM上下文窗口扩展方法。

🎯 应用场景

LongRoPE2可应用于需要处理长序列数据的各种场景，如长文档摘要、代码生成、对话系统和知识库问答。通过扩展LLM的上下文窗口，LongRoPE2能够提升模型在这些任务上的性能，使其能够更好地理解和利用长距离依赖关系，从而产生更准确、更连贯的结果。该方法降低了长文本处理的计算成本，具有广泛的应用前景。

📄 摘要（原文）

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.