摘要
在大型语言模型（LLMs）日益成为自然语言处理领域的核心技术背景下，如何有效扩展其上下文窗口成为了一个重要的研究方向。本文提出了一种新颖的方法——LongRoPE2，旨在在保持原有短上下文性能的同时，显著扩展预训练语言模型的有效上下文窗口。通过三个主要贡献，LongRoPE2成功实现了这一目标：首先，提出了一个假设，认为在更高的RoPE维度上训练不足是导致现有方法在处理超出分布（OOD）问题时表现不佳的根本原因；其次，设计了一种有效的RoPE重标定算法，通过「针驱动」的困惑度评估来解决训练不足的问题；最后，采用混合上下文窗口训练方法，微调模型权重以适应长上下文序列的重标定RoPE，同时保持短上下文的性能。
1. 引言
随着LLMs的快速发展，长上下文窗口已成为其重要特征之一。以GPT-4o和LLaMA3.1为例，128k的上下文窗口已成为标准。然而，扩展上下文窗口的过程面临着两个主要挑战：首先，现有的RoPE重标定方法往往无法有效解决超出分布（OOD）问题，导致在扩展上下文时性能显著下降；其次，扩展上下文窗口通常会导致原有短上下文任务的性能显著下降。
LongRoPE2通过全面解决这些挑战，提供了一种有效的解决方案。我们的实验表明，LongRoPE2能够将LLaMA3-8B扩展至128K的有效上下文长度，同时保持超过98.5%的短上下文性能，且所需的训练数据仅为10B个token，远低于Meta的800B个token。
2. LongRoPE2方法论
2.1 RoPE重标定理论
Rotary Position Embedding（RoPE）在LLMs中广泛应用，但在上下文长度外推时面临挑战。RoPE的周期性特征导致在高维度上，未见完整的周期时会出现OOD值。LongRoPE2提出了一种新的RoPE重标定理论，旨在通过调整每个维度的旋转角度来确保高维度RoPE-OOD值保持在预训练范围内。
2.2 RoPE重标定因子搜索
LongRoPE2采用基于进化搜索的策略来识别实际的关键维度和最优重标定因子。通过合成「针驱动」的困惑度评估，LongRoPE2能够有效捕捉长上下文的能力，同时避免无关token对评估结果的干扰。
2.3 混合上下文窗口训练
为了同时提升长上下文和短上下文的性能，LongRoPE2引入了混合上下文窗口训练方法。该方法在短上下文训练中重用原始RoPE，而在长上下文训练中则应用重标定RoPE。这种方法不仅保留了短上下文的性能，还提升了长上下文的理解能力。
3. 实验与结果
3.1 实验设置
我们在LLaMA3-8B和Phi3-mini（3.8B. 上应用LongRoPE2，评估其在多个基准上的表现。实验结果表明，LongRoPE2在RULER基准测试中表现优异，尤其是在128k上下文长度的任务中，LongRoPE2的平均得分达到了82.03，显著优于其他现有方法。✅
3.2 真实世界基准测试
在真实世界的基准测试中，LongRoPE2也展现出了强大的性能，尤其是在LOFT、InfiniteBench和LongBench等任务中，均取得了显著的提升。这表明LongRoPE2不仅在合成任务中表现优异，也具备良好的泛化能力。
3.3 短上下文性能保持
与其他RoPE重标定方法不同，LongRoPE2在扩展上下文窗口的同时，能够保持短上下文任务的高性能。实验结果显示，LongRoPE2在Phi3-mini和LLaMA3-8B上分别保留了97.6%和98.6%的预训练性能，成为首个实现「近无损」扩展的方法。
4. 结论
LongRoPE2为大型语言模型的上下文窗口扩展提供了一种有效的解决方案，通过解决高维度RoPE训练不足的问题，结合进化搜索和混合上下文窗口训练，LongRoPE2实现了128k的有效上下文长度，并保持了短上下文的优异性能。未来的工作将探索将LongRoPE2扩展至完全无损和无限上下文窗口的可能性。