YaRN：大语言模型的高效上下文窗口扩展

摘要

旋转位置嵌入（RoPE）已被证明可以在基于 Transformer 的语言模型中有效编码位置信息。然而，这些模型无法泛化到超出其训练时所使用的序列长度。我们提出了 YaRN（Yet another RoPE extensioN method），一种计算高效的方法来扩展这类模型的上下文窗口，所需 token 数量比先前方法少 10 倍，训练步数少 2.5 倍。使用 YaRN，我们展示了 LLaMA 模型可以有效地利用并外推到远超其原始预训练所允许的上下文长度，同时在上下文窗口扩展方面也超越了先前的最先进水平。此外，我们证明了 YaRN 具有超越微调数据集有限上下文的外推能力。使用 YaRN 微调的模型已在 https://github.com/jquesnelle/yarn 上公开并可在线复现，上下文长度最高可达 128k。

1 引言

基于 Transformer 的大型语言模型[40]（LLMs）已成为许多自然语言处理（NLP）任务近乎无处不在的选择，在这些任务中，上下文学习（ICL）等长距离能力至关重要。在执行 NLP 任务时，由其训练过程决定的序列最大长度（上下文窗口）已成为预训练 LLM 的主要限制之一。通过少量微调（或无需微调）动态扩展上下文窗口的能力变得越来越重要。为此，Transformer 的位置编码成为讨论的中心。

最初的 Transformer 架构使用绝对正弦位置编码，后来改进为可学习的绝对位置编码[15]。此后，相对位置编码方案[32]进一步提高了 Transformer 的性能。目前，最流行的相对位置编码是 T5 相对偏置[30]、RoPE[34]、XPos[35]和 ALiBi[27]。

位置编码的一个反复出现的局限性是无法泛化到训练期间看到的上下文窗口之外。虽然像 ALiBi 这样的一些方法能够进行有限的泛化，但没有一种方法能够泛化到明显长于其预训练长度的序列[22]。

为了克服这一限制，已经进行了一些工作。[9]和同期[21]提出通过位置插值（PI）轻微修改 RoPE 并使用少量数据进行微调来扩展上下文长度。作为替代方案，[6]提出了"NTK-aware"插值法，考虑了高频信息的损失。此后，针对"NTK-aware"插值法提出了两种改进，各有侧重：

• "Dynamic NTK"插值法[14]，适用于无需微调的预训练模型。
• "NTK-by-parts"插值法[7]，在少量长上下文数据上进行微调时表现最佳。

"NTK-aware"插值法和"Dynamic NTK"插值法已经在开源模型中得到应用，例如 Code Llama[31]（使用"NTK-aware"插值法）和 Qwen 7B[2]（使用"Dynamic NTK"）。 在本文中，除了全面介绍先前未发表的关于"NTK-aware"、"Dynamic NTK"和"NTK-by-part"插值法的工作外，我们还提出了 YaRN（Yet another RoPE extensioN method），这是一种改进的方法，可以有效地扩展使用旋转位置嵌入（RoPE）训练的模型的上下文窗口，包括 LLaMA[38]、GPTNeoX[5]和 PaLM[10]系列模型。 YaRN 在仅使用约原始预训练数据 0.1% 的情况下进行微调后，在上下文窗口扩展方面达到了最先进的性能。同时，通过结合称为动态缩放（Dynamic Scaling）的推理时技术，Dynamic-YaRN 可以在没有任何微调的情况下实现超过 2 倍的上下文窗口扩展。

2 背景与相关工作

2.1 旋转位置嵌入

我们工作的基础是在[34]中介绍的旋转位置嵌入（RoPE）。我们在一个隐藏层上工作，其中隐藏神经元的集合表示为 $D$。给定一个向量序列 $\mathbf{x}_{1}, \cdots, \mathbf{x}_{L} \in \mathbb{R}^{|D|}$，遵循[34]的表示法，注意力层首先将向量转换为查询向量和键向量：

接下来，注意力权重计算如下：

