在去年的文章《Transformer升级之路:2、博采众长的旋转式位置编码》中,笔者提出了旋转位置编码(RoPE),当时的出发点只是觉得用绝对位置来实现相对位置是一件“很好玩的事情”,并没料到其实际效果还相当不错,并为大家所接受,不得不说这真是一个意外之喜。后来,在《Transformer升级之路:4、二维位置的旋转式位置编码》中,笔者讨论了二维形式的RoPE,并研究了用矩阵指数表示的RoPE的一般解。

既然有了一般解,那么自然就会引出一个问题:我们常用的RoPE,只是一个以二维旋转矩阵为基本单元的分块对角矩阵,如果换成一般解,理论上效果会不会更好呢?本文就来回答这个问题。

指数通解 #

《Transformer升级之路:4、二维位置的旋转式位置编码》中,我们将RoPE抽象地定义为任意满足下式的方阵
\begin{equation}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n=\boldsymbol{\mathcal{R}}_{n-m}\label{eq:re}\end{equation}
然后,我们探讨了如下矩阵指数形式的解
\begin{equation}\boldsymbol{\mathcal{R}}_n=\exp n\boldsymbol{B}\end{equation}
这里的矩阵指数,不是像Softmax那样的激活函数式的element-wise运算,而是按照泰勒级数定义的“Matrix Exponential”。根据“Baker–Campbell–Hausdorff公式”,我们有
\begin{equation}\begin{aligned}
\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n=&\,\big(\exp m\boldsymbol{B}\big)^{\top}\big(\exp n\boldsymbol{B}\big) = \big(\exp m\boldsymbol{B}^{\top}\big)\big(\exp n\boldsymbol{B}\big) \\
=&\,\exp\left(m\boldsymbol{B}^{\top} + n\boldsymbol{B} + \frac{1}{2}mn\left[\boldsymbol{B}^{\top}, \boldsymbol{B}\right]+\cdots\right)
\end{aligned}\end{equation}
这里$\left[\boldsymbol{A}, \boldsymbol{B}\right]=\boldsymbol{A}\boldsymbol{B}-\boldsymbol{B}\boldsymbol{A}$,$\cdots$省略的都是$m,n$的三次或三次以上的项。按照式$\eqref{eq:re}$,那么上式指数部分应该等于$(n-m)\boldsymbol{B}$,这就推出
\begin{equation}\boldsymbol{B}^{\top} = - \boldsymbol{B}\end{equation}
即要求$\boldsymbol{B}$是反对称矩阵。

正交通解 #

进一步地,我们有$(\exp \boldsymbol{B})^{\top}(\exp \boldsymbol{B})=\exp(\boldsymbol{B}-\boldsymbol{B}) = \boldsymbol{I}$和$\exp n\boldsymbol{B} = \left(\exp\boldsymbol{B}\right)^n$,前者说明$\exp\boldsymbol{B}$是正交矩阵,后者则启示我们这是不是可以推广到任意正交矩阵?不难验证,答案是肯定的,我们有结论:

对于任意正交矩阵$\boldsymbol{O}$,$\boldsymbol{\mathcal{R}}_n=\boldsymbol{O}^n$是满足式$\eqref{eq:re}$的解。

值得指出的是,在实数域内,并不是所有正交矩阵都可以写成$\exp\boldsymbol{B}$的形式,所以$\boldsymbol{O}^n$实际上是比矩阵指数形式更宽泛的解。从《恒等式 det(exp(A)) = exp(Tr(A)) 赏析》中我们知道$\det(\exp(\boldsymbol{A})) = \exp(\text{Tr}(\boldsymbol{A})) > 0$,所以能写成矩阵指数形式的正交矩阵行列式必然大于0(即等于1),事实上这个结果反过来也成立,即行列式等于1的正交矩阵,必然可以写成$\exp\boldsymbol{B}$的形式,其中$\boldsymbol{B}$是反对称矩阵。(参考《Why can any orthogonal matrix be written as O=e^A》)。

