RoFormer:使用旋转位置编码增强Transformer.

1. 相对位置编码

大部分相对位置编码通过修改Softmax前的Attention矩阵来引入相对位置信息,比如通过相对位置矩阵$B$引入:

\[\begin{aligned} \alpha_{ij} &= \text{softmax}\{x_iW^Q (W^K)^T x_j^T+B \} \\ z_i &= \sum_{j=1}^{n} \alpha_{ij}x_jW^V \end{aligned}\]

相对位置矩阵$B$是由查询向量$q_i$与键向量$k_j$共同决定的。作者的出发点是“通过绝对位置编码的方式实现相对位置编码”,因此向查询向量$q$与键向量$k$添加绝对位置信息:

\[\tilde{q}_m = f(q,m), \tilde{k}_n = f(k,n)\]

因此通过设计函数$f(\cdot,m),f(\cdot,n)$,经过上述操作后查询向量与键向量带有位置$m,n$的绝对位置信息。若希望通过注意力的内积运算后实现相对位置信息的引入,因此假设存在恒等关系:

\[<f(q,m), f(k,n)> = g(q,k,m-n)\]

下面试图求出上述恒等式的一个简洁的解。假设存在初始条件$f(q,0)=q,f(k,0)=k$。首先考虑二维情况,并假设查询向量$q$与键向量$k$为复数,此时二者的内积$<q,k>=\text{Re}[qk^*]$,则有:

\[\text{Re}[f(q,m)f^*(k,n)] = g(q,k,m-n)\]

将$f$和$g$表示为复指数形式:

\[\begin{aligned} f(q,m) &= R_f(q,m)e^{i\Theta_f(q,m)} \\ f(k,n) &= R_f(k,n)e^{i\Theta_f(k,n)} \\ g(q,k,m-n) &= R_g(q,k,m-n)e^{i\Theta_g(q,k,m-n)} \end{aligned}\]

代入恒等关系后得到方程组:

\[\begin{aligned} R_f(q,m)R_f(k,n) &= R_g(q,k,m-n) \\ \Theta_f(q,m)-\Theta_f(k,n) &= \Theta_g(q,k,m-n) \end{aligned}\]

对于第一个方程令$n=m$得:

\[R_f(q,m)R_f(k,m) = R_g(q,k,0) =R_f(q,0)R_f(k,0) = ||q|| \cdot ||k ||\]

不妨取$R_f(q,m)=||q||,R_f(k,m)=||k ||$。

对于第二个方程令$n=m$得:

\[\Theta_f(q,m)-\Theta_f(k,m) = \Theta_g(q,k,0) = \Theta_f(q,0)-\Theta_f(k,0) = \Theta(q)-\Theta(k)\]

注意到$\Theta_f(q,m)-\Theta(q)=\Theta_f(k,m) -\Theta(k)$,则$\Theta_f(q,m)-\Theta(q)$形式上应只与$m$有关而与$q$无关,记为$\phi(m)$,因此有:

\[\Theta_f(q,m)-\Theta(q)=\phi(m)\]

此时对于第二个方程令$n=m-1$得:

\[\phi(m) -\phi(m-1) = \Theta_g(q,k,1) +\Theta(k) -\Theta(q)\]

因此$\phi(m)$是等差数列。设上式右端为$\theta$,则解得$\phi(m)=m\theta$。

2. 旋转式位置编码(Rotary Position Embedding,RoPE)

根据上面的讨论,有$R_f(q,m)=||q||$,$\Theta_f(q,m)=\Theta(q)+m\theta$,则$f$表示为:

\[f(q,m) = R_f(q,m)e^{i\Theta_f(q,m)} = ||q||e^{i(\Theta(q)+m\theta)} = qe^{im\theta}\]

根据复数乘法的几何意义,该变换实际上对应着向量$q$的旋转,所以称之为“旋转式位置编码”(Rotary Position Embedding,RoPE),还可以写成矩阵形式:

\[f(q,m) = \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \end{pmatrix}\]

由于内积满足线性叠加性,因此任意偶数维向量$q$的RoPE,都可以表示为二维情形的拼接:

\[\begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \cdots & 0 & 0 \\ 0 & 0 & \sin m\theta_1 & \cos m\theta_1 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{\lfloor \frac{d}{2}-1 \rfloor} & -\sin m\theta_{\lfloor \frac{d}{2}-1 \rfloor} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{\lfloor \frac{d}{2}-1 \rfloor} & \cos m\theta_{\lfloor \frac{d}{2}-1 \rfloor} \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{pmatrix}\]

将上式左端矩阵记为旋转矩阵$\mathcal{R}$,则$\mathcal{R}$是一个正交矩阵,不会改变向量的模长,因此通常也不会改变原模型的稳定性。

旋转位置编码RoPE相当于给位置为$m$的向量$q$乘上矩阵$\mathcal{R}_m$、位置为$n$的向量$k$乘上矩阵$\mathcal{R}_n$,然后再计算注意力矩阵,此时注意力矩阵包含了相对位置信息:

\[(\mathcal{R}_mq)^T(\mathcal{R}_nk) = q^T\mathcal{R}_m^T\mathcal{R}_nk = q^T\mathcal{R}_{n-m}k\]

由于旋转矩阵是稀疏矩阵,因此直接用矩阵乘法实现RoPE会很浪费算力,在实践中通过逐位对应相乘$\otimes$来实现RoPE

\[\begin{pmatrix} q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{pmatrix} \otimes \begin{pmatrix} \cos m\theta_0 \\ \cos m\theta_0 \\ \cos m\theta_1 \\ \cos m\theta_1 \\ \vdots \\ \cos m\theta_{\lfloor \frac{d}{2}-1 \rfloor} \\ \cos m\theta_{\lfloor \frac{d}{2}-1 \rfloor} \end{pmatrix} + \begin{pmatrix} -q_1 \\ q_0 \\ -q_3 \\ q_2 \\ \vdots \\ -q_{d-1} \\ q_{d-2} \end{pmatrix} \otimes \begin{pmatrix} \sin m\theta_0 \\ \sin m\theta_0 \\ \sin m\theta_1 \\ \sin m\theta_1 \\ \vdots \\ \sin m\theta_{\lfloor \frac{d}{2}-1 \rfloor} \\ \sin m\theta_{\lfloor \frac{d}{2}-1 \rfloor} \end{pmatrix}\]

$\theta_i$选择三角函数编码$\theta_i=10000^{-2i/d}$。