Transformer升级之路：二维位置的旋转式位置编码

©PaperWeekly 原创 · 作者 |

苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

在之前的文章

中我们提出了旋转式位置编码 RoPE 以及对应的 Transformer 模型 RoFormer。由于笔者主要研究的领域还是 NLP，所以本来这个事情对于笔者来说已经完了。但是最近一段时间，Transformer 模型在视觉领域也大火，各种 Vision Transformer（ViT）层出不穷，于是就有了问题：二维情形的 RoPE 应该是怎样的呢？

咋看上去，这个似乎应该只是一维情形的简单推广，但真实情况却远比我们想象中复杂，本文就对此做一个分析，从中深化我们对 RoPE 的理解。

二维RoPE

什么是二维位置？对应的二维 RoPE 又是怎样的？它的难度在哪里？在这一节中，我们先简单介绍二维位置，然后直接给出二维 RoPE 的结果和推导思路，在随后的几节中，我们再详细给出推导过程。

1.1 二维位置

在 NLP 中，语言的位置信息是一维的，换句话说，我们需要告诉模型这个词是句子的第几个词；但是在 CV 中，图像的位置信息是二维的，即我们需要告诉模型这个特征是在第几行、第几列。这里的二维指的是完整描述位置信息需要两个数字，并不是指位置向量的维数。

有读者可能想：简单展平后当作一维的处理不行吗？确实不大行，比如一个的 feature map，位置展平后就变成了，而位置和展平后就分别变成了和，两者与的差分别是 h 和 1。然而，按照我们直观的认识，、它们与的距离应该是一样的才对，但是展平后却得到了不一样的 h 和 1，这未免就不合理了。

所以，我们需要专门为二维情形设计的位置编码，不能简单地展平为一维来做。

1.2 标准答案

经过后面的一番推导，得到二维 RoPE 的一个解为：

其中这个解很容易理解，它是两个一维 RoPE 组成的分块矩阵，实现上它就是将输入向量分为两半，一半施加 x 的一维 RoPE，一半施加 y 的一维 RoPE。由此形式我们也不难类比三维、四维等位置的 RoPE。

矩阵 (1) 是一个正交矩阵，它满足两个关键性质：

1. 相对性： 即，也正是由于这个性质，RoPE 才具有通过绝对位置实现相对位置的能力；

2. 可逆性： 给定可以反解出，这意味着对位置信息的编码是无损的；

某种意义上来说，式 (1) 是满足上述两个性质的最简单解，也就是说，虽然存在略有不同的解满足上述两个性质，但是它们形式上和实现上都相对复杂些。

1.3 推导思路

事后来看，RoPE 其实就是找到了矩阵，使得满足“相对性”条件：

所以，不难想到，二维 RoPE 的基本要求也是满足相对性，即要找到矩阵，使得它满足二维的相对性条件。不过，如果仅仅是这个要求的话，可行解就很多了，比如直接让：

但这个解的问题是，我们无法从逆向推出，这意味着这个选择对位置信息来说是有损的，所以我们需要多一个“可逆性”，保证可以从位置矩阵中无损地重构出原始位置信号。

对此，我们有两个比较自然的途径选择：1）四元数；2）矩阵指数。接下来我们将会逐一介绍它们。

四元数

在一维 RoPE 的推导中，我们主要以复数为工具，而四元数（Quaternion）[1] 是复数的推广，它同时也保留了复数的很多性质，所以用它来推导二维 RoPE 也算是一个自然的思路。不过很遗憾，这是一条走不通的途径，但笔者仍然将思考过程放置在此，供大家参考。

2.1 复数与矩阵

在高中时我们就学习过，复数跟二维向量是一一对应的（为了跟后面的四元数对齐，这里把虚数单位也加粗了），但这个对应关系只能保持加减法对应（因为向量没有通用的乘法运算）。更精妙的对应是将复数与矩阵对应：

在这个映射下，复数的加减乘除与矩阵的加减乘除一一对应，比如：

所以，矩阵映射才是复数域的完全同构，而向量映射只是直观的几何理解。

复数的矩阵映射也是 RoPE 的重要基础，在 Transformer 升级之路：博采众长的旋转式位置编码中我们已经推导得到 RoPE 的复数表示是，所以根据复数的矩阵映射，就对应着矩阵：

从而得到了一维 RoPE 的矩阵形式。

2.2 四元数简介

前面说了，四元数是复数是一种推广，事实上它还是矩阵的“鼻祖”，从历史上来说，是先有四元数然后才有一般的矩阵运算，并且四元数启发了矩阵的很多运算。早年笔者也写了《与向量的渊源极深的四元数》[3] 、《几何的数与数的几何：超复数的浅探究》[4] 等文章来介绍四元数，欢迎读者参考。

如果说复数是一个二维向量，那么四元数就是一个四维向量，表示为，这里的，但是它们本身都各不相等，几个基之间的运算规则是：

在当时它给人们最大的冲击是非交换性，比如。但除此之外，它跟复数运算其实是高度相似的。

比如类似复数的欧拉公式：

这里。此外，还有类似的矩阵映射：

2.3 违背相对性

关于这些公式背后的起源，那就说来话长了，这里也不打算细谈，有兴趣的读者请自行搜索资料阅读。有了欧拉公式和指数映射后，读者或许反应过来：一维 RoPE 无非是对应的矩阵映射，那么二维 RoPE 的话将映射为矩阵形式不就行了？

笔者一开始也是这样想的，很遗憾，这是错的。错在哪里呢？在一维 RoPE 中，我们利用了内积的复数表示：

