YaRN: Efficient Context Window Extension of Large Language Models

POSITION ENCODING CONTEXT EXTENSION Peng et al. (Nous Research) · 2023
arXiv:2309.00071

上一代的问题：RoPE 无法外推到训练长度之外 — Position Interpolation (PI) 对所有维度均匀拉伸，抹除了高频分量的旋转差异，导致模型在 s > 8 时无法区分相邻 token 的位置。

这代改了什么：NTK-by-parts 通过波长与训练长度的比值将 RoPE 维度分为三类：编码绝对位置的维度做插值，编码相对位置的维度不做插值，中间维度平滑过渡。同时引入 Attention Temperature Scaling 补偿变长上下文导致的注意力分布扁平化。

效果：仅需约 400 步微调（约原始训练数据的 0.1%）即可将 LLaMA 7B/13B/33B/65B 扩展到 128K 上下文。Llama 3 从 8K 到 128K 的扩展直接采用 YaRN。长文本困惑度接近理论最优，短文本基准无损。

一、核心背景

1.1 RoPE 的外推困境

旋转位置编码 (Rotary Position Embedding, RoPE) 是目前大语言模型最广泛采用的位置编码方案。其核心思想是将位置信息通过旋转矩阵注入 attention 计算：对位置 m 和 n，query 和 key 的内积只依赖于它们的相对位置 m - n。具体地，RoPE 将第 d 维的旋转角度定义为 θ_d = b^-2d/|D|，其中 b 是 base（通常取 10000），|D| 是总维度数。

RoPE 的优雅之处在于它提供了 相对位置编码 的天然实现：维度越高（d 越大），θ_d 越小，旋转越慢，波长越长。低维的快速旋转编码精细的相对位置信息，高维的慢速旋转编码粗略的绝对位置信息。这种多分辨率结构在训练长度内工作完美——但在训练长度之外，LLM 无法泛化。

关键洞察：RoPE 的外推失败并非因为位置编码本身超出范围，而是因为模型从未见过某些 旋转角度组合。高维度的长波长分量在训练长度内从未完成一个完整周期，低维度的短波长分量虽然完成了多个周期，但模型只在训练长度覆盖的相位区间上见过它们。超出这个区间，模型遇到的是从未见过的相位模式。

1.2 Position Interpolation (PI) 及其局限

Position Interpolation (Chen et al., 2023; kaiokendev) 是第一个系统的上下文扩展方案。其思路直接：如果要把上下文从 L 扩展到 L'，定义缩放因子 s = L'/L，然后将所有位置索引除以 s，使得最大位置 L' 被压缩回训练长度 L 内：

f'(m, d) = f(m/s, d)

公式 1: Position Interpolation — 对所有维度统一缩放

PI 的优势是简单直观：既然模型没见过超出 L 的位置，那就把所有输入都映射回 [0, L) 区间。但它有一个致命问题：均匀缩放抹除了高频信息。

考虑低维度（d 小，θ_d 大）的快速旋转分量：它们在训练时每一两个 token 就能完成一次完整旋转，因此编码的是高度精细的 相对位置 信息——相邻 token 间的细微距离。当对所有维度统一除以 s 后，这些高频分量的旋转速度被整体减慢，相邻 token 间的角度差异被压缩。当 s 较大（如 8 倍以上）时，低维分量的分辨率严重退化，模型无法区分相邻位置。

实验观察：PI 在 s ≤ 8 时效果尚可，但 s > 8 后困惑度急剧上升。这本质上是因为高频信息被均匀稀释到了 "不该被稀释" 的维度上——那些原本编码相对位置的维度被迫去编码更宽的范围，丧失了精细分辨能力。

二、核心创新

YaRN (Yet another RoPE extensioN) 的核心贡献是认识到：并非所有 RoPE 维度都需要被插值。波长 λ_d 与训练长度 L 的关系决定了每个维度在模型中扮演的角色——长波编码绝对位置，短波编码相对位置。对它们施加相同的处理既不正确也不必要。

2.1 高频信息丢失的深层原因

要理解为什么 PI 的高频退化不可避免，需要从神经正切核 (Neural Tangent Kernel, NTK) 的理论视角来看。NTK 理论揭示了深度神经网络在学习过程中的一个根本特性：低频偏好——网络倾向于优先学习输入中的低频分量，高频分量需要更多的训练样本才能被充分学习。

RoPE 的构造方式恰好类似于 傅里叶特征映射 (Fourier Features)：它将一维的位置索引 m 映射到 |D| 维的复平面嵌入，每个维度对应一个特定频率的旋转。这种映射将位置信息的 "分辨率" 分散到了不同频率的信道上。

