大模型上下文扩展之YaRN解析：从直接外推ALiBi、位置插值PI、NTK-aware插值、YaRN

v_JULY_v

已于 2024-03-09 15:30:58 修改

阅读量4.1k

点赞数 23

分类专栏：七月论文大模型：含论文的审稿、阅读、写作、修订文章标签：大模型上下文扩展 RoPE ALiBi 位置插值 YaRN

于 2023-12-19 12:26:03 首次发布

本文链接：https://blog.csdn.net/v_JULY_v/article/details/135072211

版权

七月论文大模型：含论文的审稿、阅读、写作、修订专栏收录该内容

11 篇文章 11 订阅

订阅专栏

前言

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

在今年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扩展方法(至于RoPE详见此文)，可以比较高效的扩展大模型的上下文窗口，本文的写就基于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/
提出了NTK-Aware插值
[7] bloc97. Add NTK-Aware interpolation "by parts" correction, 2023. URL https://github.

com/jquesnelle/scaled-rope/pull/1.
提出了 “NTK-by-parts”插值
[9] S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation, 2023. arXiv: 2306.15595.
该研究团队来自Meta，该篇论文提出了 位置内插PI

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

第一部分背景知识：从进制表示谈到直接外推、线性内插、进制转换

1.1 从进制表示到直接外推

注，本部分内容援引自苏剑林博客中的部分内容，为更易懂，我在其基础上做了一定的修改、解读

1.1.1 进制表示

假设我们有一个1000以内（不包含1000）的整数 $n$ 要作为条件输入到模型中，那么要以哪种方式比较好呢？

最朴素的想法是直接作为一维浮点向量输入，然而0～999这涉及到近千的跨度，对基于梯度的优化器来说并不容易优化得动。那缩放到0～1之间呢？也不大好，因为此时相邻的差距从1变成了0.001，模型和优化器都不容易分辨相邻的数字
进一步，对于一个整数，比如759，这是一个10进制的三位数，每位数字是0～9。既然我们自己都是用10进制来表示数字的，为什么不直接将10进制表示直接输入模型呢？也就是说，我们将整数n以一个三维向量[a,b,c]来输入，a,b,c分别是n的百位、十位、个位
至于如果想要进一步缩小数字的跨度，我们还可以进一步缩小进制的基数，如使用8进制、6进制甚至2进制，代价是进一步增加输入的维度

1.1.2 直接外推

苏剑林说，假设我们还是用三维10进制表示训练了模型，模型效果还不错。然后突然来了个新需求，将n上限增加到2000以内，那么该如何处理呢？

如果还是用10进制表示的向量输入到模型，那么此时的输入就是一个四维向量了。然而，原本的模型是针对三维向量设计和训练的，所以新增一个维度后，模型就无法处理了。可能有读者想说，为什么不能提前预留好足够多的维度呢？

没错，是可以提前预留多几维，训练阶段设为0，推理阶段直接改为其他数字，这就是外推（Extrapolation）

然而，训练阶段预留的维度一直是0，如果推理阶段改为其他数字，效果不见得会好，因为模型对没被训练过的情况不一定具有适应能力。也就是说，由于某些维度的训练数据不充分，所以直接进行外推通常会导致模型的性能严重下降。

1.2 从线性内插到进制转换

1.2.1 线性内插

于是，有人想到了将外推改为内插(Interpolation)，简单来说就是将2000以内压缩到1000以内

比如通过除以2，1749就变成了874.5，然后转为三维向量[8,7,4.5]输入到原来的模型中

从绝对数值来看，新的[7,4,9]实际上对应的是1498，是原本对应的2倍，映射方式不一致；
从相对数值来看，原本相邻数字的差距为1，现在是0.5，最后一个维度更加“拥挤”
所以，做了内插修改后，通常都需要微调训练，以便模型重新适应拥挤的映射关系

当然，有读者会说外推方案也可以微调。是的，但内插方案微调所需要的步数要少得多

因为很多场景(比如位置编码)下，相对大小(或许说序信息)更加重要，换句话说模型只需要知道874.5比874大就行了，不需要知道它实际代表什么多大的数字。而原本模型已经学会了875比874大，加之模型本身有一定的泛化能力，所以再多学一个874.5比874大不会太难
不过，内插方案也不尽完美，当处理范围进一步增大时，相邻差异则更小，并且这个相邻差异变小集中在个位数，剩下的百位、十位，还是保留了相邻差异为1
换句话说，内插方法使得不同维度的分布情况不一样，每个维度变得不对等起来，模型进一步学习难度也更大

