RoPE - LLM 的位置信息指北针 | SLie's Blog|琴弦之轮

在自然语言处理（NLP）领域，大语言模型（LLM）近年来取得了显著的进展，尤其是在 Transformer 架构的推动下。然而，Transformer 模型本身存在一个关键问题：它无法天然地理解序列中各个词语的位置信息。这是因为自注意力机制（Self-Attention）主要依赖于词语本身的内容，而非其在序列中的位置。为了解决这一问题，位置编码（Positional Encoding）应运而生。本文将简单探讨位置编码的核心概念以及不同编码方法的实现。

为什么需要位置编码

在 NLP 中，序列的顺序和位置信息至关重要。例如，句子 张三打李四 与 李四打张三 在语义上是完全不同，仅仅是由于词语的位置发生了变化。然而，原始的 Transformer 模型无法区分这些位置信息，因为它仅依赖于词嵌入（Word Embedding）来表示词语，得到词语的 token。随后的 Self-attention 操作将并行地处理 token 与 token 之间的关联信息，而这个操作是对称的，无法与位置信息关联。

具体来说，关于序列的函数 $f$ 具有置换不变性，满足：

$f(\cdots,\boldsymbol{x}_m,\cdots,\boldsymbol{x}_n,\cdots)=f(\cdots,\boldsymbol{x}_n,\cdots,\boldsymbol{x}_m,\cdots)$

其中 $\boldsymbol{x}_m,\boldsymbol{x}_n$ 分别表示第 $m,n$ 个输入序列 token。将二者进行交换该函数值不变。

而位置编码的引入通过在词嵌入 $\boldsymbol{x}_k$ 中添加位置编码向量 $\boldsymbol{p}_k$ ，打破了这种对称性。以“在每个位置上都加上一个不同的编码向量”的这个方案为例：

$\tilde{f}(\cdots,\boldsymbol{x}_m,\cdots,\boldsymbol{x}_n,\cdots)=f(\cdots,\boldsymbol{x}_m + \boldsymbol{p}_m,\cdots,\boldsymbol{x}_n + \boldsymbol{p}_n,\cdots)$

当然除了以上这个简单的策略，对于位置编码总体来说还可以划分为两大类：

想办法将位置信息融入到输入中，这构成了绝对位置编码的一般做法；
想办法微调 Attention 的结构，使得它有能力分辨不同位置的 Token，这构成了相对位置编码的一般做法。

我们将在后文继续介绍这两类编码方式。总而言之，说位置编码通过引入额外的信息使得模型能够更好地理解序列的上下文关系，提升了模型的表现。

外推性与显式编码

在 Transformer 的原论文中，位置编码向量 $\boldsymbol p$ 的计算方式是使用 Sinusoidal 函数：