当 PI 对所有维度进行均匀缩放时，其实是在做 低频重采样——所有频率被同步拉低。这相当于在傅里叶空间中丢弃了高频分量：因为原来对应高频的旋转角度被压缩后，相邻位置间本应存在的快速相位变化被抹平了。这正是 NTK 理论所警示的情况——神经网络在低维输入下难以学习高频模式，而均匀插值正好削弱了这些模式。

NTK-aware 插值（即 "NTK-aware" 方法）试图缓解这一问题：不采用均匀缩放，而是对高频维度缩放少、低频维度缩放多，从而 "分散压力"。具体做法是修改 base b：

b' = b · s^|D|/(|D|-2)

公式 2: NTK-aware base scaling — 高频少缩，低频多缩

但这个方案存在一个根本问题：最佳 base 值难以确定，且单一超参数无法完美适应所有维度——高频和低频之间存在一个连续的频谱，而 NTK-aware 只是整体偏移了这个频谱，没有区分不同维度各自的信息角色。

2.2 NTK-by-parts：核心突破

YaRN 的突破性洞察来自一个简单的观察：每个 RoPE 维度的波长 λ_d 与训练长度 L 的比值，决定了该维度编码的是 绝对位置 还是 相对位置。

RoPE 维度 d 的波长定义为：

λ_d = 2π / θ_d = 2πb^2d/|D|

公式 3: RoPE 维度 d 的波长

关键比值：

r(d) = L / λ_d

公式 4: 维度 d 在训练长度内完成的旋转周数

r(d) 衡量维度 d 在预训练长度 L 内完成了多少次完整旋转。这个比值决定了维度的行为：

当 λ_d > L（即 r(d) < 1）：该维度在训练长度内从未完成一次完整旋转。因为模型只见过该维度在半个周期内的行为，它只能通过相位的绝对偏移来推断位置——也就是说，这个维度编码的是绝对位置信息。如果要扩展上下文，它必须被插值。
当 λ_d << L（即 r(d) >> 1）：该维度在训练长度内完成了多次完整旋转。模型见过该维度在所有相位上的表现，因此它学习到的是相位的相对变化而非绝对值——这个维度编码的是相对位置信息。对它做插值会破坏这种精细的相对分辨能力，它绝对不能插值。
中间区域：部分维度介于两者之间，需要平滑过渡。

基于这个洞察，YaRN 引入了一个 三段式分区策略，通过两个超参数 α 和 β（默认 α=1, β=32）来定义三个区域：

低比值区域 r(d) < α：波长大于训练长度 → 编码绝对位置 → 做线性插值。
高比值区域 r(d) > β：波长远小于训练长度 → 编码相对位置 → 不做插值。
中间区域 α ≤ r(d) ≤ β：平滑过渡。

用斜坡函数 (ramp function) 来定义插值权重：

γ(r) =
  0, if r < α
  1, if r > β
  (r - α)/(β - α), otherwise

公式 5: Ramp 函数 — 定义插值权重

最终每个维度的旋转角调整为：

h(θ_d) = (1 - γ(r(θ_d))) · θ_d/s + γ(r(θ_d)) · θ_d

公式 6: YaRN 的维度自适应旋转角调整

其中 γ = 0 对应完全插值（低频/长波维度），γ = 1 对应完全不插值（高频/短波维度）。这种设计使得每个维度根据自身的信息角色被区别对待，保留了高频维度的相对位置分辨能力。

直觉理解：想象一个钟表盘。时针（低频率）在 12 小时内只走一圈，你需要知道它指向几点来推断时间——这类似于绝对位置编码。秒针（高频率）每秒钟都在快速转动，你看不到具体数字，但可以通过它的连续运动感知时间流逝——这类似于相对位置编码。如果强行把整个表盘均匀压扁，秒针的快速运动信息就会丢失。NTK-by-parts 的做法是：保留秒针的转速不变（不插值），只调整时针的读数（插值）。

2.3 Attention Temperature Scaling

YaRN 的第二个重要创新是 Attention Temperature Scaling。这个改进源自一个被忽视的问题：随着上下文长度的增加，attention softmax 的分布会变得更平坦。

标准的 scaled dot-product attention 为：

softmax(q^Tk / √d)

公式 7: 标准 attention 计算

当上下文从 L 扩展到 L'，每个位置的 attention 需要分布在更多 token 上。理论上，如果模型保持相同的注意力模式（即关注相同的相对位置），那么 softmax 的归一化常数会变大，导致每个位置的注意力权重降低，整体分布变得更均匀——这表现为 注意力熵的增加。

