YaRN：一种高效RoPE扩展方法，可推理更长上下文并达到SOTA

摘要

旋转位置编码（RoPE）已被证明可以有效地在基于Transformer的语言模型中编码位置信息。然而，这些模型在超过它们训练的序列长度后无法推广。我们提出了YaRN（另一种RoPE扩展方法），这是一种计算高效的方法，可以扩展此类模型的上下文窗口，所需token减少10倍，训练步骤减少2.5倍。使用YaRN，我们展示了LLaMA模型可以有效地利用和推断出比其原始预训练允许的上下文长度长得多的上下文长度，并且在上下文窗口扩展中达到了SOTA。此外，我们证明YaRN表现出了超越微调数据集有限上下文的能力。

背景知识

旋转位置编码

位置插值

附加符号

方法

然而，尽管PI在所有RoPE维度上实现了均等的拉伸，但我们发现PI所描述的理论插值界限在预测RoPE与LLM内部嵌入之间的复杂动态时不足够准确。在接下来的小节中，我们分别描述了我们在PI中找到的主要问题，并解决了这些问题，以便向读者提供每种方法的背景、起源和理论依据，以此来获得完整的YaRN方法。

高频信息丢失 - “NTK 感知”插值

为了解决在插值RoPE嵌入时丢失高频信息的问题，[4]中开发了"NTK-aware"插值。与同样乘以因子s的方式相比，我们通过在多个维度上缩放高频率较小并且低频率较大的方式来分散插值压力。可以通过多种方式获得这样的转换，但最简单的方式是对θ的值进行基变换。我们希望最低频率的缩放与线性位置缩放尽可能一致，并且最高频率保持不变，因此我们需要找到一个新的基b′，使得最后一个维度与具有缩放因子s的线性插值的波长匹配。由于原始的RoPE方法跳过奇数维度，以便将cos(2πx/λ)和sin(2πx/λ)分量连接为单个嵌入，所以最后一个维度d ∈ D为|D|−2，解b′的值如下：

根据符号表示法，"NTK-aware"插值方案只需简单地应用基变换公式，即：

在测试中，与PI方法相比，这种方法在扩展非精调模型的上下文大小方面表现更好。然而，这种方法的一个主要缺点是，由于它不仅仅是一个插值方案，某些维度会略微外推到“超出界限”的值，因此使用“NTK-aware”插值进行精细调节得到的结果比PI方法更差。此外，由于“超出界限”的值，理论缩放因子s并不能准确描述实际上的上下文扩展比例。在实践中，为了给定上下文长度的扩展，必须将缩放值s设得比预期的缩放值更高。

相对局部距离的损失 - “NTK-by-parts”插值

RoPE嵌入的一个有趣观察是，给定上下文大小L，在一些维度d中，波长比预训练期间观察到的最大上下文长度要长（λ > L），这表明某些维度的嵌入可能在旋转域中分布不均匀。在PI和 “NTK-aware” 插值的情况下，我们将所有的RoPE隐藏维度视为相等。然而，实验中发现网络对一些维度与其他维度有所不同。如前所述，给定上下文长度L，某些维度的波长λ大于或等于L。考虑到当隐藏维度的波长大于或等于L时，所有位置对编码唯一的距离，我们假设绝对位置信息被保留下来，而当波长较短时，网络只能获得相对位置信息。当我们通过缩放因子s或使用基础变换b′拉伸所有RoPE维度时，所有的标记都会彼此更加接近，因为较小旋转角度下两个向量的点积较大。这种缩放严重影响了LLM理解其内部嵌入之间的小而局部关系的能力。假设这种压缩导致模型对接近的标记的位置顺序感到困惑，从而损害了模型的能力。为了解决这个问题，根据我们的观察，选择根本不插值较高频率的维度。

为了确定我们所需要的维度d，在给定某个上下文长度L下，我们可以按照如下的方式扩展方程（14）：

我们还提出对于所有满足r < α的维度d，我们会线性插值一个比例s（就像À一样，避免任何外推），而对于r>β的维度d，则不进行插值（总是进行外推）。定义斜坡函数γd为：

在斜坡函数的帮助下，我们将新的波长定义为：

α和β的值应该根据具体情况进行调整。例如，我们经实验证明，对于Llama模型家族而言，α和β的合适取值为α = 1和β = 32。将λd转换为θd后，该方法可以描述为：

使用本节中描述的技术，我们发布了一种名为“NTK-by-parts”插值的改进方法。这个改进的方法比之前的PI和“NTK-aware”插值方法在未微调模型和微调模型上的表现都更好。由于该方法避免了在旋转领域中具有不均匀分布的维度进行外推，它避免了之前方法中的所有微调问题。

动态缩放 - “动态 NTK”插值

当使用RoPE插值方法来扩展上下文大小而无需进行微调时，我们希望模型在更长的上下文大小下能够逐渐降级，而不是在设置的比所需值更高的比例s时在整个上下文大小上完全降级。回想一下，s = L′/L表示PI中的比例，其中L是训练的上下文长度，L′是新扩展的上下文长度在“动态NTK”方法中，我们动态计算比例s如下：

在推理过程中，当上下文大小超过训练的上下文限制L时，动态改变尺度允许所有模型以平滑方式退化，而不是立即失败。

在使用动态尺度调整和kv缓存时需要注意，因为在某些实现中，RoPE嵌入是被缓存的。正确的实现应该在应用RoPE之前缓存kv嵌入，因为每个标记的RoPE嵌入在s改变时也会改变。