1.2.2 进制转换

有没有不用新增维度，又能保持相邻差距的方案呢？有，那就是进制转换

三个数字的10进制编码可以表示0～999
如果是16进制呢？它最大可以表示 $16^{3}-1=4095>1999$
所以，只需要转到16进制，如1749变为 $[6,13,5]$ (咋计算得来的？很简单，1749不断除以16，并记录下每次的余数， $6 \times 16^{2}+13 \times 16^{1}+5 \times 16^{0}$ ，这些余数从最后一个到第一个组成了1749的十六进制表示，即“6 13 5”)，那么三维向量就可以覆盖目标范围，代价是每个维度的数字从0～9变为0～15

刚才说到，我们关心的场景主要利用序信息

原来训练好的模型已经学会了 $875>874$ ，而在16进制下同样有 $875>874$ ，比较规则是一模一样的
唯一担心的是每个维度超过9之后(10～15)模型还能不能正常比较，但事实上一般模型也有一定的泛化能力，所以每个维度稍微往外推一些是没问题的。所以，这个转换进制的思路，甚至可能不微调原来模型也有效

另外，为了进一步缩窄外推范围，我们还可以换用更小的 $\lceil\sqrt[3]{2000} \mid=13$ 即13进制而不是16进制

第二部分从RoPE、直接外推、位置内插到NTK-aware/NTK-by-parts/Dynamic NTK插值

基于transformer的大型语言模型已经成为许多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感知”插值和“Dynamic NTK”插值已经在开源模型中出现，如Code Llama[31](使用“NTK感知”插值)和Qwen 7B[2](使用“动态NTK”)

2.1 旋转位置嵌入

2.1.1 RoPE的快速回顾

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

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

所以下面只参照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}}$

2.1.2 位置 $n$ 的旋转位置编码(RoPE)，本质上就是数字 $n$ 的 $\beta$ 进制编码

首先，如苏剑林所说，位置 $n$ 的旋转位置编码(RoPE)，本质上就是数字 $n$ 的 $\beta$ 进制编码

为了理解这一点，我们首先回忆一个10进制的数字 $n$ ，我们想要求它的 $\beta$ 进制表示的(从右往左数)第 $m$ 位数字，方法是根据下面的公式计算得到(记为公式1)

$\left\lfloor\frac{n}{\beta^{m-1}}\right\rfloor \bmod \beta$

也就是先除以 $\beta^{k-1}$ 次方，然后求模(余数)

以上咋推导得来的呢？为方便大家一目了然，我再解释下

例如，让我们找到十进制数12345中从右边数的第三位的数字，相当于 $n=12345$ ， $\beta =10$ (因为是十进制)， $m=3$ (因为要找的是第三位)

按照公式，我们首先计算 $\beta^{m-1}=10^{3-1}=10^{2}=100$
然后，我们计算 $n$ 除以 $\beta^{m-1}$ ，即 $12345 \div 100=123.45$
接下来，我们取这个结果的向下取整值，也就是去掉小数部分，得到 $\lfloor 123.45\rfloor=123$
最后，我们对 $\beta$ 取模，得到 $123 \bmod 10=3$

所以，12345的从右边数的第三位数字是3

其次，苏剑林在其博客中再说道

RoPE的构造基础是Sinusoidal位置编码，可以改写为下面的公式(记为公式2)

$\left[\cos \left(\frac{n}{\beta^{0}}\right), \sin \left(\frac{n}{\beta^{0}}\right), \cos \left(\frac{n}{\beta^{1}}\right), \sin \left(\frac{n}{\beta^{1}}\right), \cdots, \cos \left(\frac{n}{\beta^{d/2 -1}}\right), \sin \left(\frac{n}{\beta^{d/2 -1}}\right)\right]$

其中， $\beta=10000^{\frac{2}{d}}$

可能有的读者还是有点问题，可能还是得再解释下

首先，我们通过上文已多次提到的此文《一文通透位置编码：从标准位置编码、欧拉公式到旋转位置编码RoPE、ALiBi》，来回顾下transformer原始论文中的Sinusoidal位置编码

$PE_{(pos,2i+1)} = cos\left ( \frac{pos}{10000^{\frac{2i}{d_{model}}}} \right )$

