长语境窗口扩展：LongRoPE技术解析

 人工智能咨询培训老师叶梓 转载标明出处

由于大模型高昂的微调成本、长文本的稀缺性，以及新引入的标记位置可能导致的灾难性值，目前扩展的语境窗口通常被限制在大约128k标记。为了克服这些限制，微软研究院的研究团队提出了一种名为LongRoPE的新技术，首次将预训练的大模型的语境窗口扩展到2048k标记，同时在256k训练长度内仅需1k微调步骤，同时保持原始短语境窗口的性能。

在大模型（LLMs）中，上下文窗口的扩展是一个至关重要的研究方向，因为它直接关系到模型处理长文本的能力。然而，这一过程并不简单，存在三个主要的技术障碍。

1、当模型尝试处理超出其训练时见过的位置索引时，会引入许多灾难性的值。这些值会导致模型输出分布出现问题，使得微调过程难以收敛。特别是在从4k扩展到超过1000k的位置索引时，新引入的位置超过90%，这对模型是一个巨大的挑战。

2、长文本的稀缺性也是一个问题。现有的数据集中，特别是超过1000k的长文本非常有限。而且，即使存在这样的文本，使用它们进行训练计算成本极高，需要大量的时间和GPU资源。

3、当上下文窗口扩展到极长时，模型的注意力会因为分布在过多的标记位置上而变得分散，这会降低模型在原始短上下文上的性能。

为了缓解这些挑战，研究者提出了一种基于RoPE（Rotary Positional Embedding）位置嵌入的插值方法。RoPE是一种在最近的LLMs中广泛使用的位置编码方式，它通过一组正余弦值来表示每个标记的位置信息。在RoPE中，每个维度的旋转角度与位置索引的乘积相关联，形成了一种基于频率的编码方式。

图2展示了不同插值方法下的RoPE嵌入示例。在直接外推中，新位置索引的RoPE被直接扩展到未训练的范围。在线性位置插值中，所有位置的旋转角度都按比例缩小。而在LongRoPE中，通过识别和利用RoPE维度和标记位置的两种非均匀性，实现了在不同维度和位置的多样化插值和外推。

方法

研究者首先介绍了两个关键的非均匀性：RoPE维度和token位置。基于这些非均匀性，LongRoPE利用高效的搜索算法来寻找最优的RoPE重新缩放因子。这些因子用于调整每个RoPE维度的旋转角度，以实现上下文窗口的扩展。

问题被构建为一个搜索问题，目标是为LLM找到一个目标上下文窗口大小L'的最优重新缩放因子。研究者定义了一个优化问题，通过最小化使用RoPE的LLM在输入样本X上的下一个token预测损失（即困惑度）来找到最优解。

为了解决这个问题，研究者设计了一个大型的搜索空间，包括RoPE维度和起始token位置的非均匀性。例如，λi可以在1.0到s×1.25之间搜索，其中s是目标上下文窗口扩展比例，步长为0.01。nˆ控制着不进行位置插值的初始token数量，可以从{0, 1, 2, ..., 256}中搜索。

为了高效地探索这个庞大的搜索空间，研究者采用了基于进化算法的搜索方法，并引入了两种优化技术。首先，他们不是随机初始化种群，而是在初始种群中加入了PI、NTK和YaRN对应的RoPE重新缩放因子。其次，他们引入了单调非递减约束，确保了RoPE维度的重新缩放因子λi随着维度的增加而单调递增，这与NTK理论相符。

研究者介绍了一种逐步扩展策略，首先将预训练的LLM扩展到256k，然后在此基础上进行微调，最终通过二次搜索达到2048k的上下文窗口大小。这种方法避免了直接在极长文本上进行微调的高成本。

图3 展示了不同方法在不进行微调的情况下，对LLaMA2-7B在PG19和Proof-Pile数据集上的困惑度的影响。LongRoPE方法能够将上下文窗口从4k扩展到32k，而不会像其他方法那样在扩展后导致困惑度激增。

表4 则具体展示了RoPE重新缩放因子的搜索空间，包括RoPE维度的λi和起始token位置的nˆ的取值范围和步长。

研究者还考虑了在扩展到极长上下文窗口后，原始上下文窗口性能下降的问题。为了解决这个问题，他们对扩展后的LLM进行了额外的进化搜索，以调整短上下文长度（例如4k和8k）的RoPE重新缩放因子。在推理过程中，LLM会根据序列长度动态调整相应的RoPE重新缩放因子，以恢复短上下文窗口的性能。

通过这些方法，LongRoPE不仅能够显著扩展LLM的上下文窗口，还能够保持在原始短上下文窗口中的性能，为LLM的应用开辟了新的可能性。

实验

研究者将LongRoPE应用于LLaMA2-7B和Mistral-7B两种模型上，并从三个方面评估了其性能：在长文档上的困惑度、从大量无关文本中检索简单密码的能力，以及在短的4096上下文窗口大小的标准LLM基准测试中的表现。