其中 $\mathbf{q}_{m}, \mathbf{k}_{n}$ 被视为列向量，因此 $\mathbf{q}_{m}^{T} \mathbf{k}_{n}$ 仅仅是欧几里得内积。在 RoPE 中，我们首先假设 $|D|$ 是偶数，并将嵌入空间和隐藏状态识别为复向量空间：

其中内积 $\mathbf{q}^{T} \mathbf{k}$ 成为标准埃尔米特内积的实部 $\operatorname{Re}\left(\mathbf{q}^{*} \mathbf{k}\right)$。更具体地说，同构交错实部和虚部：

为了将嵌入 $\mathbf{x}_{m}, \mathbf{x}_{n}$ 转换为查询向量和键向量，我们首先给定 $\mathbb{R}$-线性算子：

在复数坐标中，函数 $f_{q}, f_{k}$ 由下式给出：

其中 $\theta=\operatorname{diag}\left(\theta_{1}, \cdots, \theta_{|D| / 2}\right)$ 是对角矩阵，$\theta_{d}=b^{-2 d /|D|}$ 且 $b=10000$。这样，RoPE 将每个（复数值）隐藏神经元与一个单独的频率 $\theta_{d}$ 相关联。这样做的好处是查询向量和键向量之间的点积仅取决于相对距离 $m-n$，如下所示：

在实数坐标中，RoPE 可以使用以下函数编写： $f_{\mathbf{W}}\left(\mathbf{x}_{m}, m, \theta_{d}\right)=\left(\begin{array}{ccccccc}\cos m \theta_{1} & -\sin m \theta_{1} & 0 & 0 & \cdots & 0 & 0 \\ \sin m \theta_{1} & \cos m \theta_{1} & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m \theta_{2} & -\sin m \theta_{2} & \cdots & 0 & 0 \\ 0 & 0 & \sin m \theta_{2} & \cos m \theta_{2} & \cdots & 0 & 0 \\ 0 & 0 & 0 & 0 & \cdots & \cos m \theta_{l} & -\sin m \theta_{l} \\ 0 & 0 & 0 & 0 & \cdots & \sin m \theta_{l} & \cos m \theta_{l}\end{array}\right) \mathbf{W} \mathbf{x}_{m}$, 使得

2.2 位置插值

由于语言模型通常使用固定的上下文长度进行预训练，因此很自然地会问如何通过在相对较少的数据上进行微调来扩展上下文长度。对于使用 RoPE 作为位置嵌入的语言模型，Chen 等人[9]和同期的 kaiokendev[21]提出了位置插值（PI）来将上下文长度扩展到预训练限制之外。虽然直接外推在序列 $w_{1}, \cdots, w_{L}$（其中 $L$ 大于预训练限制）上表现不佳，但他们发现在预训练限制内插值位置索引在少量微调的帮助下效果很好。具体来说，给定一个使用 RoPE 的预训练语言模型，他们通过以下方式修改 RoPE：

其中 $L^{\prime}>L$ 是超出预训练限制的新上下文窗口。使用原始预训练模型加上修改后的 RoPE 公式，他们在数量级更少的 token（Chen 等人[9]中为几十亿）上进一步微调了语言模型，并成功实现了上下文窗口扩展。

2.3 附加符号

扩展上下文长度与原始上下文长度之间的比率特别重要，我们引入符号 $s$ 定义为：

我们称 $s$ 为比例因子。 我们还将公式 10 重写并简化为以下通用形式：

其中 $g(m), h\left(\theta_{d}\right)$ 是依赖于具体方法的函数。对于 PI，我们有 $g(m)=m / s, h\left(\theta_{d}\right)=\theta_{d}$。在后续章节中，当我们介绍新的插值方法时，有时我们只指定函数 $g(m)$ 和 $h\left(\theta_{d}\right)$。 此外，我们将 $\lambda_{d}$ 定义为 RoPE 嵌入在第 $d$ 个隐藏维度上的波长：

波长描述了 RoPE 嵌入在维度 $d$ 上完成一次完整旋转 $(2 \pi)$ 所需的 token 长度。 鉴于某些插值方法（例如 PI）不关心维度的波长，我们将这些方法称为"盲"插值方法，而其他方法（例如 YaRN）则关心，我们将其归类为"目标"插值方法。

