深度探讨大模型位置编码及其外推性

作者：王嘉宁，转载请注明出处：https://wjn1996.blog.csdn.net/article/details/120607050

现如今很多大模型都开始支持超过4096长度的推理，例如GPT-4支持超过30k，ChatGLM2-6B也支持最长为32K的文本，但是由于显存资源的限制，这些大模型在真正在训练过程中不一定要训练这么长的文本，通常在预训练时只会设计到4k左右，因此如何确保在模型推理阶段可以支持远远超过预训练时的长度，是目前大模型的核心问题之一，我们将这一问题归为大模型的外推性。

大模型的外推性目前主要在这两个方面考虑，也是提升最有效的两个角度：

寻找或设计合适的位置编码；
设计局部注意力机制。

本文从这两方面进行深度探讨大模型的位置编码和外推性问题。

一、位置编码基础介绍

对于一个token $w_i$ ，其表征向量记作 $\mathbf{x}_i$ ，对于一个句子 $w=\{w_1, \cdots\}$ 则表示为 $\mathbf{x}$ 。那么可以通过一个映射函数将这个句子中的token表征为 $\mathbf{q}_m$ 、 $\mathbf{k}_n$ 、 $\mathbf{v}_n$ ：

其中 $m$ 和 $n$ 表示第 $m$ 和第 $n$ 个token。

1.1 绝对位置编码

在Transformer中，采用正余弦函数来表示绝对位置，公式如下：

这种编码方式又叫做Sinusoidal编码。直觉上看，第 $i$ 个位置的表征向量维度是 $d$ ，这个向量的奇数位置元素使用余弦值，偶数位置元素使用正弦值。
可视化图如下所示：

相邻的位置编码向量很相似，较远的位置编码向量差异很大，说明基于正余弦函数的绝对位置可以表征位置的相关性；
不需要显式地学习位置，提高效率。

最后 $f$ 映射函数可以定义如下所示。即输入表征为token的表征和其对应的绝对位置表征。

该表征通常是直接将位置表征与Word表征直接相加。

1.2 相对位置编码

（1）显式的相对位置
对于第 $m$ 和第 $n$ 个位置的token，其相对位置可以表示为 $r=clip(m-n, r_{min}, r_{max})$ ，即两个token之间的相对距离，且由最大最小值做约束（相对位置不能超过最大值或小于最小值）。

因此，相比于绝对位置，相对位置只需要有 $r_{max}-r_{min}+1$ 个表征向量即可，即在计算两个token之间的attention值时，只需要在attention计算过程中注入这两个位置对应的相对位置 $r$ 对应的相对位置表征向量即可：

这样一来，只需要有限个位置编码，就可以表达出任意长度的相对位置（因为进行了截断），不管是选择可训练式的还是三角函数式的，都可以达到处理任意长度文本的需求。

该表征通常是直接将位置表征与Word表征直接相加。
参考论文：《Self-Attention with Relative Position Representations》

（2）Transformer-XL（XLNet）

将第 $m$ 和第 $n$ 个位置的QK计算进行了分解。引入了一些可学习的参数：

其中 $\mathbf{p}_m$ 表示待学习的第 $m$ 个token的相对位置向量， $\mathbf{p}_n$ 表示待学习的第 $n$ 个token的相对位置向量。通过分解让相对位置注入在attention的计算过程中。

该表征通常是在Attention计算过程中融入绝对位置。

（3）Transformer-XL的改进：
第二和第四项使用相对位置表征 $\tilde{\mathbf{p}}_{m-n}$ 来替换绝对位置表征。同时加入新的可训练参数 $\mathbf{u}$ 和 $\mathbf{v}$ 。
该表征方法在T5模型中被首次使用，参考论文：《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》

该表征通常是在Attention计算过程中融入相对位置。

（4）可训练偏置项

Transformer-XL中的位置表征是将QK分解为4项，而后面3个够项都有跟位置有关的参数，可以直接将后面的3项之和抽象为一个偏置：

进一步改造还可以得到：

以及：

这三种方法均是对Transformer-XL的表征形式进行改进。

二、RoPE旋转位置编码

