Transformer升级之路：6、旋转位置编码的完备性分析
原文链接 | 2022-12-28 | 苏剑林
在去年的文章《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不会损失一般性。