更高的注意力熵意味着模型在区分相关和无关 token 时能力下降——信息的 "聚焦度" 被稀释了。直观上，这就像一个本来只照亮几个点的探照灯，被拉宽后照亮了整个舞台。

YaRN 的解决方案简洁而有效：在 softmax 中引入一个温度参数 t：

softmax(q^Tk / (t · √d))

公式 8: Attention Temperature Scaling — t > 1 降低温度

t > 1 时，softmax 的输入被缩小（等效于降低温度），使得输出分布更尖锐——即模型能更有选择性地关注少数关键 token。注意这里的 "温度" 使用与物理学惯例相反：在统计力学中，温度越低分布越集中；在 YaRN 的语境中，t > 1 相当于 除以更大的数，使得 softmax 分布更集中。

通过实验，YaRN 的作者拟合出温度与缩放因子的经验关系：

√(1/t) = 0.1 · ln(s) + 1

公式 9: Attention Temperature 经验拟合（LLaMA 7B/13B/33B/65B）

这个公式的含义是：随着缩放因子 s 增大，温度 t 需要增大（即 softmax 除以更大的数），且增长趋势与 ln(s) 成正比。这符合直觉——上下文越长，注意力分布被稀释得越严重，就需要更强的 "聚焦" 补偿。

有趣的事实：这个温度拟合在 LLaMA 全系列（7B 到 65B）上表现一致，说明注意力熵的变化是上下文扩展的 固有特性，而非特定模型的 artifact。这意味着 temperature scaling 很可能对所有 RoPE-based LLM 都适用。

2.4 Dynamic Scaling

传统的上下文扩展方法（包括 PI 和原始的 NTK-aware）都使用 固定的 缩放因子 s = L'/L。这意味着模型在训练时看到的上下文始终是 L' 长度，推理时也只能在接近 L' 的长度下表现良好。如果输入的实际长度小于 L'，模型会因为位置编码被过度缩放而出现性能退化。

YaRN 提出了 Dynamic Scaling：不再使用固定 s，而是在每次前向传播时动态计算：