RoPE（Rotary Position Embedding）的出发点就是“通过绝对位置编码的方式实现相对位置编码”，或者可以说是实现相对位置编码和绝对位置编码的结合。这样做既有理论上的优雅之处，也有实践上的实用之处，比如它可以拓展到线性Attention中就是主要因为这一点。

2.1 复数的引入

假设 $\mathbf{q}_m$ 和 $\mathbf{k}_n$ 是对应位置的二维行向量（即每个位置都有两个向量来表征位置），因此这个二维向量可以用复数来代替（包括实部和虚部），因此他们的内积可以作为其对应的Attention值。内积的计算可以由两个复数实现： $\langle\mathbf{q}_m, \mathbf{k}_n\rangle=RE[\mathbf{q}_m\mathbf{k}_n^*]$ ，其中 $^*$ 表示共轭复数， $RE []$ 表示取复数中的实部。

两个二维向量的内积，等于把它们当复数看时，一个复数与另一个复数的共轭的乘积的实部。

因此当分别对 $\mathbf{q}_m$ 和 $\mathbf{k}_n$ 融入绝对位置时，即得到：

RoPE求解过程推导了如何确定获得每个位置编码是 $\mathbf{q}_me^{im\theta}$ 。

可以发现，当乘以绝对位置 $e^{im\theta}$ 和 $e^{in\theta}$ 时，等价于复数运算中乘以 $e^{i(m-n)\theta}$ ，即相当于在复数空间中是相对位置 $m - n$ ，这样就巧妙地通过复数运算的形式将绝对位置转换为相对位置。

复数乘法的几何意义是向量的旋转，假设 $f_{q,k}(\mathbf{x}, m)$ 表示向量 $\mathbf{x}$ 在 $m$ 位置的位置编码，则有：

其中 $\mathbf{q}=\mathbf{W}_q\mathbf{x}_m$ ， $\mathbf{k}=\mathbf{W}_k\mathbf{x}_n$

等价于

后面两项的乘积本质就是向量 $\mathbf{q}$ （或 $\mathbf{k}$ ）的两个二维行向量。

当向量维度为 $d$ 时（ $d$ 为偶数），则可以扩展为：

每两个为一组二维向量，一共有 $d /2$ 个组合，可以直接拼接作为 $d$ 维度的旋转位置编码。
二维扩展到多维的解读：Transformer升级之路：4、二维位置的旋转式位置编码 - 科学空间|Scientific SpacesTransformer升级之路：6、旋转位置编码的完备性分析 - 科学空间|Scientific Spaces

2.2 RoPE的实现

当输入一个句子“Enhanced Transformer with Rotary Position Embedding”时，首先获得其Query和Key向量 $\mathbf{q}$ 、 $\mathbf{k}$ ，其对应的维度均为 $d$ ，然后对于向量中相邻两个元素为一组，可以得到 $d /2$ 组（图中左下部分相同颜色的两个元素作为一组，对于每一组，一个文本则可以得到两个行向量）；
获得每个词的绝对位置编号（该句子由6个词，位置编号分别为1，2，3，4，5，6），假设取“Enhanced”单词为例，其第一组元素为 $\theta_1$ ，位置为 $m = 1$ ，那么通过旋转位置编码可以的到新的元素值 $x_1', x_2'$ 。
所有单词的 $d /2$ 个组合都按照这种形式进行“旋转”，即可得到新的位置编码（右下角）

RoPE一种线性的实现如下所示：

RoPE的性质

（1）远程衰减

从图中我们可以可以看到随着相对距离的变大，内积结果有衰减趋势的出现。因此，选择 $θ_i=10000^{−2i/d}$ ，确实能带来一定的远程衰减性。当然，能带来远程衰减性的不止这个选择，几乎任意的光滑单调函数都可以。如果以 $θ_i=10000^{−2i/d}$ 为初始化，将 $θ_i$ 视为可训练参数，然后训练一段时间后发现 $θ_i$ 并没有显著更新，因此干脆就直接固定 $θ_i=10000^{−2i/d}$ 了。

（2）优势

用一个旋转矩阵rotation matrix来对绝对位置进行编码，于此同时，meanwhile；
在自注意力机制中导入显式的位置依赖。
自由的序列长度；
随着相对位置的增大，而逐步延缓退化（=衰减）的inter-token dependency；
用相对位置编码来“武装”线性自注意力。