2.4 相关工作

ReRoPE [33] 也旨在扩展使用 RoPE 预训练的现有模型的上下文大小，并声称无需任何微调即可实现"无限"上下文长度。这一说法得到了在 Llama 2 13B 模型上上下文长度增加至 16k 时损失单调递减的支持。它通过修改注意力机制来实现上下文扩展，因此并非纯粹的嵌入插值方法。由于目前它与 Flash Attention 2 [13] 不兼容，并且在推理过程中需要两次注意力传递，我们不将其纳入比较。 与我们的工作同期，LM-Infinite [16] 提出了与 YaRN 类似的想法，但专注于非微调模型的"即时"长度泛化。由于它们也修改了模型的注意力机制，因此它不是嵌入插值方法，并且不能立即与 Flash Attention 2 兼容。

3 方法论

PI 同等地拉伸所有 RoPE 维度，而我们发现 PI [9] 描述的理论插值界限不足以预测 RoPE 与 LLM 内部嵌入之间的复杂动态。在以下小节中，我们描述了我们单独识别并解决的 PI 的主要问题，以便为读者提供我们协同使用的每种方法的背景、起源和理由，以获得完整的 YaRN 方法。

3.1 高频信息丢失 - "NTK-aware"插值

如果我们仅从信息编码的角度来看 RoPE，[36] 中使用神经切线核（NTK）理论表明，如果输入维度较低且相应的嵌入缺乏高频分量，深度神经网络在学习高频信息方面会遇到困难。这里我们可以看到相似之处：token 的位置信息是一维的，而 RoPE 将其扩展为 n 维复向量嵌入。

RoPE 在许多方面与傅里叶特征 [36] 非常相似，因为可以将 RoPE 定义为傅里叶特征的特殊一维情况。不加选择地拉伸 RoPE 嵌入会导致丢失重要的​​高频细节，而网络需要这些细节来解析既非常相似又非常接近的 token（描述最小距离的旋转需要不能太小，以便网络能够检测到它）。 我们假设，在 PI [9] 中看到的，在较大上下文大小上微调后，短上下文大小的困惑度略有增加可能与此问题有关。在理想情况下，没有理由在较大上下文大小上进行微调会降低较小上下文大小的性能。

为了解决插值 RoPE 嵌入时丢失高频信息的问题，[6] 中开发了"NTK-aware"插值法。我们不是将 RoPE 的每个维度都按比例因子 $s$ 等比例缩放，而是通过减少高频缩放和增加低频缩放，将插值压力分散到多个维度上。可以通过多种方式获得这种变换，但最简单的方法是对 $\theta$ 的值进行基数变换。 更准确地说，遵循第 2.3 节中设定的表示法，我们定义"NTK-aware"插值方案如下（推导细节见附录 A.1）。

定义 1 "NTK-aware"插值是对 RoPE 的一种修改，通过使用公式 12 和以下函数实现。

其中

根据[6]的结果，与 PI [9] 相比，该方法在扩展非微调模型的上下文大小方面表现要好得多。然而，该方法的一个主要缺点是，由于它不仅仅是一种插值方案，某些维度会被轻微外推到"超出边界"的值，因此使用"NTK-aware"插值[6]进行微调得到的结果劣于 PI [9]。此外，由于"超出边界"的值，理论比例因子 $s$ 并不能准确描述真实的上下文扩展比例。在实践中，对于给定的上下文长度扩展，比例值 $s$ 必须设置得高于预期比例。

我们注意到，在本文发布前不久，Code Llama [31] 发布了，它通过将基数 $b$ 手动缩放到 1M 来使用"NTK-aware"缩放。

3.2 相对局部距离丢失 - "NTK-by-parts"插值

在像 PI 和"NTK-aware"插值这样的盲插值方法中，我们平等地对待所有 RoPE 隐藏维度（即它们对网络具有相同的影响）。然而，有强烈的线索指向我们需要目标插值方法。