$\left\{\begin{aligned}&\boldsymbol{p}_{k,2i}=\sin\Big(k/10000^{2i/d}\Big)\\ &\boldsymbol{p}_{k, 2i+1}=\cos\Big(k/10000^{2i/d}\Big) \end{aligned}\right.$

其中 $\boldsymbol{p}_{k,2i},\boldsymbol{p}_{k,2i+1}$ 分别表示位置 $k$ 的第 $2i,2i+1$ 个分量， $d$ 是位置向量的维度。位置编码以相加的方式加在输入 token 当中。

按上一节的分类来说这是一种绝对位置编码，因为它为每一个位置的输入都给出了专属的固定的 Embedding。同时，这也是一种显式编码，因为位置编码由一个确定的函数生成。与之对应的，还有训练式编码。

训练式编码同样是一种绝对位置编码，但是采用的是一种最朴素方案：不特意去设计显式的函数，而是直接将位置编码当作可训练参数加上去。比如序列最大长度为512，编码维度为768，那么就初始化一个 $512×768$ 的矩阵作为位置向量，让它随着训练过程更新。现在的 BERT、GPT 等模型所用的就是这种位置编码，事实上它还可以追溯得更早，比如2017年 Facebook 的《Convolutional Sequence to Sequence Learning》就已经用到了这种方案。

一般的认为训练式编码的缺点是没有外推性。外推性是指大模型在训练时和预测时的输入长度不一致，导致模型的泛化能力下降的问题。例如，如果一个模型在训练时只使用了 512 个 token，那么在预测时如果输入超过 512 个 token，模型可能无法正确处理（当然，也可以将超过512的位置向量随机初始化，然后继续微调）。这就限制了大模型在处理长文本或多轮对话等任务时的效果。

绝对位置和相对位置

绝对位置编码

一般来说，绝对位置编码会加到输入中：在输入的第 $k$ 个向量 $\boldsymbol{x}_k$ 中加入位置向量 $\boldsymbol{p}_k$ 变为 $\boldsymbol{x}_k+\boldsymbol{p}_k$ ，其中 $\boldsymbol{p}_k$ 只依赖于位置编号 $k$ 。形式上来看，绝对位置编码是相对简单的一种方案，但即便如此，也不妨碍各路研究人员的奇思妙想，也有不少的变种。

三角式：如 Sinusoidal 位置编码
训练式：从略
递归式：FLOATER 模型. 通过 $\boldsymbol{p}_{k+1}=f(\boldsymbol{p}_k)$ 递归地得到编码向量，更极致地，FLOATER 模型使用微分方程（ODE） $d\boldsymbol{p}_t/dt=\boldsymbol{h}(\boldsymbol{p}_t,t)$ 来建模位置编码，函数 $h$ 使用神经网络模拟，故而这种微分方程也被称为神经微分方程。（📃原论文）
相乘式：使用哈达玛积（按位相乘）来代替 $\boldsymbol{x}_k+\boldsymbol{p}_k$ （乘性位置编码 - 知乎📃）

相对位置编码

相对位置编码顾名思义关注词语之间的相对距离，具体来说是在算 Attention 的时候考虑当前位置与被Attention的位置的相对距离。显然相对位置编码的外推性更强（更适合处理变长的序列），并且它的灵活性更大，更加体现出了研究人员的“天马行空”。

一般认为，相对位置编码是由绝对位置编码启发而来，考虑一般的带绝对位置编码的 Attention：

$\left\{\begin{aligned} \boldsymbol{q}_i =&\, (\boldsymbol{x}_i + \boldsymbol{p}_i)\boldsymbol{W}_Q \\ \boldsymbol{k}_j =&\, (\boldsymbol{x}_j + \boldsymbol{p}_j)\boldsymbol{W}_K \\ \boldsymbol{v}_j =&\, (\boldsymbol{x}_j + \boldsymbol{p}_j)\boldsymbol{W}_V \\ a_{i,j} =&\, \text{softmax}\left(\boldsymbol{q}_i \boldsymbol{k}_j^{\top}\right)\\ \boldsymbol{o}_i =&\, \sum_j a_{i,j}\boldsymbol{v}_j \end{aligned}\right.$

其中 $\text{softmax}(\cdot)$ 对 $j$ 那一维归一化。注意，这里的向量都是指行向量。

Attention 最核心的计算就在于计算 $\boldsymbol{q}_i \boldsymbol{k}_j^{\top}$ ，我们将其展开：

$\begin{aligned} \boldsymbol{q}_i \boldsymbol{k}_j^{\top} &=(\boldsymbol{x}_i + \boldsymbol{p}_i)\boldsymbol{W}_Q\big((\boldsymbol{x}_j + \boldsymbol{p}_j)\boldsymbol{W}_K\big)^{\top}\qquad(\S1)\\ &= \left(\boldsymbol{x}_i \boldsymbol{W}_Q + \boldsymbol{p}_i \boldsymbol{W}_Q\right)\left(\boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top} + \boldsymbol{W}_K^{\top}\boldsymbol{p}_j^{\top}\right) \\ &= \underbrace{\boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top}}_{token-token} + \underbrace{\boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{p}_j^{\top}}_{token-pos} + \underbrace{\boldsymbol{p}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top}}_{pos-token} + \underbrace{\boldsymbol{p}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{p}_j^{\top}}_{pos-pos}\qquad(\S2) \end{aligned}$

针对上式的改动也使得相对位置编码存在各种变体。

原初式

相对位置编码起源于Google的论文《Self-Attention with Relative Position Representations》，华为开源的 NEZHA 模型也用到了这种位置编码。

具体来说，它将上式中的 $\S1$ 部分中关于 $\boldsymbol{q}_i$ 的位置编码 $\boldsymbol{p}_i$ 去掉，然后把关于 $\boldsymbol{k}_j$ 的第二项 $\boldsymbol{p}_j \boldsymbol{W}_K$ 整体替换成了一个二元位置向量 $\boldsymbol{R}_{i,j}^{K}$ ，即：

$\boldsymbol{q}_i \boldsymbol{k}_j^{\top} = \boldsymbol{x}_i \boldsymbol{W}_Q\big(\boldsymbol{x}_j\boldsymbol{W}_K + {\color{green}{\boldsymbol{R}_{i,j}}^{K}}\big)^{\top}$

另外，最后的输出部分也把关于 $\boldsymbol{v}_i$ 的 $\boldsymbol{p}_j \boldsymbol{W}_V$ 也整体替换成了一个二元位置向量 $\boldsymbol{R}_{i,j}^{V}$ ，即：

$\boldsymbol{o}_i = \sum_j a_{i,j}\left(\boldsymbol{x}_j\boldsymbol{W}_V + {\color{green}{\boldsymbol{R}_{i,j}^{V}}}\right)$

之所以称为相对位置，就是因为该操作将本来依赖于二元坐标 $(i,j)$ 的向量 $\boldsymbol{R}_{i,j}^{K}$ 和 $\boldsymbol{R}_{i,j}^{V}$ 改为了只依赖于相对距离 $i-j$ 的了。并且通常来说会进行截断，以适应不同任意的距离：