增加长距离平均最小余弦相似度-YaRN

即使我们解决了在前面描述的局部距离的问题，更大的距离也必须在阈值α处进行插值以避免外推。直观来说，这似乎不是一个问题，因为全局距离不需要高精度来区分token的位置。然而，我们发现，由于平均最小距离随着令牌数的增加而变得更加接近4，它使得注意力softmax分布变得“更尖锐”（即降低了注意力softmax的平均熵）。换句话说，由于插值导致远距离衰减的效果减弱，网络“更关注”较多的token。这种分布偏移导致LLM的输出出现退化，这与之前的问题无关。

由于在将RoPE嵌入插值到更长的上下文大小时，注意力softmax分布的熵减小，我们的目标是反向减小熵（即增加注意力logits的“温度”）。这可以通过将中间注意力矩阵乘以一个大于1的温度t，然后应用softmax来实现，但由于RoPE嵌入被编码为旋转矩阵，我们可以简单地通过一个常量因子√t来调整RoPE嵌入的长度。这种“长度缩放”技巧使我们避免了对注意力代码的任何修改，从而显着简化了与现有的训练和推断流程集成的过程，并且具有O(1)的时间复杂度。

由于我们的RoPE插值方案不均匀地插值RoPE维度，因此很难计算所需的温度比例t相对于比例s的解析解。幸运的是，我们通过最小化困惑度实验发现，所有的Llama模型都大致遵循同一拟合曲线：

上面的方程是通过将 perplexity 最低的√t 与比例扩展因子 s 对拟合得到的，在 LLaMA 7b、13b、33b 和 65b 模型上，使用在 3.2 中描述的插值方法。我们还发现，这个方程也适用于Llama 2 模型 (7b、13b 和 70b)，只有轻微的差异。这表明增加熵的这个属性在不同的模型和训练数据中是普遍且可推广的。

外推和迁移学习

在Code Llama中，使用了一个包含16k个上下文的数据集，并将比例因子设置为s≈88.6，这对应于一个上下文大小为355k。他们表明，网络在训练过程中从未见过这些上下文尺寸的情况下，可以推断出多达100k个上下文。YaRN也支持使用比数据集长度更高的比例因子s进行训练。由于计算约束，我们仅测试s=32，通过对具有64k上下文的相同数据集进行200个步骤的微调来进一步调优s=16模型。

我们展示了s=32模型在训练过程中使用仅有64k上下文成功推断出128k上下文的情况。与以前的“盲目”插值方法不同，YaRN在增加比例因子s时的迁移学习效率更高。这证明了从s=16到s=32的成功迁移学习，而网络无需重新学习插入的嵌入，因为s = 32模型在整个上下文尺寸上等效于s=16模型，尽管仅在s=32上进行了200个步骤的训练。

评估

长序列语言建模

为了评估长序列语言建模的性能，我们使用了GovReport和Proof-pile数据集，这两个数据集都包含许多长序列样本。对于所有评估，我们只使用了这两个数据集的测试集进行计算。所有困惑度评估都是使用来自的滑动窗口方法进行计算的，其中S=256。

首先，我们评估了模型在上下文窗口增大时的性能。我们从Proof-pile中选择了10个至少为128k标记长度的随机样本，并在从2k标记的序列长度到128k标记的情况下截断了每个样本的前2k个步骤，并评估了每个样本的困惑度。图1显示了与其他上下文窗口扩展方法相比的结果。特别是，我们与Together.ai的公开可用的32k PI模型 和"NTK-aware" Code Llama进行了比较。这些结果总结在表1中。

表2显示了在32k上下文窗口下，对至少含有16k个标记的50个未截断的GovReport文档进行评估的最终困惑度。

我们观察到该模型在整个上下文窗口中表现出色，优于所有其他上下文窗口扩展方法。特别值得注意的是YaRN（s=32）模型，在128k中继续呈现出下降的困惑度，尽管微调数据仅限于长度为64k的标记。这表明该模型能够推广到未见过的上下文长度。

标准化基准

Hugging Face Open LLM排行榜对一系列LLM模型在四个公共基准测试上进行了比较。具体来说，有25-shot ARC-Challenge [9]，10-shot HellaSwag，5-shot MMLU和0-shotTruthfulQA。

为了测试模型性能在上下文扩展下的退化情况，我们使用这套测试套件评估了我们的模型，并将其与Llama 2基准模型的已建立分数以及公开可用的PI和“NTK-aware”模型进行了比较，结果总结在表3中。

我们观察到YaRN模型和它们相应的Llama 2基线之间存在最小的性能下降。我们还观察到在YaRN s=16和s=32模型之间的平均分数下降了0.49%。因此我们得出结论，从64k到128k的迭代扩展导致的性能损失可以忽略不计。

总结

综上所述，我们展示了YaRN改进了所有现有的RoPE插值方法，并可以作为PI的替代方案，其没有缺点且实施起来所需的工作量很小。经过微调的模型在多个基准测试中保留其原有的能力，同时能够关注非常大的上下文范围。此外，在较短的数据集上进行微调的YaRN能够高效地进行外推，并且能够利用迁移学习实现更快的收敛，这对于计算限制的情况至关重要。最后，我们展示了YaRN进行外推的有效性，其中它能够"训练短，测试长"。