从 ROPE 到 Yarn：通用公式解析长文本大模型的位置编码

ROPE 的核心亮点是每个隐藏维度的旋转频率由 $\theta$ 决定。

$\theta$ 定义为：

$$\theta_m = \frac{1}{10000^{2m/d}}$$

这里：

• $d$ 是一个固定的基数（通常取 10000）。
• $d$ 是隐藏层的维度数。
• $m$ 表示隐藏层的某一具体维度。

频率参数 $\theta$ 决定了每个维度的旋转速度：

• 对于低维度（即 $m$ 值较小的维度），$\theta$ 更接近于 1。
• 对于高维度（即 $m$ 值较大的维度），$\theta$ 会迅速衰减到接近 0。

旋转角度 $\theta$ 决定了频率：

• 当 $\theta$ 接近 1（低维度），变化较大，旋转更快。
• 当 $\theta$ 接近 0（高维度），变化较小，旋转更慢。

因此低维度具有更快的旋转（对应局部细节捕捉），高维度具有更慢的旋转（对应长距离依赖）。这种设计巧妙地结合了长距离和短距离的信息编码能力。

2.3 为什么 ROPE 直接外推在长文本外推中受限？

尽管 ROPE 在预训练窗口范围内表现优异，其主要限制在于：

• 频率不变性：在预训练时被固定，无法适应更长的上下文长度。
• 频率分布的刚性：所有维度的频率分布固定，不支持动态调整，导致当序列长度超出预训练范围时，旋转编码出现混乱。

当上下文窗口从预训练的 $L$ 扩展到 $N$ 时，相对位置 $j-i$ 的值可能远超预期范围。此时，旋转频率无法捕捉新的位置信息，导致模型性能显著下降。

3. Position Interpolation (PI)：均匀拉伸的位置插值

PI 尝试通过重新定义 $i$ 将位置索引拉伸到预训练窗口内：

$$i' = \frac{i}{\alpha}, \quad \alpha = \frac{N}{L}$$

其中 $\alpha$ 是上下文扩展比例。

$i'$ 保持不变，保持频率参数不变。

这使得每个位置索引被均匀拉伸到预训练窗口内，公式变为：

$$\text{RoPE}(x_{i'}, i')$$

优点：简单有效，训练开销小。 缺点：对所有维度频率统一缩放，导致高频信息丢失，影响局部关系建模。

4. NTK-Aware Interpolation：非均匀频率缩放

NTK-Aware 方法对 $\theta$ 引入动态调整，$i$ 保持不变：

$$\theta'_m = \theta_m \cdot \alpha^{\frac{m}{d-1}}$$

这里 $\alpha=1$ 时，$\theta'_m = \theta_m$，有没有 $\alpha$ 都不影响，因此也叫直接外推。

这里 $\alpha > 1$ 时，$\theta'_m$ 变小，变成了线性内插。

所以该方法可以理解为介于直接外推和线性内插之间的平滑方法。它主要通过对低频维度更大幅度地插值，高频维度保持不变，保留了更多高频细节。

优点：在未微调模型中显著提升长文本建模能力。 缺点：部分频率超出预训练范围，可能引发性能不稳定。

5. Dynamic Scaling：动态适配插值比例

Dynamic Scaling 进一步动态调整 $\alpha$ 以适配不同上下文长度：

$$\alpha = \left( \frac{N}{L} \right)^{\frac{d-m}{d}}$$

或者采用指数形式调整：

$$\alpha = \left( \frac{N}{L} \right)^{a \cdot \frac{m}{d} + b}$$

其中，参数 $a$ 和 $b$ 的取值决定了缩放函数的动态灵活性。且 $a,b$ 需要满足约束，且 $a$ 的推荐范围为 $[0, 1]$，公式可有效提升上下文适配能力。

相比于传统线性内插方法，上述公式差不多可以实现 1%-2% 的性能增益，其优点在于：通过 $a$ 和 $b$ 调整缩放范围，位置索引从 $i$ 被映射到更合理的区间范围，避免传统方式中索引越大越不充分的问题。

6. NTK-by-parts Interpolation：基于波长局部分段插值

基于波长的维度分类：'NTK-by-parts'方法着重考虑 RoPE 公式中定义的波长。对于每个隐藏维度，计算其波长（其中，$\lambda_m = \frac{2\pi}{\theta_m}$）。根据波长与原始上下文大小的比值，将隐藏维度分为不同类别。

不同类别维度的插值策略：

• 高频维度：对于波长明显小于上下文大小的维度（即 $\lambda_m < L$），这些维度被认为包含较多高频信息，采用线性插值，且插值比例为尺度因子。这种方式类似于 PI 方法，但避免了对高频信息的过度拉伸，从而减少高频细节的丢失。
• 低频维度：当波长大于等于上下文大小（$\lambda_m \ge L$）时，认为这些维度对绝对位置信息更敏感，因此不对这些维度进行插值，以保持绝对位置信息的完整性，避免因插值导致的位置信息混淆。
• 中间维度：对于介于上述两种情况之间的维度，采用一种混合策略，通过引入一个 ramp 函数来确定插值程度。具体而言，这些维度的插值是在原始频率和拉伸后的频率之间进行线性插值，插值程度由 $ramp$ 决定。

根据上述策略，'NTK-by-parts'方法对 RoPE 的修改可以用以下公式表示：定义 $\theta'_m$，保持位置索引不变；$\theta'_m$ 根据维度的不同类别对频率参数进行相应的变换。其中，ramp 函数定义为：

$$ramp(x) = \max(0, \min(1, x))$$

核心优势：与之前的 PI 和'NTK-aware'插值方法相比，'NTK-by-parts'方法在处理 RoPE 维度时有更强的针对性。PI 方法对所有维度同等插值，容易丢失高频信息；'NTK-aware'方法虽然尝试通过改变频率缩放方式来缓解问题，但会导致某些维度的外推，产生'越界'值，影响模型性能。而'NTK-by-parts'方法通过根据波长区分维度并采用不同插值策略，能够更好地平衡高频信息保留和位置关系理解，实验中可以表现得更好。

7. Yarn

YaRN 是基于 NTK-aware 方法的进一步拓展，通过结合温度缩放和 NTK-by-parts 插值技术，全面提升长文本外推能力。它核心解决的问题是线性内插导致的 self-attention 点积的值增大。由于线性内插会改变旋转向量转动的幅度，原来距离较远的 q,k 点积由于旋转幅度变小，他们的点积结果会增大，进而导致 Softmax 操作过于'锐化'，使得注意力分布集中于少数位置，削弱模型对全局上下文的关注能力。Yarn 在 NTK-by-parts 基础上，引入注意力温度因子 $\tau$ 来调整注意力分布：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k} \cdot \tau}\right)V$$

优点：

• 超低训练成本（0.1% 预训练数据）。
• 几乎兼容目前所有主流的 Transformer 实现。
• 性能优越，无论是否微调均能在 128k 上下文中表现出色。

8. 实验

最后展示 Yarn 中的实验结果。Yarn 无论在资源利用率还是 128K 长度性能上都超过其他 PI、NTK 类方法，无怪 Yarn 成为目前 Long Context LLM 的标配。实验表明，在长文本任务上，Yarn 能够有效维持模型的 perplexity 和准确率，解决了传统外推方法在长窗口下的性能崩塌问题。