使用 RoPE 的模型在预训练长度内学会了不同频率的位置旋转。输入突然拉长,旋转角度会进入训练时没见过的范围。位置插值可以把更长的序号压回旧范围,却会同时挤压近距离的位置差,让模型熟悉的局部关系变得难分。
我对照了 YaRN 论文的方法、训练配置和长文评测。YaRN 的名字展开为 Yet another RoPE extensioN method。它把 RoPE 的频率分开处理,保留对局部距离重要的高频部分,对需要避免外推的低频部分做插值,中间频率平滑过渡,随后再调整注意力分数的尺度。
改频率以后还要继续训练
YaRN 没有改 Transformer 层的主要结构。实现只需改变最初生成 RoPE 频率的算法,推理阶段可以继续配合 FlashAttention 2,也不增加一套随长度增长的新参数。现成权重仍需要少量继续训练,让模型适应新的位置分布。
论文用 Llama 2 的 7B 与 13B 模型做实验。扩展到 64K 的版本在 PG19 长文本上训练 400 步,全局批次为 64,每条训练片段长 64K。作者估算所用数据约为原预训练语料的 0.1%,相比论文选取的既有办法少用十倍 Token,训练步数少 2.5 倍。
128K 版本的过程更能说明边界。作者先得到扩展倍数为 16 的 64K 模型,再把倍数提高到 32,只追加 200 步训练,训练样本仍为 64K。模型随后在最长 128K 的 Proof-pile 文档上评测。7B 版本在 131072 Token 处的滑窗困惑度为 2.37,13B 版本为 2.24,论文据此认为它能外推到训练片段以外。
这里的 128K 是特定模型和评测方法下的结果。实验从 Proof-pile 选择了十篇至少 128K 的文档,困惑度采用步长 256 的滑窗计算。另一个 GovReport 测试只在 32K 窗口比较了五十篇长文。数字能证明方法值得尝试,不能说明任意 128K 输入都能被均匀理解。
GovReport 的结果也值得单独看。32K 窗口下,YaRN 64K 版本的 7B 与 13B 模型困惑度分别为 3.59 和 3.35;128K 版本分别为 3.64 和 3.39。更大的标称窗口在这组较短评测上没有自动得到更低数值。扩展倍数应按目标长度选择,继续放大频率范围并非免费收益。
短任务也要守住。论文把扩展后的模型放回常见基准,目的正是检查位置改造有没有损伤原有能力。长文困惑度下降只能说明下一词预测更稳,问答、代码和短指令仍需各自的分数。
长窗口要测信息有没有被用到
部署长上下文模型时,能放入多少 Token 只是第一项。还要在不同位置放置关键信息,测试模型能否找回;跨章节的问题应单独测,短任务也要复查,避免延长窗口以后损伤原有能力。显存、预填充时间和首词元延迟同样会随输入增长。
YaRN 降低了改造 RoPE 模型的训练成本,也把频率处理做得更细。它没有替模型增加无限记忆。标称窗口、困惑度和真实问答各自回答不同问题,三项一起测,才知道扩出来的长度有多少能用。