在本节中，我们主要考虑 RoPE 公式中公式 13 定义的波长 $\lambda_{d}$。为简单起见，我们省略 $\lambda_{d}$ 中的下标 $d$，并鼓励读者将 $\lambda$ 视为任意周期函数的波长。

RoPE 嵌入的一个有趣观察是，给定上下文大小 $L$，存在一些维度 $d$，其波长大于预训练期间看到的最大上下文长度 $(\lambda>L)$，这表明某些维度的嵌入可能在旋转域中分布不均匀。在这种情况下，我们假设拥有所有唯一的位置对意味着绝对位置信息保持不变。相反，当波长很短时，网络只能访问相对位置信息。 此外，当我们通过比例 $s$ 或使用基数变换 $b^{\prime}$ 拉伸所有 RoPE 维度时，所有 token 都会彼此更靠近，因为旋转量较小的两个向量的点积更大。这种缩放严重损害了 LLM 理解其内部嵌入之间微小和局部关系的能力。我们假设这种压缩导致模型对邻近 token 的位置顺序感到困惑，从而损害模型的能力。

为了解决这个问题，基于我们发现的两个先前观察，我们选择完全不插值较高频率的维度，而总是插值较低频率的维度。特别是：

• 如果波长 $\lambda$ 远小于上下文大小 $L$，我们不进行插值；
• 如果波长 $\lambda$ 等于或大于上下文大小 $L$，我们只想进行插值并避免任何外推（与之前的"NTK-aware"方法不同）；
• 介于两者之间的维度可以兼具两者，类似于"NTK-aware"插值。

因此，引入原始上下文大小 $L$ 和波长 $\lambda$ 之间的比率 $r=\frac{L}{\lambda}$ 更为方便。在第 $d$ 个隐藏状态中，比率 $r$ 以下列方式依赖于 $d$：

为了定义上述不同插值策略的边界，我们引入了两个额外的参数 $\alpha, \beta$。所有满足 $r(d)<\alpha$ 的隐藏维度 $d$ 是我们按比例 $s$ 进行线性插值的地方（与 PI 完全相同，避免任何外推），而满足 $r(d)>\beta$ 的 $d$ 是我们根本不进行插值的地方。定义斜坡函数 $\gamma$ 为：

借助斜坡函数，"NTK-by-parts"方法可以描述如下。 定义 2 "NTK-by-parts"插值是对 RoPE 的一种修改，通过使用公式 12 和以下函数 ${ }^{4}$ 实现：

$\alpha$ 和 $\beta$ 的值应根据具体情况进行调整。例如，我们通过实验发现，对于 Llama 系列模型，$\alpha$ 和 $\beta$ 的良好值是 $\alpha=1$ 和 $\beta=32$。

使用本节中描述的技术，该方法的一个变体以"NTK-by-parts"插值的名称发布[7]。如[7]所示，这种改进的方法在非微调模型和微调模型上都优于先前的 PI [9] 和"NTK-aware" 3.1 插值方法。

3.3 动态缩放 - "Dynamic NTK"插值

在许多用例中，会执行多个前向传递，序列长度从 1 到最大上下文大小不等。一个典型的例子是自回归生成，其中序列长度在每一步后增加 1。应用使用比例因子 $s$ 的插值方法（包括 PI、"NTK-aware"和"NTK-by-parts"）有两种方式：

1. 在整个推理周期中，嵌入层是固定的，包括比例因子 $s=L^{\prime} / L$，其中 $L^{\prime}$ 是固定的扩展上下文大小。
2. 在每次前向传递中，位置嵌入更新比例因子 $s=\max \left(1, l^{\prime} / L\right)$，其中 $l^{\prime}$ 是当前序列的序列长度。

(1) 的问题在于，模型在长度小于 $L$ 时可能会出现性能折扣，并且当序列长度长于 $L^{\prime}$ 时会出现突然的性能下降。但是通过执行动态

[^1]缩放如 (2)，它允许模型在达到训练的上下文限制 $L^{\prime}$ 时平稳地性能下降，而不是立即崩溃。我们将这种推理时方法称为动态缩放方法。当它与"NTK-aware"插值结合使用时，我们称之为"Dynamic NTK"插值。它首次以 reddit 帖子 [14] 的形式公开出现。

