查看原文
其他

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

苏剑林 PaperWeekly
2024-08-23


©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[100];
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, 000];
B2 = B[000, 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 默认每篇文章都是首发,均会添加“原创”标志


📬 投稿邮箱:

• 投稿邮箱:hr@paperweekly.site 

• 所有文章配图,请单独在附件中发送 

• 请留下即时联系方式(微信或手机),以便我们在编辑发布时和作者沟通



🔍


现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧



·

继续滑动看下一个
PaperWeekly
向上滑动看下一个

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存