YaRN：解密扩展大模型上下文的方法

1. 深入理解旋转位置编码（RoPE）

随着像 GPT-4、Llama、Mistral 等模型的涌现，我们对它们的能力叹为观止。但一个长久以来的痛点始终存在：它们能处理的文本长度是有限的。这个限制，我们称之为“上下文窗口”（Context Window）。上下文窗口决定了模型在一次推理中能“记住”的最大文本长度。

扩展上下文窗口，意味着模型可以“读”得更长，“记”得更多，从而能够处理更长的文档、进行更复杂的对话、编写更完整的代码。而 YaRN (Yet another RoPE extensioN method) 就是一个非常亮眼的解决方案，它能够在很少的微调（甚至零微调）下，将 Llama 等模型的上下文窗口扩展数倍，同时保持很高的性能。

要理解 YaRN，我们必须先深入了解它所改进的核心技术：**旋转位置编码（Rotary Position Embedding, RoPE）

1.1 深入理解旋转位置编码（RoPE）

之前有文章简单介绍过旋转位置编码（RoPE）《transformer(三)旋转位置编码RoPE详解》
由于YaRN是基于RoPE的，必须先讲RoPE的原理，才能理解YaRN的扩展上下文窗口的原理。

RoPE 的简单引入

相对位置编码它的核心思想是，在计算注意力时，只关心 query token 和 key token 之间的相对距离。而 RoPE 就是其中最优雅和高效的实现之一。

而RoPE 的思想非常巧妙。它没有直接将位置信息“加”到词嵌入上，而是通过旋转 query (q) 和 key (k) 向量来注入位置信息。

定义 $q_m,k_m,v_m$ 分别为 query, key, value 向量在位置 $m$ 的表示
用 $R_m$ 表示位置 $m$ 处的旋转矩阵，讲旋转矩阵嵌入到 $q,k$ 中，计算注意力分数时

$(R_{m}q{)}^T(R_{n}k)=q^T{R}_{n-m}k$

内积的结果中，query $q_m$ 被一个只依赖于相对位置 $n-m$ 的旋转矩阵 $R_{n-m}$ 变换后，再与 key $k_n$ 进行内积。这就完美地将相对位置信息融入了注意力计算中。

1.2 完整的旋转矩阵

之前写的$R_m$ 表示位置 $m$ 处的旋转矩阵，其实里面还有个$\theta$，表示旋转的角度。这个角度是根据位置 $m$ 计算得到的，这个角度是不需要学习的，根据具体公式计算即可:

${\theta }_k=b^{-2k/d}$

其中 $b=10000$，$d$ 是模型的维度（如1024维），$k$ 是维度的索引。

这个角度的计算，确保了不同位置之间的旋转角度是不同的，从而能够注入不同的位置信息。

完整的第m个位置（m是上下文的位置索引，上下文长度可能是32k）上的RoPE矩阵应该为

$$R_{\theta ,m}=\left[\begin{array}{ccccccc}\cos (m{\theta }_1)& -\sin (m{\theta }_1)& 0& 0& ...& 0& 0\\ \sin (m{\theta }_1)& \cos (m{\theta }_1)& 0& 0& ...& 0& 0\\ 0& 0& \cos (m{\theta }_2)& -\sin (m{\theta }_2)& ...& 0& 0\\ 0& 0& \sin (m{\theta }_2)& \cos (m{\theta }_2)& ...& 0& 0\\ ...& ...& ...& ...& ...& ...& ...\\ 0& 0& 0& 0& ...& \cos (m{\theta }_{d/2})& -\sin (m{\theta }_{d/2})\\ 0& 0& 0& 0& ...& \sin (m{\theta }_{d/2})& \cos (m{\theta }_{d/2})\end{array}\right]$$

1.3 RoPE的分析

关键是这个角度 ${\theta }_k$ 的设计非常精妙，它使得 RoPE 能够在不同维度上编码不同尺度的位置信息：

• 高频信息（感知近距离）: 当维度索引 $k$ 很小（即在向量的靠前部分）时， ${\theta }_k$ 的值很大。这意味着旋转的频率很高。位置 $m$ 的微小变化都会导致旋转角度 $m{\theta }_k$ 的显著改变。这部分维度负责编码精细的、近距离的相对位置信息。模型可以利用这些高频信号来区分紧邻的词。

• 低频信息（感知远距离）: 当维度索引 $k$ 很大（即在向量的靠后部分）时，${\theta }_k$ 的值很小。这意味着旋转的频率很低。即使位置 $m$ 变化很大，旋转角度 $m{\theta }_k$ 的变化也相对平缓。这部分维度负责编码粗粒度的、远距离的相对位置信息。模型可以利用这些低频信号来判断两个词相距很远。

