人工智能咨询培训老师叶梓 转载标明出处
由于大模型高昂的微调成本、长文本的稀缺性,以及新引入的标记位置可能导致的灾难性值,目前扩展的语境窗口通常被限制在大约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的上下文窗口方面是有效的,并且在不同长度的上下文上都能保持较低的困惑度,展示了其在长文档处理任务中的潜力。