LLM 位置编码及外推

RoPE

https://zhuanlan.zhihu.com/p/629681325

PI

位置插值（POSITION INTERPOLATION）显著改善RoPE的外推能力。你只需要对PT（pretraining)模型fine-turing最多1000步就能实现。PI是通过线性的缩小了输入位置的索引使其匹配原始上下文窗口大小，而不是外推超出训练的上下文长度（超出训练的上下文长度，会导致attention score的爆炸性高分，破坏了分数）。不用PI方法，使用FT（fine-turning）在PT（pretraining）模型上训练长文本，即使训练1万步，能有效扩展的窗口长度也十分有限。

扩展后的模型在长文本任务中有较好的性能，困惑度没有上升，在原来长度的任务中困惑度略有下降（实验中下降了2%）。

考虑使用2048上下文窗口长度预训练的Llama模型。左上角是LLM模型的正常用法:输入位置指数(蓝点)在预训练范围内。右上方显示长度外推，其中模型需要操作未见位置(红点)，最高可达4096。左下角展示了位置插值，我们将位置索引(蓝色和绿色点)本身从[0,4096]缩小到[0,2048]，以迫使它们驻留在预训练的范围内。

RoPE的外推能力有限，理论上，使用2048长度训练的模型，如果输入3000的长度，它就无法使用0位置的信息，但可以使用2900位置的信息，但实际上，如果答案就在2900位置处，也无法得到。

左:一个拟合的注意力评分函数(红色)，形式为公式（1）,
$$d=d_{model}/n_{head}=\frac{4096}{32}=128$$
(LLaMA 7B设置)。圆点为待拟合的随机输入点，红色曲线为最小二乘法拟合的分数函数，其近似在[−1,1]内。右:虽然拟合函数在[0,L]内似乎有很好的边界，其中L = 2048，但在这个区域之外，它可能会超过8000，导致注意力计算出现灾难性问题。请注意，这里我们根本没有挑选:几乎每个从[0,L]内随机生成的输入点集合中学习的曲线都有外推问题。

$$a(s)=Re[\sum _{j=0}^{\frac{d}{2}-1}{h}_j{e}^{is{\theta }_j}]\text{\hspace{1em}(1)}$$

位置插值公式：

$$f^{{}^{\prime }}(x,m)=f(x,\frac{mL}{L^{{}^{\prime }}})$$

$f$ 是RoPE的位置函数，$x$是$m$位置的embedding向量，$L$ 是原来的上下文窗口长度，$L^{{}^{\prime }}$是扩展的窗口长度

passkey retrieval

这是一种测试扩展后窗口的有效长度，即一个token可以注意多远。

12345可以被替换成任何其他的词。

PS: 为了将7B, 13B和33B模型扩展到8192上下文窗口大小，我们使用32个A100 gpu和64个全局批处理大小

对原始2048窗口的模型使用FT，发现训练10000步的时候，只有效扩展到了2560的长度，但是使用PI，训练200步的时候就能扩展到我们想要的窗口长度（实验中有抖动，步数多了有效窗口长度反而变短）。

ALiBi

ALiBi（Attention with Linear Biases）不为token嵌入添加位置嵌入（相比其他位置编码），而是用一个和query, key之间的距离成比例的一个“惩罚项”来偏置query-key的attention score。

效果：可以加快11%的训练速度，以及减少11%的内存使用。

使用ALiBi，LM可以在短序列上进行训练，在长序列上推理。

m叫做坡度（slope）

按照论文中的描述，其做法是：

例如，8个heads的时候，m的取值为：

1/2, 1/4, 1/8, 1/16, 1/32, 1/64, 1/128, 1/256

如果是16个heads，则m的取值为：

1/sqrt(2), 1/2, 1/(2*sqrt(2)), 1/4, …, 1/256

相当于追加了一半的1/sqrt(2)到原来的8个head的每个m的取值。

扩展到一般情况就是：

对于n个head的话，m的取值就是 $2^{\frac{-8}{n}}$：

$2^{\frac{-8}{1}},2^{\frac{-8}{2}},2^{\frac{-8}{3}}...2^{\frac{-8}{n}}$， 这样的m个坡度了。

NTK ALiBi

• 针对ALiBi位置编码在长文本中注意力视野有限的问题，提出了基于ALiBi编码的两种插值方法：内插值和NTK-ALiBi插值。
• 实验表明：无需微调情况下，插值方法能够有效扩大ALiBi编码的注意力视野，提高长文本上的任务效果

NTK-ALiBi插值

• 频域：RoPE编码或ALiBi编码其共同点，都是将位置空间编码为频域空间，其中三角函数（RoPE）或偏置项（ALiBi）的系数，即为频域值。

• NTK-RoPE插值：NTK-RoPE位置编码的改进，在于保持分辨率的情况下（高频），实现了频域空间缩放（低频），从而实现位置空间的插值。

• NTK-ALiBi插值：受NTK编码的启发，我们也可以对ALiBi的频域空间进行缩放，实现NTK-ALiBi的位置插值。改进后的偏置项系数为：

• • $m_h=\frac{1}{2^{8h/H}\ast a^{(h-1)/(H-1)}}$
  • 令$b=a^{1/(H-1)}$, 则有：$m_h=b/(2^{8/H}\ast b{)}^h$
  • NTK改进后可以实现高频分辨率不变，低频视野放大的效果
  • 注：通常可以将缩放系数a设为：a = 推理长度 / 训练长度

• 解释：NTK-ALiBi的公式看起来可能有些难懂，但核心思想与苏建林大佬所说的“高频外推，低频内插”相同。下面从两种情况考虑：

• • h=1时，视野较小，为高频情况。m_h = 1 / 2^(8/H)，与原始偏置系数相同，相当于直接外推，因此是高频外推。高频视野分辨率不变。
  • h=H时，视野较大，为低频情况。m_h = 1 / {2^8 * a}，在原始偏置基础上缩减了a倍，等价于对位置进行了内插值，因此是低频内插。低频视野变大a倍。

参考

https://zhuanlan.zhihu.com/p/632780188

NTK-ALiBi：通过插值实现大模型ALiBi位置编码的长文本外推