具体来说，RoPE 使用旋转矩阵对绝对位置进行编码，同时将显式的相对位置依赖性纳入自注意公式中。【核心的两个点，一个是“旋转矩阵”，一个是“显式的相对位置依赖”】。

三、长文本外推性

外推性的含义是在长文本表征过程中，如何在训练阶段只需要学习有限的长度，即可以在推理阶段能够延伸长度至若干倍且依然保持不错的性能和效果。
长度外推性是一个训练和预测的长度不一致的问题，主要体现在两个方面：

预测的时候用到了没训练过的位置编码（不论是绝对位置还是相对位置）；
预测的时候注意力机制所处理的token数量远超训练时的数量。

解决长文本外推性问题的一个简单有效的方法是Attention Mask，如图所示：

通过类似滑动窗口的结构，约束一个每个token只能对局部区域的token计算Attention值，因此对于相对位置大小不会超过窗口大小，解决了第一个问题；
Attention只会在窗口内计算，避免了对大量的token的Attention进行加权平均导致最终权重过度“平滑”现象。

在实现过程中，本质上是在计算完 $\mathbf{q}_m^T\mathbf{k}_n$ 之后减去一个矩阵，即 $\mathbf{q}_m^T\mathbf{k}_n-M$ ，其中 $M$ 的形状如下图所示：

可以看出，蓝色区域（即滑动窗口内的局部区域）为0，说明保持原始的Attention归一化前的值；其他区域则为一个INT内最大的整数，说明Attention值是一个非常小的数（在softmax归一化后几乎为0）。

3.1 ALIBI

论文：《Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation》

其与上面的思想一样，只是改进了上面的 $M$ 矩阵为 $\lambda|m-n|$ ，即Attention在归一化前的计算为： $\mathbf{q}_m^T\mathbf{k}_n-\lambda|m-n|$ ，其中 $\lambda>0$ 为超参数，Transformer的多头注意力中的每个头的值可设置不同。矩阵 $\lambda|m-n|$ 的形状如下所示：

相比于原始的方法，相对距离越长， $\lambda|m-n|$ 值就越大，越远的位置Attention值被归一化后就约小，相对于“滑动窗口”采用的方法是hard（在窗口内就计算attention，不在窗口内不计算），AIBLI是比较soft的（离得近attention就会比较大，离得远就比较小）。

3.2 KERPLE

论文：《KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation》

其对ALIBI进行了一些改进，引入了两个可学习的参数 $r_1$ 和 $r_2$ 来“动态”学习局部区域。如下图所示，左侧为原始的 $\mathbf{q}_m^T\mathbf{k}_n$ ，通过引入参数来动态减去AIBLI中的 $\lambda|m-n|$ 矩阵：

定义了两种模式，分别是power和logarithmic，分别对应没有对数和有对数的形式：

在logarithmic模式中， $r_1$ 控制了整体的量， $r_2$ 相当于ALIBI中的 $\lambda$ ， $c$ 是一个常数。苏神版简化写作：

3.3 Sandwich

论文：《Receptive Field Alignment Enables Transformer Length Extrapolation》

Sandwich与KEPRLE是同一个作者提出的，其对KEPRLE进行了少量改进，即对应的公式改写为： $\mathbf{q}_m^T\mathbf{k}_n+\lambda\mathbf{p}_m^T\mathbf{p}_n$ ，其中 $\mathbf{p}_m$ 和 $\mathbf{p}_n$ 可以使用Sinusoidal编码表示，即：

由于Sinusoidal编码在单调性上等价于 $∣ m - n ∣$ ，都是线性递增形式，因此Sandwich只是改头换面了。

3.4 XPOS

论文：《A Length-Extrapolatable Transformer》
参考解读：Transformer升级之路：7、长度外推性与局部注意力

其在RoPE的基础上引入了局部注意力。RoPE的本质是：

其中：

在第2大节中已经介绍了RoPE是通过使用复数来实现绝对位置表示相对位置的。XPOS通过引入一个新的标量 $\xi$ ，即有：

