LLM上下文长度扩展方案：YaRN

题目： YaRN: Efficient Context Window Extension of Large Language Models
论文地址： YaRN: Efficient Context Window Extension of Large Language Models

I. 前言

在之前的两篇文章中分别介绍了上下文长度扩展的两种方案：

1. LLM上下文长度扩展方案：Position Interpolation
2. LLM上下文长度扩展方案：NTK-aware interpolation

其中PI线性缩放所有位置索引，并平等地对待每个维度；NTK-aware则实现了高频外推和低频内插，缓解了高频也就是低维的压力，使模型能够更好地区分不同位置间的精细区别。

为了引出上述两种方案存在的问题，这里先引入一个新的定义：给定RoPE中一个具体的维度$i$，则这个维度的波长被定义为：
$${\lambda }_i=\frac{2\pi }{{\theta }_i}=\frac{2\pi }{b^{-2(i-1)/d}}=2\pi b^{2(i-1)/d}$$基于波长公式，我们有如下结论：

1. 波长描述了为了在维度$i$处嵌入的旋转位置执行全旋转($2\pi$)所需的token的长度。
2. 维度越高波长越长。

从上图我们知道第$i$个维度（两两一组）处二维复向量的旋转角度为$m{\theta }_i$，为了能够在该维度处执行全旋转，令$m{\theta }_i=2\pi$，则${\lambda }_i=m$，所以波长描述了为了在维度$i$处执行全旋转($2\pi$)所需的token的长度。

像PI这种类型的插值方案不关心波长的维数，我们将这些方法称为“盲”插值方法(blind interpolation)，比如像PI和“NTK-aware”插值这样的blind interpolation方法中，我们面对所有RoPE隐藏维度没有做任何针对性的处理。而其他方法如这篇文章提出的YaRN，我们将其归类为“有针对性的”插值方法，即有对RoPE的不同维度做出不同处理。

II. NTK-by-parts

为了针对不同维度做出不同的处理，YaRN中首先定义了NTK-by-parts插值方法。具体来说，关于RoPE中不同维度的波长，有如下结论：

存在某些维度$i$，其波长${\lambda }_i$大于在预训练期间看到的最大上下文长度$L$

从理论上来讲，RoPE是一种编码绝对位置的方法。然而，我们前面说过，波长描述了为了在维度$i$处嵌入的旋转位置执行全旋转($2\pi$)所需的token的长度，如果某些维度的波长大于上下文长度$L$，这说明该维度无法执行全旋转。在这种情况下，由于维度在预训练期间至少不会完全旋转一次，如果我们选择第一个令牌作为基准，那么在预训练期间每隔一个令牌到它的距离是唯一的，神经网络可以用它来确定它的绝对位置信息。相反，如果某个维度波长小于$L$，那该维度就执行了至少一次全旋转，我们就无法在这个维度描述绝对距离，只能描述相对位置信息。

考虑到上述观察，作者认为，不要对只编码相对位置信息的维度（${\lambda }_i<L$）进行内插破坏，因为它们对于模型区分附近令牌的相对顺序至关重要。同时，应该始终对仅编码绝对位置信息的维度（${\lambda }_i>L$）进行内插，因为较大的距离将超出之前模型能够编码的距离。因此，可以制定一种考虑以上所有因素的显式且有针对性的插补方法，即：

1. 如果波长远小于$L$，即编码相对位置的维度，不进行内插
2. 如果波长等于或大于$L$，即编码绝对位置的维度，应该进行内插以防止超出绝对位置的最大可编码范围
3. 波长介于上述之间的维度，采用NTK-aware方法

为了定义上述不同内插策略的边界，引入了两个额外的参数$\alpha$和$\beta$，并将维度$i$处的波长定义为：
$$r(i)=\frac{2\pi }{{\theta }_i}=\frac{2\pi }{b^{-2i/d}}=2\pi b^{2i/d}$$上述维度的取值从0开始。接着，定义一个分段函数如下：
$$\gamma (r)=\{\begin{array}{ll}0& \text{if }r<\alpha ,\\ 1& \text{if }r>\beta ,\\ \frac{r-\alpha }{\beta -\alpha }& \text{otherwise}.\end{array}$$基于上述分段函数，NTK-by-parts被定义为：
$$\begin{array}{rl}{f}_{parts}(x_m,m,\theta )& =f(x_m,g(m),h({\theta }_i))\\ g(m)& =m\\ h({\theta }_i)=(1-& \gamma (r(i)))\frac{{\theta }_i}{S}+\gamma (r(i)){\theta }_i\end{array}$$在原文中，作者针对LLaMA系列模型给出的建议参数值为$\alpha =1$和$\beta =32$。简单来说，如果当前维度的波长远小于上下文长度$L$(即远小于L/32)时不内插，即$h({\theta }_i)={\theta }_i$；当波长大于等于上下文长度L(即大于$\alpha L$)时执行内插，此时有$h({\theta }_i)=\frac{{\theta }_i}{S}$。

笔者曾在Qwen2系列模型上尝试多个组合值，发现效果始终不如$\alpha =1$和$\beta =32$，这说明两个参数具有较好的代表性。

III. YaRN

YaRN的作者注意到，无论数据样本或者扩展上下文窗口上的令牌位置怎么样，在计算自注意力公式中引入温度系数$t$对困惑度ppl有着一致的影响。为此，将该公式改为：
$$\text{softmax}(\frac{q_m^{\top }{k}_n}{t\sqrt{d}})$$为了实现上述公式，只需要将两个位置的旋转位置嵌入各自缩放为原来的$1/\sqrt{t}$即可，而无需做其他任何操作。

最后，将NTK-by-parts技术和上面自注意力计算公式结合，便得到了YaRN。

作者在原文中基于实验给出了$1/\sqrt{t}$的表达式：
$$\sqrt{\frac{1}{t}}=0.1\ln S+1$$笔者曾根据作者的方法对Qwen2系列模型的温度系数参数进行了拟合，发现效果远差于上式。因此建议使用YaRN扩展除LLaMA以外的模型时，也可以直接采用上式。

IV. Dynamic NTK

在真正使用模型时，我们无法保证每次输入的长度都是固定的。在这种情况下，有以下两种方法得到扩展因子$S$：

1. 在整个推理周期中，$S=\frac{L^{{}^{\prime }}}{L}$是固定的，保证模型输入的最大长度不超过扩展后的长度$L^{{}^{\prime }}$即可
2. 根据每次输入的长度来动态调整$S$，进而动态内插。扩展比被定义为$S=\max (1,l^{{}^{\prime }}/L)$，其中$l^{{}^{\prime }}$为当前输入的长度。

第一种方式的问题是，当序列长度小于L时，模型的性能可能会出现一些折扣，而当序列长度大于L时，模型的性能会退化。但通过像方式2那样进行动态缩放，它允许模型在达到训练的上下文限制时较为缓慢地退化性能，而不是骤降。

将上述动态缩放和NTK-aware插值相结合时，我们称之为动态NTK插值。