大模型上下文扩展之YaRN解析：从RoPE、到ALiBi、位置插值、到YaRN

前言

下半年以来，我全力推动我司大模型项目团队的组建，我虽兼管整个项目团队，但为了并行多个项目，最终分成了三个项目组，每个项目都有一个项目负责人，分别为霍哥、阿荀、朝阳

在今年Q4，我司第一项目组的第一个项目「AIGC模特生成平台」得到CSDN蒋总的大力支持，并亮相于CSDN举办的1024程序员节，一上来就吸引了很多市里领导、媒体、观众的关注，如今该平台的入口链接已在七月官网右上角
而第二项目组的论文审稿GPT，我和阿荀则一直全程推动整个流程的开发(第一版详见此文的第三部分、第二版详见：七月论文审稿GPT第2版：从Meta Nougat、GPT4审稿到Mistral、LongLora Llama)
到12月中旬，进入了模型训练阶段，选型的时候最开始关注的两个模型，一个Mistral 7B，一个Llama-LongLora，但考虑到前者的上下文长度是8K，面对一些论文时可能长度还是不够，于是我们便考虑让Mistral结合下YaRN
所以本文重点介绍下YaRN，顺带把位置编码外推ALiBi、线性插值等相关的方法一并总结下
至于第三项目组的知识库问答项目则也一直在并行推进，核心还是一系列各种细节问题的优化，而这个优化过程还是比较费时的

YaRN本质上是一种新的RoPE扩展方法(至于RoPR详见此文：)，可以比较高效的扩展大模型的上下文窗口，本文的写就基于YaRN论文：YaRN: Efficient Context Window Extension of Large Language Models，且为方便大家更好的理解本文，特地再列下本文重要的几个参考文献(当下文出现带中括号的[6]、[7]、[9]时，便特指的以下相关文献)

[6] bloc97. NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation., 2023.
URL https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_ scaled_rope_allows_llama_models_to_have/
[7] bloc97. Add NTK-Aware interpolation "by parts" correction, 2023. URL https://github.

com/jquesnelle/scaled-rope/pull/1 .
[9] S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation, 2023. arXiv: 2306.15595.

有何问题欢迎随时留言评论，thanks

第一部分上下文窗口扩展史：从RoPE到ALiBi、位置插值、NTK-aware Scaled RoPE

基于transformer的大型语言模型(llm)已经成为许多自然语言处理(NLP)任务的选择，其中远程能力(如上下文学习(ICL))至关重要。在执行NLP任务时，其上下文窗口的最大长度一直是预训练LLM的主要限制之一。能够通过少量的微调(或不进行微调)来动态扩展上下文窗口已经变得越来越受关注。为此，transformer的位置编码是核心焦点

最初的Transformer架构使用了绝对正弦位置编码，后来被改进为可学习的绝对位置编码[Convolutional sequence to sequence learning]。此后，相对位置编码方案[Self-attention with relative position representations]进一步提升了transformer的性能。目前，最流行的相对位置编码是T5 relative Bias[Exploring the limits of transfer learning with a unified text-to-text transformer]、RoPE[34]、XPos[35]和ALiBi[Attention with linear biases enables input length extrapolation]
位置编码的一个反复出现的限制是无法对「训练期间看到的上下文窗口之外的情况」进行泛化
One reoccurring limitation with positional encodings is the inability to generalize past the context window seen during training

虽然ALiBi等一些方法能够进行有限的泛化，但没有一种方法能够泛化到明显长于预训练长度的序列
好在已经有一些工作正在尝试克服这种限制。比如位置插值(Position Interpolation, PI)[Extending context window of large language models
via positional interpolation]通过对RoPE进行轻微修改，并对少量数据进行微调，从而扩展上下文长度
作为一种替代方案，Reddit一网友bloc97通过该帖子，提出了“NTK-aware”插值方法[NTK-Aware Scaled RoPE allows LLaMA models to have extended(8k+) context size without any fine-tuning and minimal perplexity degradation]，该方法考虑到高频信号的损失

此后，对“NTK感知”插值提出了两项改进

无需微调的预训练模型的“动态NTK”插值方法[14]
在对少量较长的上下文数据进行微调时表现最佳的“NTK-by-parts”插值方法[7]

“NTK感知”插值和“动态NTK”插值已经在开源模型中出现，如Code Llama[31](使用“NTK感知”插值)和Qwen 7B[2](使用“动态NTK”)。

在本文中，除了对先前未发表的关于“NTK感知”、“动态NTK”和“NTK-by-parts”插值的技术进行完整的描述外，我还将介绍YaRN(另一种RoPE扩展方法)，这是一种改进的方法，可以有效地扩展使用旋转位置嵌入(RoPE)训练的模型的上下文窗口，包括LLaMA[38]、GPT-NeoX[5]和PaLM[10]家族的模型

