Transformer升级之路：博采众长的旋转式位置编码

最新推荐文章于 2024-09-11 11:50:00 发布

PaperWeekly

最新推荐文章于 2024-09-11 11:50:00 发布

阅读量1k

点赞数

文章标签：人工智能深度学习机器学习计算机视觉自然语言处理

原文链接：https://kexue.fm/archives/8265

版权

©PaperWeekly 原创 · 作者｜苏剑林

单位｜追一科技

研究方向｜NLP、神经网络

上一篇文章中，我们对原始的 Sinusoidal 位置编码做了较为详细的推导和理解，总的感觉是 Sinusoidal 位置编码是一种“想要成为相对位置编码的绝对位置编码”。

一般来说，绝对位置编码具有实现简单、计算速度快等优点，而相对位置编码则直接地体现了相对位置信号，跟我们的直观理解吻合，实际性能往往也更好。

由此可见，如果可以通过绝对位置编码的方式实现相对位置编码，那么就是“集各家之所长”、“鱼与熊掌兼得”了。Sinusoidal 位置编码隐约做到了这一点，但并不够好。

本文将会介绍我们自研的 Rotary Transformer（RoFormer）模型，它的主要改动是应用了笔者构思的“旋转式位置编码（Rotary Position Embedding，RoPE）”，这是一种配合 Attention 机制能达到“绝对位置编码的方式实现相对位置编码”的设计。而也正因为这种设计，它还是目前唯一一种可用于线性 Attention 的相对位置编码。

RoFormer：

https://github.com/ZhuiyiTechnology/roformer

基本思路

在之前的文章让研究人员绞尽脑汁的 Transformer 位置编码中我们就简要介绍过 RoPE，当时称之为“融合式”，本文则更加详细地介绍它的来源与性质。在 RoPE 中，我们的出发点就是“通过绝对位置编码的方式实现相对位置编码”，这样做既有理论上的优雅之处，也有实践上的实用之处，比如它可以拓展到线性 Attention 中就是主要因为这一点。

为了达到这个目的，我们假设通过下述运算来给添加绝对位置信息：

也就是说，我们分别为设计操作，使得经过该操作后，就带有了位置 m, n 的绝对位置信息。Attention 的核心运算是内积，所以我们希望的内积的结果带有相对位置信息，因此假设存在恒等关系：

所以我们要求出该恒等式的一个（尽可能简单的）解。求解过程还需要一些初始条件，显然我们可以合理地设和。

求解过程

同上一篇思路一样，我们先考虑二维情形，然后借助复数来求解。在复数中有，所以我们有：

简单起见，我们假设存在复数，使得，然后我们用复数的指数形式，设：

那么代入方程后就得到方程组：

对于第一个方程，代入 m=n 得到：

最后一个等号源于初始条件和。所以现在我们可以很简单地设，即它不依赖于 m。至于第二个方程，同样代入 m=n 得到：

这里的是本身的幅角，最后一个等号同样源于初始条件。根据上式得到，所以应该是一个只与 m 相关、跟无关的函数，记为，即。接着代入 n=m-1，整理得到：

即是等差数列，设右端为，那么就解得。

编码形式

综上，我们得到二维情况下用复数表示的 RoPE：

根据复数乘法的几何意义，该变换实际上对应着向量的旋转，所以我们称之为“旋转式位置编码”，它还可以写成矩阵形式：

由于内积满足线性叠加性，因此任意偶数维的 RoPE，我们都可以表示为二维情形的拼接，即：

也就是说，给位置为 m 的向量乘上矩阵、位置为 n 的向量乘上矩阵，用变换后的序列做 Attention，那么 Attention 就自动包含相对位置信息了，因为成立恒等式：

值得指出的是，是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性。

由于的稀疏性，所以直接用矩阵乘法来实现会很浪费算力，推荐通过下述方式来实现 RoPE：

其中是逐位对应相乘，即 Numpy、Tensorflow 等计算框架中的*运算。从这个实现也可以看到，RoPE 可以视为是乘性位置编码的变体。

远程衰减

可以看到，RoPE 形式上和 Sinusoidal 位置编码有点相似，只不过 Sinusoidal 位置编码是加性的，而 RoPE 可以视为乘性的。在的选择上，我们同样沿用了 Sinusoidal 位置编码的方案，即，它可以带来一定的远程衰减性。

具体证明如下：将两两分组后，它们加上 RoPE 后的内积可以用复数乘法表示为：

记，并约定，那么由 Abel 变换（分部求和法）可以得到：

所以：