对于$sin(x\theta )$，它的周期（或者叫波长）是$\frac{2\pi }{\theta }$，一个周期就会回到原始角度。

举例：

• 例如中间向量维度是1024，上下文的长度为32k=32768，批量输入transformer的形状就是[batch_size, 32k, 1024]
• 讲旋转矩阵应用于向量上时，对于这个公式是：${\theta }_k=b^{-2k/d}$
• 这个$d$ 是模型的维度，即向量的维度就是1024，$k$ 是维度的索引，因为每两个值共享一个角度，所以从0开始，到511结束。
• 考虑向量的前两个维度，计算出的角度分别是：${\theta }_0=b^{-2(0)/1024}=10000^0=1$，它的周期是$\frac{2\pi }{1}=2\pi$，也就是相对位置的距离超过6个位置，就会回到原始角度。
• 考虑向量的后两个维度，计算出的角度分别是：${\theta }_{511}=b^{-2(511)/1024}=10000^{-0.9980}\approx 10^{-4}$，它的周期是$\frac{2\pi }{10^{-4}}=2\pi \times 10^4$，也就是相对位置的距离超过62.8k个位置，就会回到原始角度。（即使相对距离达到预定义的上下文长度32k，也不会回到原始角度，因为周期太长）

过于高频的信号会导致模型对远距离位置的感知能力下降，周期太短可能过几个位置又回到了和原始角度相近的角度，所以需要非常低频信号来维护距离信息，周期特别长，根本无法走完整个周期，角度不会重复，从而维护了远位置的相对距离信息

通过在不同维度位置上组合使用不同频率的旋转，RoPE 就像是为模型提供了一把“多尺度”的尺子。它既能用高频信息精确测量邻近词语的距离，又能用低频信息感知到远方词语的大致顺序，从而对整个序列的位置关系有一个完整而丰富的表征。这正是 RoPE 优雅而强大的地方。

这样学习的一个效果就是导致$q,k$ 向量在不同维度上的旋转角度不同，从而能够注入不同的位置信息。向量维度较低的位置负责编码精细的、近距离的相对位置信息，而维度较高的位置负责编码粗粒度的、远距离的相对位置信息。

Part 2: 揭秘 YaRN 的扩展魔法

2.1 基本思想

其实扩展上下文窗口最直观的解决方法是：“把长的文本假装成短的”。

这就是位置内插（Position Interpolation, PI的思想。

假设我们要把窗口从 $L_{train}=4096$ 扩展到 $L_{test}=32768$，扩展倍数 $s=\frac{32768}{4096}=8$。
我们可以把新的位置索引 $m$ 除以 $s$，变成 $m^{\prime }=m/s$。

这样，原本 $0\sim 32768$ 的位置范围，就被压缩回了 $0\sim 4096$。

这样做的优点是：所有的位置都在模型“见过”的范围内。

但是缺点是：分辨率下降，特别是在高频部分，因为高频信号的周期很短，强行拉长会导致模型对位置的感知能力下降。

回到 RoPE 的视角，把位置 $m$ 变成 $m/s$，等价于把所有的频率 ${\theta }_k$ 都缩小了 $s$ 倍（波长拉长了 $s$ 倍）。
即：
${\theta }_k^{\prime }=\frac{{\theta }_k}{s}=\frac{b^{-2k/d}}{s}$

还记得 RoPE 的高频部分负责近距离识别吗？

如果我们把高频信号也强行拉长，相邻 Token 之间的旋转角度变化就变得非常小，模型可能分不清“谁在前，谁在后”了。这严重损害了模型的微观感知能力。

2.2 NTK 感知内插 (NTK-aware Interpolation)

这时候，NTK (Neural Tangent Kernel) 理论给出了一个更聪明的方向：
不要对所有维度一视同仁！

• 高频维度（低维索引）：包含丰富的信息，对分辨率要求极高。它们旋转得很快，即便外推到更远的距离，随机性也比较强，模型反而对外推比较鲁棒。不要去压缩它们！
• 低频维度（高维索引）：波长很长，模型没见过完整的周期。这些才是真正需要“内插/压缩”的部分。

NTK-aware 的做法不是直接除以 $s$，而是修改 RoPE 的基数 $b$（原为 10000）。通过增大 $b$，可以非线性地改变不同维度的频率：让高频部分保持原样，只把低频部分拉长。

计算公式为：
$b^{\prime }=b\cdot s^{\frac{d}{d-2}}$

这样，新的频率 ${\theta }_k^{\prime }=(b^{\prime }{)}^{-2k/d}$ 在高频部分变化很小，在低频部分变化较大，完美符合我们的需求。

