Transformer升级之路：二维位置的旋转式位置编码
©PaperWeekly 原创 · 作者 | 苏剑林
单位 | 追一科技
研究方向 | NLP、神经网络
在之前的文章 Transformer 升级之路：博采众长的旋转式位置编码中我们提出了旋转式位置编码 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 团队。
文章小结
本文介绍了我们对 RoPE 的二维推广，主要以“相对性”、“可逆性”为出发点来确定二维 RoPE 的最终形式，尝试了四元数和矩阵指数两种推导过程，最终通过矩阵指数来给出了最终的解，从推导过程中我们还可以深化对 RoPE 的理解。
参考文献
[1] https://en.wikipedia.org/wiki/Quaternion
[2] https://kexue.fm/archives/8397
[3] https://kexue.fm/archives/898
[4] https://kexue.fm/archives/2291
[5] https://en.wikipedia.org/wiki/Matrix_exponential
[6] https://kexue.fm/archives/6377
[7] https://blog.eleuther.ai/rotary-embeddings/
[8] https://arxiv.org/abs/2104.09864
特别鸣谢
感谢 TCCI 天桥脑科学研究院对于 PaperWeekly 的支持。TCCI 关注大脑探知、大脑功能和大脑健康。
更多阅读
#投 稿 通 道#
让你的论文被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学习心得或技术干货。我们的目的只有一个，让知识真正流动起来。
???? 来稿标准：
• 稿件确系个人原创作品，来稿需注明作者个人信息（姓名+学校/工作单位+学历/职位+研究方向）
• 如果文章并非首发，请在投稿时提醒并附上所有已发布链接
• PaperWeekly 默认每篇文章都是首发，均会添加“原创”标志
???? 投稿邮箱：
• 投稿邮箱：[email protected]
• 所有文章配图，请单独在附件中发送
• 请留下即时联系方式（微信或手机），以便我们在编辑发布时和作者沟通
????
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