而对于$\det(\boldsymbol{O}) = -1$的正交矩阵,我们有$\boldsymbol{O}=\boldsymbol{O}_+ \boldsymbol{I}_-$,其中$\boldsymbol{I}_-$是对角线元素有一个-1、剩下都是1的对角阵,$\boldsymbol{O}_+$则是$\det(\boldsymbol{O}_+) = 1$的正交矩阵,它可以写成$\exp\boldsymbol{B}$的形式,此时$\boldsymbol{O}^n = (\boldsymbol{O}_+ \boldsymbol{I}_-)^n = \boldsymbol{I}_-^n\exp n\boldsymbol{B}$。这也就是说,即便对于$\det(\boldsymbol{O}) = -1$的$\boldsymbol{O}^n$,也只是$\exp n\boldsymbol{B}$的简单变换,所以接下来我们主要研究$\exp n\boldsymbol{B}$形式的解。

完备分析 #

众所周知,我们平时所用的RoPE位置编码,是如下形式的分块对角矩阵:
\begin{equation}\scriptsize{\left(\begin{array}{cc:cc:cc:cc}
\cos n\theta_0 & -\sin n\theta_0 & 0 & 0 & \cdots & \cdots & 0 & 0 \\
\sin n\theta_0 & \cos n\theta_0 & 0 & 0 & \cdots & \cdots & 0 & 0 \\
\hdashline
0 & 0 & \cos n\theta_1 & -\sin n\theta_1 & \cdots & \cdots & 0 & 0 \\
0 & 0 & \sin n\theta_1 & \cos n\theta_1 & \cdots & \cdots & 0 & 0 \\
\hdashline
\vdots & \vdots & \vdots & \vdots & \ddots & \ddots & \vdots & \vdots \\
\vdots & \vdots & \vdots & \vdots & \ddots & \ddots & \vdots & \vdots \\
\hdashline
0 & 0 & 0 & 0 & \cdots & \cdots & \cos n\theta_{d/2-1} & -\sin n\theta_{d/2-1} \\
0 & 0 & 0 & 0 & \cdots & \cdots & \sin n\theta_{d/2-1} & \cos n\theta_{d/2-1} \\
\end{array}\right)}\end{equation}
它可以简写成
\begin{equation}\begin{pmatrix}
\boldsymbol{R}_{n\theta_0} & \boldsymbol{0} & \cdots & \boldsymbol{0} \\
\boldsymbol{0} & \boldsymbol{R}_{n\theta_1} & \cdots & \boldsymbol{0} \\
\vdots & \vdots & \ddots & \vdots \\
\boldsymbol{0} & \boldsymbol{0} & \cdots & \boldsymbol{R}_{n\theta_{d/2-1}} \\
\end{pmatrix} = \exp n\begin{pmatrix}
\boldsymbol{J}_{\theta_0} & \boldsymbol{0} & \cdots & \boldsymbol{0} \\
\boldsymbol{0} & \boldsymbol{J}_{\theta_1} & \cdots & \boldsymbol{0} \\
\vdots & \vdots & \ddots & \vdots \\
\boldsymbol{0} & \boldsymbol{0} & \cdots & \boldsymbol{J}_{\theta_{d/2-1}} \\
\end{pmatrix}\end{equation}
其中
\begin{equation}\boldsymbol{R}_{\theta} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix},\quad \boldsymbol{J}_{\theta} = \begin{pmatrix} 0 & -\theta \\ \theta & 0\end{pmatrix}\end{equation}
这种选择可以说是最简单的一种,其本质原因可以说是为了降低计算量。那么,所谓完备性问题,就是要回答:如上的分块对角矩阵的特例,相比全参数$\exp n\boldsymbol{B}$,是否有能力上的缺失?换句话说,如果不考虑计算量,将$\boldsymbol{B}$替换为一般的反对称矩阵,效果是否可能会有提升?