一个值得注意的事实是，"Dynamic NTK"插值在没有进行任何微调的情况下，在预训练于 $L$ 上的模型 $\left(L^{\prime}=L\right)$ 上表现异常出色。附录 B.3 中的实验支持了这一点。 通常在重复的前向传递中，应用 kv-caching [8]，以便我们可以重用先前的键值向量并提高整体效率。我们指出，在某些实现中，当缓存 RoPE 嵌入时，需要注意修改它以适应带有 kv-caching 的动态缩放。正确的实现应该在应用 RoPE 之前缓存 kv 嵌入，因为当 $s$ 改变时，每个 token 的 RoPE 嵌入都会改变。

3.4 YaRN

除了先前的插值技术，我们还观察到，在注意力 softmax 之前对 logits 引入温度 $t$ 对困惑度具有均匀的影响，无论数据样本和 token 在扩展上下文窗口中的位置如何（参见附录 A.2）。更准确地说，我们修改注意力权重的计算，而不是使用公式 2：

将 RoPE 重新参数化为一组 2D 矩阵对于实现这种注意力缩放具有明显的优势：我们可以使用"长度缩放"技巧，通过简单地将复数 RoPE 嵌入缩放相同的量，将 $\mathbf{q}_{m}$ 和 $\mathbf{k}_{n}$ 都缩放一个常数因子 $\sqrt{1 / t}$。有了这个，YaRN 可以有效地改变注意力机制而无需修改其代码。此外，它在推理和训练期间都具有零开销，因为 RoPE 嵌入是预先生成的，并在所有前向传递中重复使用。将其与"NTK-by-parts"插值相结合，我们得到了 YaRN 方法。

定义 3 我们将"YaRN 方法"称为公式 21 中的注意力缩放和第 3.2 节中介绍的"NTK-by-parts"插值的组合。

对于 LLaMA 和 Llama 2 模型，我们建议使用以下值：

上述方程是通过在没有微调的情况下，使用"NTK-by-parts"方法（第 3.2 节）在 LLaMA 7b、13b、33b 和 65b 模型上，将最低困惑度下的 $\sqrt{1 / t}$ 拟合各种比例扩展因子 $s$ 而得到的。我们注意到，相同的 $t$ 值也相当适用于 Llama 2 模型（7b、13b 和 70b）。这表明熵增加的性质和温度常数 $t$ 可能具有一定程度的"普适性"，并且可能在某些模型和训练数据之间具有泛化性。 YaRN 方法结合了我们所有的发现，并在微调和非微调场景中都超越了所有先前的方法。由于其低占用空间，YaRN 允许直接与修改注意力机制的库兼容，例如 Flash Attention 2 [13]。

4 实验

我们展示了 YaRN 成功实现了使用 RoPE 作为其位置嵌入的语言模型的上下文窗口扩展。此外，这一结果仅通过 400 个训练步骤就实现了，约占模型原始预训练语料库的 0.1%，比 Rozière 等人[31]减少了 10 倍，比 Chen 等人[9]减少了 2.5 倍的训练步骤，使其在训练方面具有极高的计算效率，且没有额外的推理成本。我们计算长文档的困惑度并在既定基准上评分 以评估生成的模型，发现它们超越了所有其他上下文窗口扩展方法。

我们大致遵循了[9]中概述的训练和评估程序。

4.1 训练

在训练中，我们扩展了 Llama 2 [39] 7B 和 13B 参数模型。除了按照 3.4 节所述，使用 $s=16$ 和 $s=32$ 计算嵌入频率外，LLaMA 模型架构没有其他更改。

我们使用了 $2 \times 10^{-5}$ 的学习率，无权重衰减，线性预热 20 步，并结合 AdamW [24] $\beta_{1}=0.9$ 和 $\beta_{2}=0.95$。对于 $s=16$，我们使用 PyTorch [26] 全分片数据并行（Fully Sharded Data Parallelism）[42] 和 Flash Attention 2 [13]，在 PG19 数据集 [29] 上进行了 400 步的微调，全局批处理大小为 64，数据集被分块成 64k 的片段，并以 BOS 和 EOS token 包裹。对于 $s=32$，我们遵循相同的程序，但从完成的 $s=16$ 检查点开始，并额外训练了 200 步。