因此我们可以考察随着相对距离的变化情况来作为衰减性的体现， Mathematica 代码如下：

结果如下图：

▲ RoPE的远程衰减性（d=128）

从图中我们可以可以看到随着相对距离的变大，内积结果有衰减趋势的出现。因此，选择，确实能带来一定的远程衰减性。当然，同上一篇文章说的一样，能带来远程衰减性的不止这个选择，几乎任意的光滑单调函数都可以，这里只是沿用了已有的选择而已。

笔者还试过以为初始化，将视为可训练参数，然后训练一段时间后发现并没有显著更新，因此干脆就直接固定了。

线性场景

最后，我们指出，RoPE 是目前唯一一种可以用于线性 Attention 的相对位置编码。这是因为其他的相对位置编码，都是直接基于 Attention 矩阵进行操作的，但是线性 Attention 并没有事先算出 Attention 矩阵，因此也就不存在操作 Attention 矩阵的做法，所以其他的方案无法应用到线性 Attention 中。

而对于 RoPE 来说，它是用绝对位置编码的方式来实现相对位置编码，不需要操作 Attention 矩阵，因此有了应用到线性 Attention 的可能性。

关于线性 Attention 的介绍，这里不再重复，有需要的读者请参考线性 Attention 的探索：Attention 必须有个 Softmax 吗？。线性 Attention 的常见形式是：

其中是值域非负的激活函数。可以看到，线性 Attention 也是基于内积的，所以很自然的想法是可以将 RoPE 插入到内积中：

但这样存在的问题是，内积可能为负数，因此它不再是常规的概率注意力，而且分母有为 0 的风险，可能会带来优化上的不稳定。考虑到都是正交矩阵，它不改变向量的模长，因此我们可以抛弃常规的概率归一化要求，使用如下运算作为一种新的线性 Attention：

也就是说，RoPE 只插入分子中，而分母则不改变，这样的注意力不再是基于概率的（注意力矩阵不再满足非负归一性），但它某种意义上来说也是一个归一化方案，而且也没有证据表明非概率式的注意力就不好（比如 Nyströmformer 也算是没有严格依据概率分布的方式构建注意力），所以我们将它作为候选方案之一进行实验，而我们初步的实验结果显示这样的线性 Attention 也是有效的。

此外，笔者在线性 Attention 的探索：Attention 必须有个 Softmax 吗？中还提出过另外一种线性 Attention 方案：，它不依赖于值域的非负性，而 RoPE 也不改变模长，因此 RoPE 可以直接应用于此类线性 Attention，并且不改变它的概率意义。

模型开源

RoFormer 的第一版模型，我们已经完成训练并开源到了 Github 中。

简单来说，RoFormer 是一个绝对位置编码替换为 RoPE 的 WoBERT [1] 模型，它跟其他模型的结构对比如下：

在预训练上，我们以 WoBERT Plus 为基础，采用了多个长度和 batch size 交替训练的方式，让模型能提前适应不同的训练场景：

从表格还可以看到，增大序列长度，预训练的准确率反而有所提升，这侧面体现了 RoFormer 长文本语义的处理效果，也体现了 RoPE 具有良好的外推能力。

在短文本任务上，RoFormer 与 WoBERT 的表现类似，RoFormer 的主要特点是可以直接处理任意长的文本。下面是我们在 CAIL2019-SCM 任务上的实验结果：

其中后面的参数是微调时截断的 maxlen，可以看到 RoFormer 确实能较好地处理长文本语义，至于设备要求，在 24G 显存的卡上跑 maxlen=1024，batch_size 可以跑到 8 以上。目前中文任务中笔者也就找到这个任务比较适合作为长文本能力的测试，所以长文本方面只测了这个任务，欢迎读者进行测试或推荐其他评测任务。

当然，尽管理论上 RoFormer 能处理任意长度的序列，但目前 RoFormer 还是具有平方复杂度的，我们也正在训练基于线性 Attention 的 RoFormer 模型，实验完成后也会开源放出，请大家期待。

文章小结

本文介绍了我们自研的旋转式位置编码 RoPE 以及对应的预训练模型 RoFormer。从理论上来看，RoPE 与 Sinusoidal 位置编码有些相通之处，但 RoPE 不依赖于泰勒展开，更具严谨性与可解释性；从预训练模型 RoFormer 的结果来看，RoPE 具有良好的外推性，应用到 Transformer 中体现出较好的处理长文本的能力。此外，RoPE 还是目前唯一一种可用于线性 Attention 的相对位置编码。