RoFormer：旋转位置编码增强Transformer

经典Transformer架构的位置编码属于绝对位置编码，即没有显式的包含相对位置信息。RoPE作者提出了旋转位置编码，利用旋转矩阵编码位置信息，且显式的包含相对位置信息。同时，对RoPE的推导和理论解释给出了详细的表述。

问题定义

为了包含相对位置信息，作者们定义了查询向量$\mathbf{q}_m$与键向量$\mathbf{k}_{n}$的内积函数$g$，可见式(1)

$$ \begin{aligned} \langle\mathbf{q}_m,\mathbf{k}_n\rangle=\langle f_q(x_m,m),f_{k}(x_n,n)\rangle=g(x_m,x_n,m-n) \end{aligned}\tag{1} $$

因此，其目标是寻找一个编码方式，从而使式(1)成立

旋转位置编码

2D

作者们先讨论了2维的情况，利用了2D平面的几何特性和复数形式，其函数为

$$ \begin{aligned} f_q(x_m,m)&=(\mathbf{W}_qx_m)e^{im\theta} \\ f_{k}(x_n,n)&=(\mathbf{W}_kx_n)e^{in\theta} \\ g(x_m,x_n,m-n)&=Re[(\mathbf{W}_qx_m)(\mathbf{W}_kx_n)^{*}e^{i(m-n)\theta}] \end{aligned}\tag{2} $$

式(2)中$Re[\cdot]$为复数的实数部分，$(\mathbf{W}_kx_n)^{*}$表示$\mathbf{W}_kx_n$的共轭复数，$\theta$为非零常量。进一步表达为

$$ \begin{aligned} f_{\{q,k\}}(x_m,x_n)=\begin{pmatrix} cosm\theta&-sinm\theta \\ sinm\theta&cosm\theta \end{pmatrix} \begin{pmatrix} W^{(11)}_{\{q,k\}}&W^{(12)}_{\{q,k\}}\\ W^{(21)}_{\{q,k\}}&W^{(22)}_{\{q,k\}} \end{pmatrix} \begin{pmatrix} x^{(1)}_m\\ x^{(2)}_m \end{pmatrix} \end{aligned}\tag{3} $$

对应的$g$为矩阵。位置信息编码体现在乘积$m\theta$

一般形式

为了把2D的情况泛化到一般形式，作者们把$d$维空间分解成$\frac{d}{2}$个子空间，利用内积线性的优点组合在一起，可见式(4)

$$ \begin{aligned} f_{\{q,k\}}(x_m,m)=\mathbf{R}^{d}_{\Theta,m}\mathbf{W}_{\{q,k\}}x_m \end{aligned}\tag{4} $$

式(4)中

$$ \begin{aligned} \mathbf{R}_{\Theta,m}^d=\begin{pmatrix} cosm\theta_1&-sinm\theta_1&0&0&\cdots&0&0\\ 0&0&cosm\theta_2&-sinm\theta_2&\cdots&0&0\\ \vdots&\vdots&\vdots&\vdots&\ddots&\vdots&\vdots\\ 0&0&0&0&\cdots&cosm\theta_{d/2}&-sinm\theta_{d/2}\\ 0&0&0&0&\cdots&sinm\theta_{d/2}&cosm\theta_{d/2} \end{pmatrix} \end{aligned}\tag{5} $$

式(5)为旋转矩阵，$\Theta=\{\theta_i=1000^{-2(i-1)/d},i\in[1,2,\ldots,d/2\}$，原理图可见图1所示。

图1 RoPE的原理图

那么，

$$ \begin{aligned} \mathbf{q}^{T}_m\mathbf{k}_n=(\mathbf{R}_{\Theta,m}^d\mathbf{W}_qx_m)^T(\mathbf{R}_{\Theta,n}^d\mathbf{W}_kx_n)=x^{T}\mathbf{W}_qR_{\Theta,n-m}^d\mathbf{W}_kx_n \end{aligned}\tag{6} $$

由于$\mathbf{R}_{\theta}^d$的稀疏性，直接的矩阵相乘效率偏低。因此，作者们建议

$$ \begin{aligned} \mathbf{R}^d_{\Theta,m}x= \begin{pmatrix} x_1\\ x_2\\ x_3\\ x_4\\ \vdots\\ x_{d-1}\\ x_d \end{pmatrix} \otimes \begin{pmatrix} cosm\theta_1\\ cosm\theta_1\\ cosm\theta_2\\ cosm\theta_2\\ \vdots\\ cosm\theta_{d/2}\\ cosm\theta_{d/2} \end{pmatrix}+ \begin{pmatrix} -x_2\\ x_1\\ -x_4\\ x_3\\ \vdots\\ -x_d\\ x_{d-1} \end{pmatrix} \otimes \begin{pmatrix} sinm\theta_1\\ sinm\theta_1\\ sinm\theta_2\\ sinm\theta_2\\ \vdots\\ sinm\theta_{d/2}\\ sinm\theta_{d/2} \end{pmatrix} \end{aligned}\tag{7} $$

最终，RoPE拥有灵活序列长度、随着位置距离增加的token间衰退的依赖、以及为线性自注意力配备相对位置编码的特性。