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

Transformer升级之路:二维位置的旋转式位置编码_第1张图片

©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 的一个解为:

Transformer升级之路:二维位置的旋转式位置编码_第2张图片

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

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

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

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

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

1.3 推导思路

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

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

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

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

四元数

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

2.1 复数与矩阵

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

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

Transformer升级之路:二维位置的旋转式位置编码_第3张图片

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

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

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

2.2 四元数简介

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

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

Transformer升级之路:二维位置的旋转式位置编码_第4张图片

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

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

这里 。此外,还有类似的矩阵映射:

Transformer升级之路:二维位置的旋转式位置编码_第5张图片

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,我们考虑:

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

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

Transformer升级之路:二维位置的旋转式位置编码_第6张图片

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

3.4 RoPE现身

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

Transformer升级之路:二维位置的旋转式位置编码_第7张图片

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

Transformer升级之路:二维位置的旋转式位置编码_第8张图片

不失一般性还可以设 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 个独立参数,不失一般性,考虑正交分解:

Transformer升级之路:二维位置的旋转式位置编码_第9张图片

解得:

求解代码:

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,此时:

Transformer升级之路:二维位置的旋转式位置编码_第10张图片

可以增加个参数 θ,完成展开,就得到:

Transformer升级之路:二维位置的旋转式位置编码_第11张图片

延伸故事

至此,关于二维 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 关注大脑探知、大脑功能和大脑健康。

更多阅读

Transformer升级之路:二维位置的旋转式位置编码_第12张图片

Transformer升级之路:二维位置的旋转式位置编码_第13张图片

Transformer升级之路:二维位置的旋转式位置编码_第14张图片

#投 稿 通 道#

 让你的论文被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。

总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学习心得技术干货。我们的目的只有一个,让知识真正流动起来。

???? 来稿标准:

• 稿件确系个人原创作品,来稿需注明作者个人信息(姓名+学校/工作单位+学历/职位+研究方向) 

• 如果文章并非首发,请在投稿时提醒并附上所有已发布链接 

• PaperWeekly 默认每篇文章都是首发,均会添加“原创”标志

???? 投稿邮箱:

• 投稿邮箱:[email protected] 

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

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

????

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

进入知乎首页搜索「PaperWeekly」

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

·

你可能感兴趣的:(人工智能,bmp,xhtml,cstring,3d)