$PE_{(pos,2i)} = sin\left ( \frac{pos}{10000^{\frac{2i}{d_{model}}}} \right )$

如阿荀所说，可知

$\cos \left(\frac{n}{10000^{2 i / d}}\right)=\cos \left(\frac{n}{10000^{(2 / d) * i}}\right)=\cos \left(\frac{n}{\left(10000^{(2 / d)}\right)^{i}}\right)=\cos \left(\frac{n}{\beta^{i}}\right)$

其中

$\beta=10000^{\frac{2}{d}}$
$i$ 的取值范围为 $[0,...,\frac{d_{model}}{2}]$ ，啥意思？意味着

pos
(0 2 4等偶数维用sin函数计算) $i$ 公式2中 $\beta$ 的指数项
0 $i = 0 // 2 = 0$ 0
1 $i = 1 //2 =0$ 0
2 $i = 2 // 2 = 1$ 1
3 $i = 3 // 2 = 1$ 1
4 $i = 4 // 2 = 2$ 2
5 $i = 5//2 = 2$ 2
6
....
510 $i = 510 // 2 = 255$ $d/2 -1 = (512/2) -1 = 256 - 1 = 255$
511 $i = 511 // 2 = 255$
$d/2 -1 = (512/2) -1 = 256 - 1 = 255$

所以，也就有了上面的公式2
$\left[\cos \left(\frac{n}{\beta^{0}}\right), \sin \left(\frac{n}{\beta^{0}}\right), \cos \left(\frac{n}{\beta^{1}}\right), \sin \left(\frac{n}{\beta^{1}}\right), \cdots, \cos \left(\frac{n}{\beta^{d/2 -1}}\right), \sin \left(\frac{n}{\beta^{d/2 -1}}\right)\right]$

现在，对比公式1、公式2，是不是也有一模一样的 $\frac{n}{\beta^{m-1}}$

至于模运算，它的最重要特性是周期性，而公式2的cos、sin是不是刚好也是周期函数？所以，除掉取整函数这个无关紧要的差异外，RoPE(或者说Sinusoidal位置编码)其实就是数字 $n$ 的 $\beta$ 进制编码

// 待更

2.2 直接外推之ALiBi

此文《一文通透位置编码：从标准位置编码、欧拉公式到旋转位置编码RoPE、ALiBi》的第4部分已经详细介绍过了ALiBi

简言之，ALiBi是对Transformers进行长度外推，即在短上下文窗口上进行训练，并在较长的上下文窗口上进行推理

好处是虽然一开始不用对模型结构做任何更改
但坏处是直接把位置外推到没有见到的地方会导致模型灾难性的崩坏(例如体现在PPL陡增)，为了弥补，需要再做一些微调

// 待更..

2.3 位置内插：基于Positional Interpolation扩大模型的上下文窗口

2.3.1 什么是位置内插Positional Interpolation

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

对于使用RoPE作为位置嵌入的语言模型，Chen等人[9]和kaiokendev[21]同时提出了位置插值(position Interpolation, PI)，将上下文长度扩展到预训练极限之外

对于后者，(Super-HOT kaiokendev(2023))，它在RoPE中插入了位置编码，将上下文窗口从2K扩展到8K

对于前者，按照该篇论文《Extending context window of large language models via positional interpolation》，可知

关键思想是，我们不是进行外推，而是直接将位置索引缩小，使最大位置索引与预训练阶段的先前上下文窗口限制相匹配
we directly down-scale the position indices so that the maximum position index matches the previous context window limit in the pre-training stage

如下图所示，下图左上角为预训练阶段的位置向量范围[0,2048]，右上角为长度外推的部分(2048,4096]
如果直接使用位置(2048,4096]进行推理，那么因为模型没有见过这一部分的位置，效果会出现灾难性的下降。那么，就把[0,4096]这个区间”压缩“到[0,2048]不就可以了嘛
于是，原先的1就变成了0.5，4096就变成了2048，这就是位置内插法，即把没见过的位置映射到见过的位置
相当于对于绝对位置 $m$ ，我们把它”缩放“一下，变成 $\frac{m L^{\prime}}{L}$ 。其中， $L$ 为原先支持的长度(如2048)， $L'$ 为需要扩展的长度(如4096)。这样，在计算query和key的时候，就有
其中 $L ' > L$ 是超出预训练限制的新上下文窗口。通过原始的预训练模型加上修改的RoPE公式，他们在几个数量级更少的token上进一步微调了语言模型，并成功实现了上下文窗口扩展
考虑到扩展的上下文长度与原始上下文长度之间的比例一直特别重要，我们以此定义
$s=\frac{L^{\prime}}{L}$
有了 $s$ 这个定义(这个 $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)$ 重写并简化为以下一般形式(其中 $g(m)=m / s$ ， $h\left(\theta_{d}\right)=\theta_{d}$ )：