由于RoPE相对位置是 $m - n$ ，而不是 $∣ m - n ∣$ ，因此XPOS约束在了单向Transformer，从而避免了负数出现。

XPOS又设计了一个局部感知的注意力机制Blockwise Causal Attention，进一步提升了局部注意力的性能，提升了长文本的外推性。

四、外推性的其他探索

（1）混合注意力Mask

在解决长文本位置表征时，典型的代表有Transformer-XL、BigBird、LongFormer，他们除了局部注意力机制以外，还引入了随机位置的性质：

如上图，第2张图为局部注意力（滑动窗口），第3章图为有限的全局感知（例如只限制前两个token可以看见所有的token）。而第一张图则是随机mask，以缓解过度hard的局部注意力。三者注意力混合起来后得到第四张图，这也是普遍训练超长文本大模型时采用的方法。

（2）随机位置表征

论文：《Randomized Positional Encodings Boost Length Generalization of Transformers》

绝对位置表征时，会存在位置上的OOV问题，随机位置编码则是通过在训练过程中采用如下策略：

对应的代码也很简单：

def random_position_ids(N, L=2048):
    """从[0, L)中随机不重复挑N个整数，并从小到大排列
    """
    return np.sort(np.random.permutation(L)[:N])

苏神对随机位置编码的新探索：

对应的代码为：

def random_position_ids(N):
    """先随机采样n，然后从[0, n]均匀取N个点
    """
    n = sample_from_xxx()
    return np.linspace(0, 1, N) * n

（3） $\log n$ Attention Scale

原始的Attention计算公式为：

$V)=softmax(\frac{QK^T}{\sqrt{d_k}})V$

只需要简单的改成下面即可：

$V)=softmax(\frac{\log_{m}n}{\sqrt{d_k}}QK^T)V$

其中 $m$ 是训练时的最大长度， $n$ 是预测时的位置，通常来说 $n >> m$ 。直觉来看，就是直接在计算 $QK^T$ 时，根据其相对位置来控制Attention的值。当 $m$ 和 $n$ 距离很远时， $log_mn$ 的值会很大，会使得整体Attention归一化后会比较平缓，有助于解决外推性问题。

（4）全局依赖

滑动窗口的方法如果在一层Transformer Layer里看，本质上类似长度为 $w$ 的N-Gram模型，即如下图所示：

如果Transformer又 $L$ 层，那么，从输入层开始，长度为 $w$ 的窗口内的信息，可以在经过 $L$ 层之后传给一个更广的区域，区域长度为 $(w - 1) L + 1$ ，如下图所示：

苏神给出的一种新的想法，就是假设我有 $L$ 层Transformer，则可以在前 $L - 1$ 层利用这种扩张特性，得到最终 $(w - 1) (L - 1) + 1$ 长度的区域后，在最后一层采用上面提到的 $\log n$ Attention Scale方法，讲前 $L - 1$ 层扩张的信息快速在最后一层与所有token进行交互。引入苏神的原文为：
这种局部注意力+Attention Scale的结合也是一种很巧妙的idea。实验也发现这种策略的外推性很惊艳。

参考文献：

《Self-Attention with Relative Position Representations》
ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING
Transformer升级之路：4、二维位置的旋转式位置编码 - 科学空间|Scientific SpacesTransformer升级之路：6、旋转位置编码的完备性分析 - 科学空间|Scientific Spaces
《Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation》
《KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation》
《A Length-Extrapolatable Transformer》
《Randomized Positional Encodings Boost Length Generalization of Transformers》
Transformer升级之路：7、长度外推性与局部注意力
【OpenLLM 009】大模型基础组件之位置编码-万字长文全面解读LLM中的位置编码与长度外推性（上）
Transformer升级之路：8、长度外推性与位置鲁棒性 - 科学空间|Scientific Spaces
Transformer升级之路：9、一种全局长度外推的新思路 - 科学空间|Scientific Spaces
Bias项的神奇作用：RoPE + Bias = 更好的长度外推性 - 科学空间|Scientific Spaces

博客记录着学习的脚步，分享着最新的技术，非常感谢您的阅读，本博客将不断进行更新，希望能够给您在技术上带来帮助。