在对不到0.1%的原始预训练数据进行微调后，YaRN在上下文窗口扩展中达到了最先进的性能。同时，通过与称为动态缩放的推理时间技术相结合，Dynamic- yarn允许超过2倍的上下文窗口扩展，而无需任何微调

1.1 旋转位置嵌入

YaRN的基础是[RoFormer: Enhanced transformer with rotary position embedding]中介绍的旋转位置嵌入(RoPE)

考虑到本博客内已有另一篇文章详细阐述了位置编码与RoPE，所以如果你对本节有任何疑问，可进一步参考此文：一文通透位置编码：从标准位置编码、欧拉公式到旋转位置编码RoPE

所以下面只参照YaRN论文做个最简单的回顾，首先，我们在一个隐藏层上工作，隐藏神经元的集合用 $d$ 表示。给定向量序列 $\mathbf{x}_{1}, \cdots, \mathbf{x}_{L} \in \mathbb{R}^{|D|}$ ，遵循RoPE的表示法，注意力层首先将向量转换为查询向量和关键向量：

$\mathbf{q}_{m}=f_{q}\left(\mathbf{x}_{m}, m\right) \in \mathbb{R}^{|D|}, \mathbf{k}_{n}=f_{k}\left(\mathbf{x}_{n}, n\right) \in \mathbb{R}^{|D|}$

接下来，注意力权重被计算为

$\operatorname{softmax}\left(\frac{\mathbf{q}_{m}^{T} \mathbf{k}_{n}}{\sqrt{|D|}}\right)$

其中 $q_m$ 、 $k_n$ 被认为是列向量，因此 $\mathbf{q}_{m}^{T} \mathbf{k}_{n}$ 就是简单的欧氏内积。在RoPE中，我们首先假设 $|D|$ 是偶数，并将嵌入空间和隐藏状态识别为complex vector spaces

$\mathbb{R}^{|D|} \cong \mathbb{C}^{|D| / 2}$

其中内积 $\mathbf{q}^{T} \mathbf{k}$ 转化为 $\operatorname{Re}\left(\mathbf{q}^{*} \mathbf{k}\right)$ 的实部「where the inner product q T k becomes the real part of the standard Hermitian inner product Re(q ∗k)，如对该点有疑问的，请参见此文的3.2.1节」，更具体地说，同构将实数部分和复数部分交织在一起(the isomorphisms interleave the real part and the complex part)

$\begin{aligned} \left(\left(\mathbf{x}_{m}\right)_{1}, \cdots,\left(\mathbf{x}_{m}\right)_{|D|}\right) & \mapsto\left(\left(\mathbf{x}_{m}\right)_{1}+i\left(\mathbf{x}_{m}\right)_{2}, \cdots,\left(\left(\mathbf{x}_{m}\right)_{|D|-1}+i\left(\mathbf{x}_{m}\right)_{|D|}\right)\right) \\ \left(\left(\mathbf{q}_{m}\right)_{1}, \cdots,\left(\mathbf{q}_{m}\right)_{|D|}\right) & \mapsto\left(\left(\mathbf{q}_{m}\right)_{1}+i\left(\mathbf{q}_{m}\right)_{2}, \cdots,\left(\left(\mathbf{q}_{m}\right)_{|D|-1}+i\left(\mathbf{q}_{m}\right)_{|D|}\right)\right) \end{aligned}$

为了将嵌入 $x_m$ 、 $x_n$ 转换为查询向量和键向量，我们首先给出了R-linear算子

$\mathbf{W}_{q}, \mathbf{W}_{k}: \mathbb{R}^{|D|} \rightarrow \mathbb{R}^{|D|}$

在复坐标中，函数 $f_q$ ， $f_k$ 由

$f_{q}\left(\mathbf{x}_{m}, m\right)=e^{i m \theta} \mathbf{W}_{q} \mathbf{x}_{m}, f_{k}\left(\mathbf{x}_{n}, n\right)=e^{i n \theta} \mathbf{W}_{k} \mathbf{x}_{n}$

这样做的好处是，查询向量和关键向量之间的点积只取决于如下所示的相对距离 $m-n$

$\begin{aligned} & \left\langle f_{q}\left(\mathbf{x}_{m}, m\right), f_{k}\left(\mathbf{x}_{n}, n\right)\right\rangle_{\mathbb{R}} \\ = & \operatorname{Re}\left(\left\langle f_{q}\left(\mathbf{x}_{m}, m\right), f_{k}\left(\mathbf{x}_{n}, n\right)\right\rangle_{\mathbb{C}}\right) \\ = & \operatorname{Re}\left(\mathbf{x}_{m}^{*} \mathbf{W}_{q}^{*} \mathbf{W}_{k} \mathbf{x}_{n} e^{i \theta(m-n)}\right) \\ = & g\left(\mathbf{x}_{m}, \mathbf{x}_{n}, m-n\right) . \end{aligned}$