按照之前的想法，所有的角度缩放为原来的 $1/s=\frac{1}{8}=0.125$，所有的角度值都乘以 $0.125$ 从而得到新的角度值。但是现在这个相当于这个比例值是按照位置进行指数下降的，高频位置缩放比例没有那么大，而低频位置缩放比例则是接近 $0.125$。

2.3 YaRN：集大成者 (Yet another RoPE extensioN)

YaRN 在 NTK-aware 的基础上，做得更加极致和精细。它的核心思想可以概括为：分而治之。

为了精确地实现这个思想，YaRN 引入了两个关键的阈值参数 $\alpha$ 和 $\beta$，用来界定什么是“高频”，什么是“低频”。

我们首先定义每个维度的波长 $\lambda$。回顾一下 RoPE，第 $k$ 个维度的旋转频率是 ${\theta }_k$，那么它的波长就是完成一次 $2\pi$ 旋转所需的长度（即转一圈需要的 Token 数量）：
$${\lambda }_k=\frac{2\pi }{{\theta }_k}$$

我们将波长 ${\lambda }_k$ 与模型原本的训练长度 $L_{train}$ 进行比较，看看这个维度在训练时“转了几圈”，以此来判断它属于哪个区域：

1. 高频区（High Frequency）：如果不插值，模型也能懂

   • 判断标准：波长很短，小于一定比例的训练长度。具体来说，$\frac{{\lambda }_k}{L_{train}}<\alpha$。
   • 含义：在训练长度 $L_{train}$ 内，这个维度的向量已经旋转了很多圈。模型已经充分见过各种相对位置关系，对它的外推能力很强。
   • 策略：不进行插值 ($\gamma =0$)。保持原汁原味，保留局部的高分辨率。

2. 低频区（Low Frequency）：必须插值，否则模型没见过

   • 判断标准：波长很长，大于一定比例的训练长度。具体来说，$\frac{{\lambda }_k}{L_{train}}>\beta$。
   • 含义：这个维度的旋转极慢，在整个训练长度 $L_{train}$ 内，可能只转了一个很小的角度。如果我们把长度扩展到 $L_{test}$，旋转角度会超出模型见过的范围（Out-of-Distribution）。
   • 策略：完全线性插值 ($\gamma =1$)。把它“压缩”回模型熟悉的范围内。

3. 过渡区（Transition）：平滑过渡

   • 判断标准：介于两者之间。$\alpha \le \frac{{\lambda }_k}{L_{train}}\le \beta$。
   • 策略：混合处理。$\gamma$ 从 0 平滑过渡到 1。

公式化描述：

我们定义一个斜坡函数 ${\gamma }_k$ 来表示插值的强度：

$${\gamma }_k=\{\begin{array}{ll}0& \text{if }\frac{{\lambda }_k}{L_{train}}<\alpha (\text{高频})\\ 1& \text{if }\frac{{\lambda }_k}{L_{train}}>\beta (\text{低频})\\ \frac{\frac{{\lambda }_k}{L_{train}}-\alpha }{\beta -\alpha }& \text{otherwise}(\text{过渡})\end{array}$$

最终的频率 ${\theta }_k^{YaRN}$ 是原始频率 ${\theta }_k$ 和插值频率 ${\theta }_k/s$ 的加权混合：

$${\theta }_k^{YaRN}=(1-{\gamma }_k){\theta }_k+{\gamma }_k\frac{{\theta }_k}{s}$$

为什么需要过渡区？
如果直接一刀切（只有0和1），在临界点的频率会发生突变，导致模型的性能不稳定。过渡区确保了频率变化是平滑连续的。

通常的参数设置：
在 Llama 2 等模型中，通常设置 $\alpha =1$, $\beta =32$。这意味着：

• 波长小于训练长度的，完全不插值。
• 波长大于 32 倍训练长度的，完全插值。
• 中间的进行平滑过渡。

2.4 温度缩放

除了上述的频率分离，YaRN 还引入了一个关键细节：温度缩放（Temperature Scaling）。

当上下文变长时，Attention 机制中的 Softmax 分母（通常是 $\sqrt{d}$）会让注意力分布变得“过于平滑”或“过于尖锐”（熵变了）。
YaRN 建议修改 Softmax 的温度参数：
$$\text{softmax}(\frac{q\cdot k^T}{t\sqrt{d}})$$
其中 $t$ 是一个根据扩展倍数计算的系数（通常 $\sqrt{t}\approx 0.1\ln (s)+1$）。
这一步简单但至关重要，它修正了长文本下的注意力分布，使得模型不需要微调就能保持原来的“注意力习惯”。