s = max(1, l' / L)

公式 10: Dynamic Scaling — 根据实际序列长度调整

其中 l' 是当前序列的实际长度。这意味着：

短序列（l' ≤ L）使用 s=1，完全不需要插值。
长序列（l' > L）使用 s = l'/L，自适应地缩放。
模型在任何长度下都能获得最合适的位置编码。

Dynamic Scaling 带来的一个关键优势是 优雅降级(graceful degradation)：当序列长度超出训练长度时，模型不会突然崩溃，而是随长度增加逐步退化。相比之下，固定缩放方法在超出预期长度时往往出现非线性的性能崩溃。

Dynamic NTK 是 Dynamic Scaling 与 NTK-aware 插值的结合，结合了 "自适应缩放" 和 "频率差异化缩放" 两者的优点。虽然 YaRN 的核心是 NTK-by-parts，但 Dynamic NTK 作为一种轻量替代也在实践中被广泛采用。

三、训练细节

YaRN 的一大亮点是它的训练极其高效。具体配置如下：

超参数	配置
学习率	2e-5
优化器	AdamW (β1=0.9, β2=0.95)
训练步数	400 步
全局批大小	64
总训练 token 数	约 1.6B (64k × 64 × 400)
数据来源	PG19 数据集分割为 64k 段
实现优化	Flash Attention 2, FSDP
硬件	单节点 8x A100 (80GB)

训练数据量仅约 16 亿 token，而 LLaMA 65B 的原始预训练数据为约 1.4 万亿 token。换句话说，YaRN 的微调只需要原始训练的 约 0.1% 的数据量，400 步微调即可完成上下文从 2K 到 128K 的扩展。这个效率在当时是前所未有的。

这么高的效率从何而来？关键原因在于 YaRN 没有改变模型的核心能力——它没有学习新的语言知识，而是在告诉模型 "同一组 RoPE 维度在新的尺度下如何解析"。NTK-by-parts 已经将位置编码调整到了接近正确的形态，微调只是让模型适应这种微调后的编码模式。本质上，这是一个 表征对齐 的过程，而不是 知识注入 的过程。

对比：PI 要达到相似的长上下文性能通常需要 5000-10000 步微调，而 YaRN 只需 400 步。差异的来源正是 PI 破坏了高频信息，模型需要额外花时间去适应这种信息损失带来的表征变化。YaRN 从一开始就保留了这些信息，因此对齐过程更加高效。

四、实验结果

4.1 长文本困惑度

在 Proof-pile 128k 数据集上测量不同上下文长度下的困惑度，这是评价上下文扩展效果的核心指标：

模型	s=16 @ 8k	@ 16k	@ 32k	@ 64k	@ 128k
LLaMA 7B YaRN	3.51	2.99	2.65	2.42	2.37
LLaMA 13B YaRN	3.25	2.79	2.50	2.29	2.24

一个有趣的现象：随着上下文长度增加，困惑度持续下降。这是因为更长的上下文提供了更多的条件信息，降低了下一个 token 的不确定性——这正是扩展上下文的收益所在。注意困惑度在 128K 时仍然在下降，说明 YaRN 没有在极限长度处出现过拟合或退化。

将 YaRN 与 PI 和 NTK-aware 对比，在 s=32（64K 上下文对 2K 训练长度）时：

PI：困惑度严重退化（高频信息丢失），在 64K 处比 YaRN 差约 0.3-0.5 PPL。
NTK-aware：比 PI 好但仍有差距，特别是在中长序列（16K-32K）区间。
YaRN：在所有长度上最优，且差距随着上下文增长而扩大。

4.2 短文本基准无损

上下文扩展的另一个关键指标是 短文本性能保持——扩展后的模型不应该在原来的任务上退化。以下是 LLaMA 7B 在四个标准短文本基准上的对比：

方法	ARC-C	HellaSwag	MMLU	TruthfulQA
原始 LLaMA 7B	51.0	77.8	35.7	34.3
YaRN (扩展后)	48.1	77.2	30.0	35.1

可以看到，YaRN 在 HellaSwag 和 TruthfulQA 上几乎完全保持了原始性能，ARC-C 下降约 3 个百分点，MMLU 下降约 5.7 个百分点。这种轻微退化是预期的，因为 400 步的高学习率微调不可避免地会对权重产生一些扰动。

值得注意的是，13B 模型 的短文本保留效果通常优于 7B，说明更大的模型对位置编码变化的鲁棒性更强。但 YaRN 相比 PI 在短文本保留上有明显优势——PI 在 7B 上的 MMLU 从 35.7 降至约 26-28，而 YaRN 保持了 30.0。

4.3 消融实验

YaRN 的消融实验验证了三个核心组件的各自贡献：

移除 Temperature Scaling：在 64K 长度上困惑度上升约 0.1-0.2，短文本基准有小幅下降。说明 temperature scaling 对长文本性能有稳定贡献。
使用均匀插值替代 NTK-by-parts（即退化为 PI）：性能大幅下降，特别是在长序列上差距显著。这验证了 NTK-by-parts 是 YaRN 的核心。
移除 Dynamic Scaling：在中等长度区间（如 4K-16K）出现性能退化，但极限长度（128K）受影响较小。说明动态缩放对鲁棒性有贡献。

五、上下游关联

5.1 前驱工作

RoPE (Rotary Position Embedding) — Su et al., 2021。RoPE 是 YaRN 的基础，也是当前 LLM 的事实标准位置编码。RoPE 通过旋转矩阵将相对位置信息编码到 attention 计算中，其多频率结构为 YaRN 的波长分析提供了理论前提。没有 RoPE 的频率分解，NTK-by-parts 的分区策略就无从谈起。

Position Interpolation (PI) — Chen et al., 2023; kaiokendev。PI 是第一个系统的上下文扩展方法，证明了通过微调可以对 RoPE 进行 "重缩放"。虽然 PI 在 s > 8 时遇到瓶颈，但它为 YaRN 提供了 "插值基线的概念验证"——证明了这个方向是可行的，只是需要更精细的粒度控制。

NTK-aware Interpolation — NTK-aware 是 PI 和 YaRN 之间的过渡方法，通过修改 base 来差异化缩放不同频率的维度。它首次引入了 "频率维度应区别对待" 的思想，但缺乏理论依据来决定具体如何区分。YaRN 的 NTK-by-parts 正是对这个问题的回答。

5.2 后继工作与工业采用

Llama 3 的 128K 上下文扩展 — Meta 在 Llama 3 技术报告中明确表示，将 Llama 3 从 8K 扩展到 128K 采用了 YaRN。这是 YaRN 最高调的工业验证——Meta 的生产级模型验证了 YaRN 的可靠性和可扩展性。Llama 3 的实现细节包括在长上下文微调中使用 YaRN 结合特定数据处理策略。

Code Llama — Meta 的代码专用模型 Code Llama 使用 NTK-aware 的变体，将 base 设为 1,000,000 而非默认的 10,000，并采用 ABF (Adjusted Base Frequency) 策略。Code Llama 中的 "NTK-aware ABF" 与 YaRN 共享相似的动机——频率差异化缩放——但实现方式不同：Code Llama 通过增大 base 来整体拉长所有维度的波长，而 YaRN 通过分区策略实现更精细的控制。

Qwen 系列 — 阿里云的 Qwen (通义千问) 模型使用 Dynamic NTK 作为其上下文扩展策略。Dynamic NTK 将 YaRN 中的 Dynamic Scaling (自适应缩放) 与 NTK-aware (频率差异化) 结合，但不包含 NTK-by-parts 的三段式分区。Qwen 2 进一步将上下文长度扩展到 128K，Dynamic NTK 在其中起到了关键作用。

模型	基础上下文	扩展上下文	使用的方法
LLaMA 2	4K	—	原始 RoPE
LLaMA 3	8K	128K	YaRN
Code Llama	16K	100K	NTK-aware ABF
Qwen / Qwen 2	8K / 32K	32K / 128K	Dynamic NTK
Mixtral 8x7B	32K	—	RoPE (参考扩展方法)

六、个人思考

6.1 为什么 YaRN 成功了？

YaRN 的成功根源于它对 RoPE 结构的 信息论尊重。它不是简单地将问题视为 "如何把位置映射到更大的范围"，而是深入理解了每个频率维度在模型中的信息角色。它的核心贡献可以总结为：

频率即信息信道：每个 RoPE 维度是一个独立的频率信道，传递不同粒度的位置信息。NTK-by-parts 识别出这些信道分为两类——绝对位置信道（低频）和相对位置信道（高频），并分别处理。
保留比扩展更重要：在上下文扩展中，最危险的不是 "未能覆盖新位置"，而是 "破坏了已有的表示能力"。YaRN 的克制——对高频维度不做插值——正是体现了保留已有能力优先于覆盖新长度的原则。
理论驱动的设计选择：NTK-by-parts 的三个分区不是经验调参的结果，而是从波长的数学定义出发的推导。这种理论驱动的设计使得 YaRN 在不同的模型系列（7B 到 65B）上表现一致，无需逐个模型调参。

6.2 Attention Temperature 的深层意义

Attention Temperature Scaling 是 YaRN 最被低估的贡献。绝大多数上下文扩展工作关注的是 "位置编码如何适配新长度"，但忽视了上下文长度增长对 attention 机制本身的直接影响。随着上下文变长，每个位置的 attention 注意力池扩大，softmax 的熵自然增加。这是一个 与位置编码无关的根本性问题——即使位置编码完美，注意力熵的增加也会降低模型对关键信息的聚焦能力。

这也暗示了一个更一般化的原理：位置编码的精度和注意力的分配效率是上下文扩展的两个正交维度。YaRN 同时解决了这两个问题——NTK-by-parts 处理精度，Temperature Scaling 处理效率，两者缺一不可。未来的上下文扩展工作如果能同时考虑这两个维度，可能会取得更好的效果。

6.3 位置编码的后续方向

YaRN 证明了 RoPE 具备极强的可扩展性——只要理解其频率结构，RoPE 可以远远超出其训练长度。但这是否意味着 RoPE 是位置编码的终极方案？

近期的研究展示了一些替代方案：

ALiBi (Press et al., 2022)：通过引入与距离相关的线性偏置来实现外推，不需要位置嵌入。但 ALiBi 在长上下文任务上的表现通常不如 RoPE + 扩展方法。
No Position Encoding (NoPE) (Kazemnejad et al., 2023)：因果注意力本身隐式编码了位置信息，完全移除显式位置编码在某些场景下可行。
CoPE (Contextual Position Encoding) (Golovneva et al., 2024)：根据内容来决定位置编码，而不是基于固定的索引。

但至少在目前（2024-2025），RoPE 仍然是 LLM 的事实标准，而 YaRN 是 RoPE 上下文扩展的最优方案之一。只要 RoPE 继续主导 LLM 架构，YaRN 的思路和方法论就会持续具有价值。

6.4 方法论启示

YaRN 的研发历程展示了一种值得学习的学术方法论：从物理/数学的已有理论（波长、频率、傅里叶分析）出发，应用到深度学习系统的工程问题中。不是盲目尝试各种插值策略，而是从 "每个频率分量的作用是什么" 这个本源问题出发，推导出自然的解决方案。

这种方法论在深度学习历史上曾被多次验证：Batch Normalization 从信号处理中的白化概念出发；ResNet 从动力系统的欧拉离散化出发；而 YaRN 从波与波长的物理图像出发。当一个问题有清晰的数学/物理结构时，从第一性原理推导往往比经验调参产生更优雅和更通用的解。