$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)$

最终，通过位置插值方法，将预训练的7B、13B、33B和65B LLaMA模型(Touvron等人，2023)扩展到大小为32768的各种上下文窗口，除了重新缩放使用位置插值扩展的模型的位置索引外，没有以任何方式修改LLaMA模型架构

2.3.2 位置内插的问题

话说，位置插值法有什么问题呢？

我们先看下三角函数 $\sin (w x)$ ，它的周期是 $T=2 \pi / \omega$
对应到RoPE里的每个维度 $\sin m \theta_{j}, \cos m \theta_{j}$ ，其中 $\theta_{j}=10000^{-2(j-1) / d}, j \in[1,2, \ldots, d / 2]$ (其中， $m$ 是指位置， $j$ 是指维度）
计算得到周期为： $\frac{2 \pi}{m} b^{\frac{2(j-1)}{d}}$ ，其中，用 $b$ 表示base，即10000
从周期计算的公式我们可以知道，针对不同的维度编码 $j$ ，每个维度对应的三角函数周期是越来越大的(即对应到低频、高频)

如果插值是针对绝对位置 $m$ ，那么对每个维度 $j$ 都同等地生效；但是周期小(高频)维度，插值之后会变得很密集(本来一个周期包含10个值，但是内插之后能包含20个值)，这样高频的维度就变的很拥挤

2.4 从“NTK-aware”插值到“NTK-by-parts”插值、"Dynamic NTK"插值

2.4.1 “NTK-aware”插值：高频外推，低频内插

为了解决RoPE嵌入插值时丢失高频信息(losing high frequency information when interpolating the RoPE embeddings)的问题，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$ ，而是通过减少高频的缩小和增加低频的放大来将插值压力分散到多个维度(Instead of scaling every dimension of RoPE equally by a factor s, we spread out the interpolation pressure across multiple dimensions by scaling high frequencies less and low frequencies more)
虽然人们可以通过许多方法获得这样的变换，但最简单的方法是对θ的值进行基础更改(One can obtain such a transformation in many ways, but the simplest would be to perform a base change on the value of θ)

啥意思呢，其实我们是要把2.1.2节中的公式2

该公式2中的最低频项 $\frac{n}{\beta^{\mathrm{d} / 2-1}}$ ，引入参数 $\lambda$ ，从而变为 $\frac{n}{(\beta \lambda)^{d / 2-1}}$ ，让它跟内插一致(内插就是将 $n$ 换成 $n/k$ ，其中 $k$ 是要扩大的倍数)，即

$\frac{n}{(\beta \lambda)^{d / 2-1}}=\frac{n / k}{\beta^{d / 2-1}}$

从而解得

$\lambda=\mathrm{k}^{2 /(\mathrm{d}-2)}$

而公式2中的最高频是 $\frac{n}{\beta }$ 项，引入 $\lambda$ 后变为 $\frac{n}{\beta \lambda }$ ，由于 $d$ 通常很大， $\lambda$ 很接近1，所以它还是接近于 $\frac{n}{\beta }$ ，即等价于外推

从而，NTK-aware便这样把外推和内插给结合起来了

YaRN论文中对“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}}$

其中 $b$ 相当于上面的 $\beta$ ，且 $s=\frac{L^{\prime}}{L}$ 类似于上面的 $k$ ， $m$ 相当于上面的 $n$

可能有同学有疑问了，咋和上面的表示不一致呢？

原因如阿荀所说

上面(进制角度)的 $d$ 和这里(YaRN论文)的 $|D|$ 才是一回事(容易把上面的 $d$ 和这里的 $d$ 混淆看做同一个)
且上面(进制角度)的计数范围是从0到d/2-1的，而这里(YaRN论文)的计数范围是从1到|D|/2的

与位置插值PI相比，该方法在扩展非微调模型的上下文大小方面表现得更好