回答这个问题不困难,事实上,对于任意偶数阶反对称矩阵,它都可以对角化为分块对角矩阵
\begin{equation}\boldsymbol{\Lambda} = \begin{pmatrix}
\boldsymbol{J}_{\theta_0} & \boldsymbol{0} & \cdots & \boldsymbol{0} \\
\boldsymbol{0} & \boldsymbol{J}_{\theta_1} & \cdots & \boldsymbol{0} \\
\vdots & \vdots & \ddots & \vdots \\
\boldsymbol{0} & \boldsymbol{0} & \cdots & \boldsymbol{J}_{\theta_{d/2-1}} \\
\end{pmatrix}\end{equation}
该结论可以参考Skew-symmetric matrix。也就是说,存在可逆矩阵$\boldsymbol{P}$,使得$\boldsymbol{B}=\boldsymbol{P}\boldsymbol{\Lambda}\boldsymbol{P}^{-1}$,于是
\begin{equation}\exp n\boldsymbol{B} = \exp \left(n\boldsymbol{P}\boldsymbol{\Lambda}\boldsymbol{P}^{-1}\right) = \boldsymbol{P}(\exp n\boldsymbol{\Lambda})\boldsymbol{P}^{-1}\end{equation}
也就是说,任意的$\exp n\boldsymbol{B}$与分块对角的$\exp n\boldsymbol{\Lambda}$,仅仅相差一个相似变换,而我们在Self Attention中应用RoPE时,是
\begin{equation}\boldsymbol{q}^{\top}\big(\exp (n-m)\boldsymbol{B}\big)\boldsymbol{k} = \big(\boldsymbol{P}^{\top}\boldsymbol{q}\big)^{\top}\big(\exp (n-m)\boldsymbol{\Lambda}\big)\big(\boldsymbol{P}^{-1}\boldsymbol{k}\big)\end{equation}
由于$\boldsymbol{q},\boldsymbol{k}$一般都是输入$\boldsymbol{x}$经过某个可学习的线性变换而来,$\boldsymbol{P}^{\top},\boldsymbol{P}^{-1}$原则上都可以吸收到线性变换的训练参数中,因此直接设为$\boldsymbol{q}^{\top}\big(\exp (n-m)\boldsymbol{\Lambda}\big)\boldsymbol{k}$理论上不会损失一般性。

所以,对于Self Attention来说,问题的答案是否定的。不过,如果是线性Attention,答案会有少许区别,因为线性Attention的$\boldsymbol{q},\boldsymbol{k}$加了个激活函数:
\begin{equation}\phi(\boldsymbol{q})^{\top}\big(\exp (n-m)\boldsymbol{B}\big)\varphi(\boldsymbol{k}) = \big(\boldsymbol{P}^{\top}\phi(\boldsymbol{q})\big)^{\top}\big(\exp (n-m)\boldsymbol{\Lambda}\big)\big(\boldsymbol{P}^{-1}\varphi(\boldsymbol{k})\big)\end{equation}
这就导致了$\boldsymbol{P}^{\top},\boldsymbol{P}^{-1}$不一定能吸收到线性变换的训练参数中,因此对线性Attention补上两个参数矩阵,是有可能带来提升的。

文章小结 #

本文简单分析了RoPE的完备性问题,表明对于Self Attention来说,目前的分块对角型RoPE不会损失一般性。

转载到请包括本文地址:https://kexue.fm/archives/9403

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Dec. 28, 2022). 《Transformer升级之路:6、旋转位置编码的完备性分析 》[Blog post]. Retrieved from https://kexue.fm/archives/9403

@online{kexuefm-9403,
        title={Transformer升级之路:6、旋转位置编码的完备性分析},
        author={苏剑林},
        year={2022},
        month={Dec},
        url={\url{https://kexue.fm/archives/9403}},
}