该恒等式在四元数中同样成立，因此可以照搬。接着我们利用了复指数：

前两个等号都可以照搬到四元数中，关键是第三个等号在四元数中并不恒成立！一般地，对于两个四元数，等式并不成立！更广义来说，对于两个乘法不满足交换律的对象，一般都有。

所以，推导到最后，由于指数乘法无法转换为加法，最后的相对性没法得到保证。因此通过四元数推导这条路，就此夭折了...

矩阵指数

四元数的矩阵映射表明四元数事实上代表了一簇特定的 4×4 矩阵，四元数推导走不通，那么或许使用一般的矩阵分析可以走得通。事实上确实如此，在这一节中，我们将利用矩阵指数给出一个推导结果。

3.1 矩阵指数

这里的矩阵指数 [5] ，并不是神经网络的用指数函数作为激活函数的逐位运算，而是是按照幂级数定义的运算：

其中是指按照矩阵乘法将 k 个连乘。关于矩阵指数，笔者之前曾写过《恒等式 det(exp(A)) = exp(Tr(A)) 赏析》[6] ，也欢迎参考阅读。

矩阵指数是非常重要的一种矩阵运算，它可以直接写出常系数微分方程组的解：

当然这跟本文的主题关系不大。对于 RoPE 的推导，我们主要利用到矩阵指数的如下性质：

也就是说，如果的乘法可以交换，那么矩阵指数就可以像数的指数一样将乘法转换为加法。不过要注意这是一个充分不必要条件。

至于怎么把矩阵指数算出来，这里没法再展开介绍了，但是很多软件库已经自带了矩阵指数运算，比如数值计算库 scipy 和 tensorflow 都有 expm 函数，而符号计算的话，Mathematica 里边有 MatrixExp 函数。

3.2 一维通解

为什么能够将 RoPE 跟矩阵指数联系起来呢？因为一维的 RoPE 存在比较简单的指数表达式：

于是笔者开始思考如下形式的矩阵作为 RoPE 的解：

其中是一个跟 n 无关的矩阵，RoPE 的必要条件是满足“相对性”条件 (2)，于是我们分析：

这里先假设是可交换的，那么根据式（13）有：

要让，只需要满足：

这便是“相对性”给出的约束条件，刚才我们还假设了是可交换的，现在可以检验满足这个等式的一定是可交换的，所以结果是自洽的。

这也就是说，对于任何满足的矩阵都是方程（2）的解，并且还可以证明它一定是正交矩阵。对于的矩阵来说，它的的通解是，于是就有了如式（14）的解。

3.3 二维约束

类似地，对于二维 RoPE，我们考虑：

作为候选解。重复上述关于“相对性”条件的推导：先假设是可交换的，那么我们可以得到如下约束条件：

然而，与可交换，意味着、、和都可交换，但是上述两个约束只能保证和可交换，不能保证后两者的交换性，所以我们需要把它作为约束条件补充上去，得到：

不难证明在前两个条件下，新增的约束条件也相当于。

3.4 RoPE现身

由于满足前两个条件的 2×2 矩阵只有一个独立参数，不满足“可逆性”，所以我们至少要考虑 3×3 矩阵，它有 3 个独立参数：

为了保证可逆性，我们不妨设是“正交”的，比如设：

不失一般性还可以设 a=1，那么由条件解得 b=0,c=0，即只能是全零解，这不符合我们的要求。Mathematica 的求解代码为：

B[a_, b_, c_] = {{0, -a, -b}, {a, 0, -c}, {b, c, 0}};

B1 = B[
1, 
0, 
0];

B2 = B[
0, b, c];

Solve[{Dot[B1, B2] == Dot[B2, B1]}, {b, c}]

因此，我们至少要考虑 4×4 矩阵，它有 6 个独立参数，不失一般性，考虑正交分解：

解得：

求解代码：

B[a_, b_, c_, d_, e_, 

   f_] = {{0, -a, -b, -d}, {a, 0, -c, -e}, {b, c, 0, -f}, {d, e, f, 

    0}};

B1 = B[
1, b, c, 
0, 
0, 
0];

B2 = B[
0, 
0, 
0, d, e, f];

Solve[{Dot[B1, B2] == Dot[B2, B1]}, {b, c, d, e, f}]

可以求解结果没有对 f 提出约束，所以从最简单起见，我们可以让 f=1，剩下的 b,c,d,e 全部为 0，此时：

可以增加个参数 θ，完成展开，就得到：

延伸故事

至此，关于二维 RoPE 的推导介绍完毕。现在读者可能想问的是效果如何？很遗憾，现在还没有很完整的实验结果，毕竟笔者之前也没做过 ViT 相关的工作，而这个二维 RoPE 的推导也刚完成没多久，所以进展比较慢，只能说初步的结果显示还是挺有效的。EleutherAI 团队的成员也实验过这个方案，效果也比已有的其他位置编码好。

说到 EleutherAI 团队，这里再多说几句。EleutherAI 团队是前段时候比较火的号称要“复现 GPT3”的那个团队，我们在文章 Transformer 升级之路：博采众长的旋转式位置编码中提出了 RoPE 及 RoFormer 后，有幸得到了EleutherAI团队的关注，他们做了很多补充实验，确认了RoPE比很多其他位置编码都更加有效（参考他们的博客《Rotary Embeddings: A Relative Revolution》[7]），这促使我们完成了英文论文《RoFormer: Enhanced Transformer with Rotary Position Embedding》[8]并提交到了 Arxiv 上。而关于二维 RoPE 的疑问，最初也是来源于 EleutherAI 团队。