4.2 外推和迁移学习

在 Code Llama [31] 中，使用了具有 16k 上下文的数据集，比例因子设置为 $s \approx 88.6$，对应于 355k 的上下文大小。他们表明，网络可以外推到 100k 的上下文，而无需在训练期间看到这些上下文大小。与 3.1 和 Rozière 等人 [31] 类似，YaRN 也支持使用比数据集长度更高的比例因子 $s$ 进行训练。由于计算限制，我们仅通过在 $s=16$ 模型的基础上使用相同的 64k 上下文数据集进一步微调 200 步来测试 $s=32$。

我们在 4.3.1 中表明，$s=32$ 模型仅使用 64k 上下文进行训练，就成功地外推到 128k 上下文。与之前的"盲"插值方法不同，YaRN 在增加比例 $s$ 时进行迁移学习效率更高。这证明了从 $s=16$ 到 $s=32$ 的成功迁移学习，网络无需重新学习插值嵌入，因为 $s=32$ 模型在整个上下文大小上等同于 $s=16$ 模型，尽管仅在 $s=32$ 上训练了 200 步。

4.3 评估

评估侧重于三个方面：

1. 具有扩展上下文窗口的微调模型的困惑度得分，
2. 微调模型上的密钥检索任务，
3. 微调模型的通用 LLM 基准测试结果，

4.3.1 长序列语言建模

为了评估长序列语言建模性能，我们使用了 GovReport [18] 和 Proof-pile [4] 数据集，这两个数据集都包含许多长序列样本。对于所有评估，仅使用了两个数据集的测试分割。所有困惑度评估均使用 Press 等人 [27] 的滑动窗口方法计算，其中 $S=256$。 首先，我们评估了模型在上下文窗口增加时的表现。我们从 Proof-pile 中随机选择了 10 个至少包含 128k token 的样本，并评估了这些样本在从 2k token 到 128k token 的序列长度下，每隔 2k 步截断时的困惑度。 表 1 显示了通过 PI (LLongMA-2 7b ${ }^{5}$ )、"NTK-aware" 和 YaRN 将 Llama-2 模型从 4096 扩展到 8192 上下文长度的并排比较。请注意，PI 和"NTK-aware"模型使用 Chen 等人 [9] 的方法进行训练，而 YaRN 使用相同的方法，但训练步数和数据量减少了 2.5 倍，如第 4 节所述。

表 1：通过 PI、NTK 和 YaRN 扩展的 Llama-2 上十个 128k Proof-pile 文档的滑动窗口困惑度 ($S=256$)

我们进一步评估了 YaRN 在比例因子 $s=16, 32$ 下的表现，并将其与一些从 Llama-2 微调并扩展到超过 32k 上下文窗口的开源模型进行了比较，例如 Together.ai [37] 和"NTK-aware"Code Llama [31]。结果总结在表 2 中（更详细的图表见图 1）。

模型大小	模型名称	上下文窗口	扩展方法	评估上下文窗口大小
				8192	32768	65536	98304	131072
7B	Together	32k	PI	3.50	2.64	$>10^{2}$	$>10^{3}$	$>10^{4}$
7B	Code Llama	100k	NTK	3.71	2.74	2.55	2.54	2.71
7B	YaRN ( $s=16$ )	64k	YaRN	3.51	2.65	2.42	$>10^{1}$	$>10^{1}$
7B	YaRN ( $s=32$ )	128k	YaRN	3.56	2.70	2.45	2.36	2.37
13B	Code Llama	100k	NTK	3.54	2.63	2.41	2.37	2.54
13B	YaRN ( $s=16$ )	64k	YaRN	3.25	2.50	2.29	$>10^{1}$	$>10^{1}$
13B	YaRN ( $s=32$ )	128k	YaRN	3.29	2.53	2.31	2.23	2.24