在实坐标中，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}$

如此，便有

$f_{q}=f_{\mathbf{W}_{q}}, f_{k}=f_{\mathbf{W}_{k}}$

1.2 位置插值

由于语言模型通常是用固定的上下文长度进行预训练的，自然会问如何通过在相对较少的数据量上进行微调来扩展上下文长度

对于使用RoPE作为位置嵌入的语言模型，Chen等人[9]和同时kaiokendev[21]提出了位置插值(position Interpolation, PI)，将上下文长度扩展到预训练极限之外
虽然直接外推在 $L$ 大于预训练极限的序列 $w_{1}, \cdots, w_{L}$ 上的表现并不好，但他们发现，在预训练极限内插入位置指标，并进行少量微调可以显著提高效果
While a direct extrapolation does not perform well on sequences w1, · · · , wL with L larger than the pre-trained limit, they discovered that interpolating the position indicies with in the pre-trained limit works well with the help of a small amount of fine-tuning

具体来说，给定一个带有RoPE的预训练语言模型，他们通过
其中 $L ' > L$ 是超出预训练限制的新上下文窗口。通过原始的预训练模型加上修改的RoPE公式，他们在几个数量级更少的token上进一步微调了语言模型，并成功实现了上下文窗口扩展
考虑到扩展的上下文长度与原始上下文长度之间的比例一直特别重要，我们以此定义 $s=\frac{L^{\prime}}{L}$
有了 $s$ 这个定义，我们便可以将公式 $f_{\mathbf{W}}^{\prime}\left(\mathbf{x}_{m}, m, \theta_{d}\right)=f_{\mathbf{W}}\left(\mathbf{x}_{m}, \frac{m L}{L^{\prime}}, \theta_{d}\right)$ 重写并简化为以下一般形式：
$f_{\mathbf{W}}^{\prime}\left(\mathbf{x}_{m}, m, \theta_{d}\right)=f_{\mathbf{W}}\left(\mathbf{x}_{m}, g(m), h\left(\theta_{d}\right)\right)$
其中 $g(m)=m / s$ ， $h\left(\theta_{d}\right)=\theta_{d}$

1.3 从“NTK-aware”插值到“NTK-by-parts”插值、

1.3.1 “NTK-aware”插值

为了解决RoPE嵌入插值时丢失高频信息的问题，Reddit一网友通过[NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation]开发了“NTK-aware”插值。我们不是将RoPE的每个维度平均缩放一个因子s，而是通过减少高频的缩放和增加低频的缩放来将插值压力分散到多个维度。人们可以通过许多方法获得这样的变换，但最简单的方法是对θ的值进行基础更改。

“NTK-aware”的内插方案如下

“NTK-aware”插值是对RoPE的修改，使用 $f_{\mathbf{W}}^{\prime}\left(\mathbf{x}_{m}, m, \theta_{d}\right)=f_{\mathbf{W}}\left(\mathbf{x}_{m}, g(m), h\left(\theta_{d}\right)\right)$ 和以下函数

$g(m)=m$

$h\left(\theta_{d}\right)=b^{\prime-2 d /|D|}$

且

$b^{\prime}=b \cdot s^{\frac{|D|}{|D|-2}}$

与PI[Extending context window of large language models via positional interpolation]相比，该方法在扩展非微调模型的上下文大小方面表现得更好。然而，这种方法的一个主要缺点是，由于它不仅仅是一种插值方案，一些维度被轻微外推到“超出边界”的值，因此使用“NTK-aware”插值[6]进行微调的结果不如PI[9]。此外，由于存在“越界”值，理论尺度因子 $s$ 并不能准确描述真实的上下文扩展尺度。在实践中，对于给定的上下文长度扩展，尺度值s必须设置得高于预期尺度

我们注意到，在本文发布前不久，Code Llama[31]发布了，并通过手动将基数b扩展到1M来使用“NTK-aware”扩展。

1.3.2 相对局部距离的损失-“NTK-by-parts”插值

在像PI和“ntk感知”插值这样的盲插值方法中，我们平等地对待所有RoPE隐藏维度(因为它们对网络有相同的影响)。然而，有强有力的线索指向我们需要有针对性的插值方法。

在本节中，我们将着重考虑RoPE公式中Eq. 13中定义的波长λd。为简单起见，我们省略了λd中的下标d，并鼓励读者将λ视为任意周期函数的波长。

我们定义λd为RoPE嵌入在第d隐维处的波长：
$\lambda_{d}=\frac{2 \pi}{\theta_{d}}=2 \pi b^{\frac{2 d}{|D|}}$
波长描述了在维d上嵌入的RoPE执行完整旋转(2π)所需的标记长度
假设一些插值方法(例如。(PI)不关心波长的维数，我们将这些方法称为“盲”插值方法，而其他方法(如YaRN)，我们将其归类为“有针对性的”插值方法。