在微调方面，对于LLaMA2，研究者使用了2e-5的学习率，并以线性衰减的方式，全局批量大小为32，对Redpajama数据集进行了400步的微调，每段长度为128k，并以BOS和EOS标记作为边界。之后，基于已完成的检查点，再进行600步的微调以实现256k的上下文窗口。128k的上下文大小在8个A100 GPU上使用分布式训练系统进行训练，而256k则需要16个A100 GPU。对于Mistral，使用了1e-6的恒定学习率和64的全局批量大小。对于128k和256k的模型，研究者遵循YaRN的设置，在Together Computer的Long-Data Collections上进行了400步的训练，使用16k序列长度，训练使用了4个A100 GPU。

在搜索方面，对于目标窗口大小在256k以内的情况，研究者设置了种群大小P=64，变异大小N1=N2=16，迭代次数T=40，变异概率p=0.3，并在每次迭代中选择前32个进行变异/交叉。使用5个随机的PG19验证集样本计算困惑度，样本的最小长度需满足目标上下文长度。对于超过512k的窗口，研究者减半了种群、变异和交叉的大小。在Pile-Books3验证集上的3个随机样本上测量困惑度。

在256k的评估长度内，研究者开始与现有的扩展LLMs进行比较。他们使用了Proof-pile和PG19两个数据集来展示模型的泛化能力。通过256的滑动窗口在不同上下文长度下评估困惑度。对于PG19，研究者使用了100份完整测试数据。对于Proof-pile，他们随机选择了10个样本，每个样本长度至少为128k。

表5 和 表7 分别比较了在Proof-pile和PG19上通过不同插值方法扩展的LLaMA2和Mistral的困惑度。研究者强调了两个关键观察结果：(1) 他们的扩展模型在从4k到256k的评估长度上显示出整体下降的困惑度趋势，证明了它们利用更长上下文的能力。 (2) 即使在通常对在较短长度上保持性能具有挑战性的16倍长上下文窗口条件下，他们的LongRoPE-2048k模型在256k上下文长度内的性能也超过了现有最先进的基线。

在评估极长文档的有效性方面，研究者使用了Books3数据集。为了评估效率，他们随机选择了20本每本超过2048k长度的书籍，并使用256k的滑动窗口。

表6 显示，LongRoPE成功地将LLaMA2-7B和Mistral-7B的上下文窗口扩展到了2048k，同时在8k到128k的较短长度上实现了与基线相当或更优越的困惑度。研究者还观察到了2048k的LLaMA2和Mistral之间的显著性能差异。Mistral在较短长度上优于基线，但当超过256k时，困惑度超过了7。而LLaMA2的性能符合预期：随着上下文的增长，困惑度显著下降，在1024k和2048k时只有轻微的增加。此外，在LLaMA2上，LongRoPE-2048k在256k的微调长度上的表现优于128k，这是由于较小的二次扩展比率（即8倍对16倍）。与此相反，Mistral在128k的微调窗口大小上表现更好。

在消融研究中，研究者探讨了第二次位置插值的有效性。在他们的逐步扩展策略中，研究者使用搜索算法在微调后的扩展LLMs上进行了第二次非均匀位置插值。通过在微调后的LLaMA2-256k模型上进行实验，将其扩展到512k、1024k和2048k，并使用PI和YaRN进行比较。表9 显示，他们的非均匀位置插值在困惑度上保持了一致的水平。相比之下，PI和YaRN下的困惑度随着扩展比率的增加而迅速上升。

为了缓解在较短上下文长度上的性能损失，研究者通过搜索算法重新调整了LongRoPE-2048k的RoPE因子。具体来说，他们减少了搜索的最大允许缩放因子，以鼓励在短的4k和8k长度上进行较少的插值。表10 显示了LongRoPE-LLaMA2-2048k在Proof-pile上的4k和8k长度上的困惑度比较，以及平均LLM基准测试的准确性。结果清楚地表明，在短上下文长度上的性能有了显著提升。

最后研究者对两种非均匀性进行了消融研究，以了解每个部分对性能的贡献。他们设置了两个实验：(i) 使用不同的方法将LLaMA2-7B扩展到短的16k和32k - PI、仅搜索RoPE维度，以及搜索两种非均匀性；(ii) 按照相同的程序将我们的微调后的256k长度的LLaMA2扩展到2048k。困惑度在没有微调的情况下进行评估。表11 显示，与PI的线性插值相比，RoPE维度的非均匀性显著降低了困惑度。在16k和32k长度上，token位置的非均匀性明显提高了性能，但在2048k时没有显示出相同的影响，可能是由于长度极长。仅保留初始token不进行插值变得没有用，研究者将此留作未来的工作。

通过这些详细的实验设置和结果，研究者证明了LongRoPE方法在扩展LLMs的上下文窗口方面是有效的，并且在不同长度的上下文上都能保持较低的困惑度，展示了其在长文档处理任务中的潜力。