表 2：十个 128k Proof-pile 文档（截断到评估上下文窗口大小）的滑动窗口困惑度 ($S=256$)

我们观察到，该模型在整个目标上下文大小上表现出强大的性能，YaRN 插值法是第一个成功将 Llama 2 的有效上下文大小扩展到 128k 的方法。特别值得注意的是 YaRN $(s=32)$ 模型，尽管微调数据长度限制在 64k token，但它们在整个 128k 上下文范围内持续表现出下降的困惑度，这表明模型能够泛化到未见过的上下文长度。

此外，在附录 B.1 中，我们在表 4 中展示了在 32k 最大上下文窗口设置下，对 50 个未截断且每个样本至少包含 16k token 的 GovReport 文档进行评估的平均困惑度结果（未使用动态缩放）。与 Proof-pile 结果类似，GovReport 结果表明使用 YaRN 进行微调在长序列上取得了良好的性能。

4.3.2 密钥检索

[25] 中定义的密钥检索任务衡量模型从大量无意义文本中检索简单密钥（即五位数字）的能力。为了评估模型，我们在 8k 到 128k 的不同上下文窗口大小上，对密钥检索任务进行了 10 次迭代，每次将密钥放置在评估上下文窗口中随机均匀分布的位置。使用 YaRN 微调的 7b 和 13b 模型在 128k 上下文大小下都以非常高的准确率（$>99\%$）通过了整个上下文窗口大小内的密钥检索任务。我们在附录 B.2 中展示了详细结果。

4.3.3 标准化基准测试

Hugging Face Open LLM 排行榜 [19] 在一组标准化的四个公共基准上比较了众多 LLM。具体来说，我们使用了 25-shot ARC-Challenge [11]、10-shot HellaSwag [41]、5-shot MMLU [17] 和 0-shot TruthfulQA [23]。

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

模型大小	模型名称	上下文窗口	扩展方法	ARC-c	Hellaswag	MMLU	TruthfulQA
7B	Llama 2	4k	无	53.1	77.8	43.8	39.0
7B	Together	32k	PI	47.6	76.1	43.3	39.2
7B	Code Llama	100k	NTK	39.9	60.8	31.1	37.8
7B	YaRN ( $s=16$ )	64k	YaRN	52.3	78.8	42.5	38.2
7B	YaRN ( $s=32$ )	128k	YaRN	52.1	78.4	41.7	37.3
13B	Llama 2	4k	无	59.4	82.1	55.8	37.4
13B	Code Llama	100k	NTK	40.9	63.4	32.8	43.8
13B	YaRN ( $s=16$ )	64k	YaRN	58.1	82.3	52.8	37.8
13B	YaRN ( $s=32$ )	128k	YaRN	58.0	82.2	51.9	37.3

表 3：上下文窗口扩展方法在 Hugging Face Open LLM 基准套件上的性能与原始 Llama 2 基线的比较

我们观察到 YaRN 模型与其各自的 Llama 2 基线之间性能下降极小。我们还观察到 YaRN $s=16$ 和 $s=32$ 模型之间的得分平均下降了 $0.49\%$。由此我们得出结论，从 64k 到 128k 的迭代扩展导致的性能损失可以忽略不计。

5 结论

总之，我们已经证明 YaRN 改进了所有现有的 RoPE 插值方法，并且可以作为 PI 的直接替代品，没有任何缺点且实现工作量极小。微调后的模型在多个基准上保留了其原始能力，同时能够关注非常大的上下文大小。此外，YaRN 允许通过在较短数据集上进行微调来实现高效外推，并可以利用迁移学习来加速收敛，这在计算受限的场景下都至关重要。最后，我们展示了 YaRN 外推的有效性，它能够"短时训练，长时测试"。

6 可复现性

为了有助于复现，我们作为补充材料提供了用于训练表 2 中 YaRN 模型的全部代码，以及生成图 1 和表 [1、2、3、4 和 5] 的评估代码。该代码还包含论文中引用的各种扩展方法的实现。为了训练 YaRN，我们使用了公开可用的 PG19 数据集 [29]，并将其 token 化为 64k token。