然而，这种方法的一个主要缺点是，由于它不仅仅是一种插值方案，一些维度被轻微外推到“超出边界”的值，因此使用“NTK-aware”插值[6]进行微调的结果不如PI[9]
此外，由于存在“越界”值，理论尺度因子 $s$ 并不能准确描述真实的上下文扩展尺度。在实践中，对于给定的上下文长度扩展，尺度值s必须设置得高于预期尺度

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

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

在本节伊始，得先介绍一个概念，即波长，所谓波长，其描述的在维 $d$ 上嵌入的RoPE，执行完整旋转(2π)所需的标记长度

一般而言，把 $\lambda_{d}$ 定义为RoPE嵌入在第 $d$ 隐维处的波长：

$\lambda_{d}=\frac{2 \pi}{\theta_{d}}=2 \pi b^{\frac{2 d}{|D|}}$

有一些插值方法(例如位置插值PI)不关心波长的维数，我们将这些方法称为“盲”插值方法(blind interpolation)，比如像PI和“NTK-aware”插值这样的blind interpolation方法中，我们面对所有RoPE隐藏维度的没有做任何针对性的处理(因为它们对网络有相同的影响)，而其他方法(如YaRN)，我们将其归类为“有针对性的”插值方法

进一步，关于RoPE嵌入的一个有趣的观察是

给定上下文大小 $L$ ，有一些维度 $d$ 的波长长于预训练期间看到的最大上下文长度( $\lambda > L$ )，这表明一些维度的嵌入可能在旋转域中不均匀分布(might not be distributed evenly in therotational domain)
在这种情况下，我们假设拥有所有唯一的位置对意味着绝对的位置信息保持完整(we presume having all unique position pairs implies that theabsolute positional information remains intact)
相反，当波长较短时，只有相对位置信息可以被网络访问(when the wavelength is short, onlyrelative positional information is accessible to the network)
此外，当我们以 $s$ 的比例或使用 $b'$ 的基数将RoPE的所有维度进行拉伸时，所有tokens都变得更接近彼此，因为两个向量的点积旋转较小的量更大(as the dot product of two vectors rotated by a lesser amountis bigger)
这种缩放严重损害了LLM理解其内部嵌入之间的小型和局部关系的能力。我们假设，这种压缩导致模型在邻近标记的位置顺序上被混淆，从而损害模型的能力
This scaling severely impairs a LLM’s ability to understand small and local relationshipsbetween its internal embeddings. We hypothesize that such compression leads to the model beingconfused on the positional order of close-by tokens, and consequently harming the model’s abilities.

为了解决上述问题，选择不插值更高频率的维度，而总是插值更低频率的维度(we choose not to interpolate the higher frequency dimensions at all while always interpolating the lower frequency dimensions)。特别是

如果波长 $\lambda$ 比上下文长度 $L$ 小得多，此时不插值
如果波长 $\lambda$ 等于或大于上下文长度 $L$ ，此时只做插值，不做任何外推(不像以前的“NTK-aware”方法)
两者之间的维度可以兼而有之(dimensions in-between can have a bit of both)，类似于“NTK-aware”插值

因此，在原始上下文大小 $L$ 和波长 $\lambda$ 之间引入比率 $r=\frac{L}{\lambda}$ ，且在第 $d$ 个隐状态下，比率 $r$ 以如下方式依赖于 $d$

$r(d)=\frac{L}{\lambda_{d}}=\frac{L}{2 \pi b^{\prime \frac{2 d}{|D|} \mid}}$

为了确定上述不同插值策略的边界，引入两个额外参数 $\alpha$ 和 $\beta$ ，且针对所有隐藏维度 $d$

如果是 $r(d)<\alpha$ (比如 $\alpha=1$ ，意味着波长大于上下文长度)，则将线性插入一个尺度 $s$ (完全像PI，避免任何外推)
至于如果是 $r(d)>\beta$ 则不插值

接下来，定义斜坡函数 $\gamma$ 为