$\begin{aligned} \boldsymbol{R}_{i,j}^{K} = \boldsymbol{p}_K\left[\text{clip}(i-j, p_{\min}, p_{\max})\right]\\ \boldsymbol{R}_{i,j}^{V} = \boldsymbol{p}_V\left[\text{clip}(i-j, p_{\min}, p_{\max})\right] \end{aligned}$

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

XLNET式

XLNET式位置编码源自 Transformer-XL 架构，它的做法是在上式中的 $\S2$ 部分调整所有涉及到位置项，即后三项（token-pos、pos-token 和 pos-pos）。

具体来说，它将 $\boldsymbol{p}_j$ 直接替换成相对位置向量 $\boldsymbol{R}_{i-j}$ ，并且直接用Sinusoidal式的生成方案赋值。至于 pos-token项和 pos-pos项中的 $\boldsymbol{p}_i\boldsymbol{W}_Q$ 干脆直接用两个新的可学习向量 $\boldsymbol{u,v}$ 取代。最终得到：

$\boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top} + \boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_{K,R}^{\top}{\color{green}{\boldsymbol{R}_{i-j}^{\top}}} + {\color{red}{\boldsymbol{u}}}\boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top} + {\color{red}{\boldsymbol{v}}} \boldsymbol{W}_{K,R}^{\top}{\color{green}{\boldsymbol{R}_{i-j}^{\top}}}$

值得注意的是，因为相对位置向量 $\boldsymbol{R}_{i-j}$ 的编码空间可能和 $\boldsymbol x_j$ 不同，因此不能都用同一个参数 $\boldsymbol{W}_K$ 对其进行投影，所以针对它单独使用一个命名为 $\boldsymbol{W}_{K,R}$ 的矩阵。此外， $\boldsymbol{v}_j$ 上的位置偏置干脆直接去掉了，最终输出项就只保留：

$\boldsymbol{o}_i = \sum_j a_{i,j}\left(\boldsymbol{x}_j\boldsymbol{W}_V \right)$

并且似乎从这个工作开始，后面的相对位置编码都只加到Attention矩阵上去，而不加到 $\boldsymbol{v}_j$ 上去了。

T5式

T5模型出自文章《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》，里边用到了一种更简单的相对位置编码。它认为输入信息与位置信息应该是独立（解耦）的，那么它们就不应该有过多的交互，从而可以将 $\S2$ 式直接删去 token-pos 和 pos-token 两项。至于最后一项 $\boldsymbol{p}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{p}_j^{\top}$ 本质上也只是一个依赖于 $(i,j)$ 的标量，可以把它直接作为一个可学习的参数 $\beta_{i,j}$ 训练出来。得到：

$\boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top} + {\color{green}{\boldsymbol{\beta}_{i,j}}}$

说白了，它仅仅是在Attention矩阵的基础上加一个可训练的偏置项而已，并且输出项 $\boldsymbol{v}_j$ 上的位置偏置也是直接去掉了。比较“别致”的是，不同于常规位置编码对将 $\beta_{i,j}$ 视为 $i-j$ 的函数然后截断，T5对其进行了一种“哈希分组”。直观理解就是，比较邻近的位置（0～7），因为需要比较得精细一些，所以给它们都分配一个独立的位置编码，至于稍远的位置（比如8～11），我们不用区分得太清楚，所以它们可以共用一个位置编码，距离越远，共用的范围就可以越大，直到达到指定范围再clip。

DeBERTa式

DeBERTa 论文为《DeBERTa: Decoding-enhanced BERT with Disentangled Attention》，与 T5 相反，它扔掉了第4项，保留第2、3项并且替换为相对位置编码：

$\boldsymbol{q}_i \boldsymbol{k}_j^{\top} = \boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top} + \boldsymbol{x}_i \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}{\color{green}{\boldsymbol{R}_{i,j}^{\top}}} + {\color{green}{\boldsymbol{R}_{j,i}}} \boldsymbol{W}_Q \boldsymbol{W}_K^{\top}\boldsymbol{x}_j^{\top}$

不过，DeBERTa比较有意思的地方，是提供了使用相对位置和绝对位置编码的一个新视角，它指出NLP的大多数任务可能都只需要相对位置信息，但确实有些场景下绝对位置信息更有帮助，于是它将整个模型分为两部分来理解。以 Base 版的 MLM 预训练模型为例，它一共有13层，前11层只是用相对位置编码，这部分称为 Encoder，后面2层加入绝对位置信息，这部分它称之为Decoder，还弄了个简称 EMD（Enhanced Mask Decoder）；至于下游任务的微调截断，则是使用前11层的Encoder加上1层的Decoder来进行。

旋转位置编码

旋转位置编码（Rotary Positional Encoding, RoPE）是一种结合了绝对和相对位置信息的编码方法，这是一种配合Attention机制能达到“绝对位置编码的方式实现相对位置编码”的设计。而也正因为这种设计，它还是目前唯一一种可用于线性Attention的相对位置编码。RoPE 于苏剑林大神和他们实验室自研的 Rotary Transformer（RoFormer）模型中首次使用。

参考

（pytorch进阶之路）四种Position Embedding的原理及实现-CSDN博客