$\gamma(r)=\left\{\begin{array}{ll} 0, & \text { if } r<\alpha \\ 1, & \text { if } r>\beta \\ \frac{r-\alpha}{\beta-\alpha}, & \text { otherwise } \end{array}\right.$

从而借助该函数，“NTK-by-parts”方法可以定义如下

“NTK-by-parts”插值是对RoPE的一种修改，基于以下函数

$\begin{array}{l} g(m)=m \\ h\left(\theta_{d}\right)=(1-\gamma(r(d))) \frac{\theta_{d}}{s}+\gamma(r(d)) \theta_{d} \end{array}$

$\alpha$ 和 $\beta$ 的值根据具体情况进行调整。当然，通过实验发现，对于Llama家族的模型而言，其中 $\alpha$ 和 $\beta$ 的比较好的取值是 $\alpha=1$ 和 $\beta = 32$

// 待更

2.4.3 "Dynamic NTK"插值

在很多用例中，以从1到最大上下文大小不等的序列长度进行多次前向传递。一个典型的例子是自回归生成，其中序列长度在每一步之后递增1

有两种方法可以应用使用比例因子 $s$ 的插值方法(包括PI、"NTK-aware" and "NTK-by-parts"):

方法1：在整个推理周期中，嵌入层是固定的，包括缩放因子 $s=L^{\prime} / L$ ，其中 $L'$ 是固定数量的扩展上下文大小
方法2：在每次前向传递中，位置嵌入更新缩放因子(the position embedding updates the scale factor)： $s=\max \left(1, l^{\prime} / L\right)$ ，其中 $l'$ 是当前序列的序列长度

上述方法中，方法1的问题在于模型在长度小于 $L$ 时可能出现性能折扣，当序列长度大于 $L'$ 时可能出现突然退化

对此，故提出了方法2，我们称这种推理时间方法为动态缩放方法，当再与“NTK-aware”插值相结合时，我们称之为“动态NTK”插值

一个值得注意的事实是，“动态NTK”插值在 $L$ 上预训练的模型上工作得非常好，而不需要任何微调( $L' =L$ )

第三部分 YaRN全面解析

介绍完“NTK-aware”插值、“NTK-by-parts”插值、"Dynamic NTK"插值之后，接下来即将介绍YaRN(另一种RoPE扩展方法)，这是一种改进的方法，可以有效地扩展使用旋转位置嵌入(RoPE)训练的模型的上下文窗口，包括LLaMA[38]、GPT-NeoX[5]和PaLM[10]家族的模型

3.1 YaRN怎么来的：基于“NTK-by-parts”插值修改注意力

除了前述的插值技术，我们还观察到，在对logits进行softmax操作之前引入温度t可以统一地影响困惑度，无论数据样本和扩展上下文窗口上的token位置如何，更准确地说，我们将注意力权重的计算修改为

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

通过将RoPE重新参数化为一组2D矩阵对，给实现注意力缩放带来了明显的好处(The reparametrization of RoPE as a set of 2D matrices has a clear benefit on the implementation of this attention scaling)

我们可以利用“长度缩放”技巧，简单地将复杂的RoPE嵌入按相同比例进行缩放，使得qm和kn都以常数因子 $\sqrt{1 / t}$ 进行缩放
这样一来，在不修改代码的情况下，YaRN能够有效地改变注意力机制
we can instead use a "length scaling" trick which scales both qm and kn by a constant factor p 1/t by simply scaling the complex RoPE embeddings by the same amount.
With this, YaRN can effectively alter the attention mechanism without modifying its code.
此外，在推理和训练期间，它没有额外开销，因为RoPE嵌入是提前生成并在所有向前传递中被重复使用的。结合“NTK-by-parts”插值方法，我们就得到了YaRN方法
Furthermore, it has zero overhead during both inference and training, as RoPE embeddings are generated in advance and are reused for all forward passes. Combining it with the "NTK-by-parts" interpolation, we have the YaRN method

对于LLaMA和LLaMA 2模型，我们推荐以下值:

$\sqrt{\frac{1}{t}}=0.1 \ln (s)+1$

上式是在未进行微调的LLaMA 7b、13b、33b和65b模型上，使用“NTK-by-parts”方法对各种因素的尺度扩展进行最小困惑度 $\sqrt{1 / t}$ 拟合得到的(The equation above is found by fitting p 1/t at the lowest perplexity against the scale extension by various factors s using the "NTK-by-parts" method)

且相同的t值也适用于Llama 2模型(7b、13b和70b)，这表明熵增加和温度常数t的性质可能具有一定的“普遍性”，并且可以推广到某些模型和训练数据中(It suggests that the property of increased entropy and the temperature constant t may have certain degree of "universality" and may be generalizable across some models and training data)

YaRN方法在微调和非微调场景中均超过以前所有方法，由于其占用空间较小，YaRN与修改注意力机制库(如Flash Attention 2[13])直接兼容
且在对不到0.1%的原始预训练数据进行微调后，YaRN在上下文窗口扩展中达到了最先进的性能
同时，如果YaRN与动态缩放的推理技术相结合而得到的Dynamic-yarn，其允许在超过2倍的上下文窗口扩展，而无需任何微调

3.2 实际应用效果

// 待更

参考文献与推荐阅读

了解几种外推方案做了什么
https://zhuanlan.zhihu.com/p/647145964
https://zhuanlan.zhihu.com/p/642398400
Transformer升级之路：7、长度外推性与局部注意力
Transformer升级之路：10、RoPE是一种β进制编码
大语言模型结构之：RoPE位置外推
大模型上下文长度扩展的一篇综述文献
The What, Why, and How of Context Length Extension Techniques in Large Language Models -- A Detailed Survey
RoPE外推的缩放法则 —— 尝试外推RoPE至1M上下文

创作、修订、完善记录

12.19，开始写本文的前两部分
12.21，修订第二部分..
12.22，新增一节：“2.1.2 位置 $n$ 的旋转位置编码(RoPE)，本质上就是数字 $n$ 的 $\beta$ 进制编码”
且结合苏剑林的博文，补充说明：2.4.1 “NTK-aware”插值：高频外推，低频内插
另，开始写：2.4.2 相对局部距离的损失-“NTK-by-parts”插值，和第三部分 YaRN的全面解析

v_JULY_v

关注

23
点赞
踩
41

收藏

觉得还不错? 一键收藏
打赏
0
评论
大模型上下文扩展之YaRN解析：从直接外推ALiBi、位置插值PI、NTK-aware插值、YaRN

下半年以来，我全力推动我司大模型项目团队的组建，我虽兼管整个项目团队，但为了并行多个项目，最终分成了三个项目组，每个项目都有一个项目负责人，分别为霍哥、阿荀、朝阳，有何问题欢迎随时留言评论，thanks了解几种外推方案做了什么然后再顺着苏剑林文章的思路来看为什么这样做但总觉得不够深刻moe我有了解过GLaM，Mistral那边的没了解过打算了解下，估计也大差不差。
复制链接

扫一扫

专栏目录

pos (0 2 4等偶数维用sin函数计算)	$i$	公式2中 $\beta$ 的指数项
0	$i = 0 // 2 = 0$	0
1	$i = 1 //2 =0$	0
2	$i = 2 // 2 = 1$	1
3	$i = 3 // 2 = 1$	1
4	$i = 4 // 2 = 2$	2
5	$i = 5//2 = 2$	2
6
....
510	$i = 510 // 2 = 255$	$d/2 -1 = (512/2) -1 = 256 - 1 = 255$
511	$i = 511 // 2 = 255$	$d/2 -1 = (512/2) -1 = 256 - 1 = 255$

大模型上下文扩展之YaRN解析：从直接外推ALiBi、位置插值PI、NTK-aware插值、YaRN

前言

第一部分 背景知识：从进制表示谈到直接外推、线性内插、进制转换

1.1 从进制表示到直接外推

1.1.1 进制表示

1.1.2 直接外推

1.2 从线性内插到进制转换

1.2.1 线性内插

1.2.2 进制转换

第二部分 从RoPE、直接外推、位置内插到NTK-aware/NTK-by-parts/Dynamic NTK插值

2.1 旋转位置嵌入

2.1.1 RoPE的快速回顾

2.1.2 位置的旋转位置编码(RoPE)，本质上就是数字的进制编码

2.2 直接外推之ALiBi

2.3 位置内插：基于Positional Interpolation扩大模型的上下文窗口

2.3.1 什么是位置内插Positional Interpolation

2.3.2 位置内插的问题

2.4 从“NTK-aware”插值到“NTK-by-parts”插值、"Dynamic NTK"插值

2.4.1 “NTK-aware”插值：高频外推，低频内插

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

2.4.3 "Dynamic NTK"插值

第三部分 YaRN全面解析

3.1 YaRN怎么来的：基于“NTK-by-parts”插值修改注意力

3.2 实际应用效果

参考文献与推荐阅读

创作、修订、完善记录

“相关推荐”对你有帮助么？

第一部分背景知识：从进制表示谈到直接外推、线性内插、进制转换

第二部分从RoPE、直接外推、位置内插到NTK-aware/NTK-by-parts/Dynamic NTK插值

2.1.2 位置 $n$ 的旋转位置编码(RoPE)，本质上就是数字 $n$ 的 $\beta$ 进制编码