RoPE

RoPE的全称是旋转式位置编码（Rotary Position Embedding，RoPE），这是一种配合Attention机制能达到“绝对位置编码的方式实现相对位置编码”的设计，本文从数学上证明RoPE的原理；

https://kexue.fm/archives/8265

1. 基本思路

为了达到这个目的，我们假设通过下述运算来给q,k添加绝对位置信息：

$q_m=f(q,m),k_n=f(k,n)$

也就是说，我们分别为q,k设计操作 $f(\cdot ,m),f(\cdot ,n)$，使得经过该操作后， $q_m,k_n$就带有了位置m,n的绝对位置信息。Attention的核心运算是内积，所以我们希望的内积的结果带有相对位置信息，因此假设存在恒等关系： $⟨f(q,m),f(k,n)⟩=g(q,k,m-n)$ 所以我们要求出该恒等式的一个（尽可能简单的）解。求解过程还需要一些初始条件，显然我们可以合理地设 $f(q,0)=q和f(k,0)=k$。

2. 求解过程

在复数中有 $⟨\mathbf{q},\mathbf{k}⟩=\text{Re}[\mathbf{q}{\mathbf{k}}^{\ast }]$，$\text{Re}[]$代表复数的实部

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

然后我们用复数的指数形式，设：

$\begin{array}{rl}\mathbf{f}(\mathbf{q},m)=& R_f(\mathbf{q},m)e^{\text{i}{\Theta }_f(\mathbf{q},m)}\\ \mathbf{f}(\mathbf{k},n)=& R_f(\mathbf{k},n)e^{\text{i}{\Theta }_f(\mathbf{k},n)}\\ \mathbf{g}(\mathbf{q},\mathbf{k},m-n)=& R_g(\mathbf{q},\mathbf{k},m-n)e^{\text{i}{\Theta }_g(\mathbf{q},\mathbf{k},m-n)}\end{array}$

那么代入方程后就得到方程组：

$\begin{array}{rl}{R}_f(\mathbf{q},m)R_f(\mathbf{k},n)=& R_g(\mathbf{q},\mathbf{k},m-n)\\ {\Theta }_f(\mathbf{q},m)-{\Theta }_f(\mathbf{k},n)=& {\Theta }_g(\mathbf{q},\mathbf{k},m-n)\end{array}$

对于第一个方程，代入m=n得到：$R_f(\mathbf{q},m)R_f(\mathbf{k},m)=R_g(\mathbf{q},\mathbf{k},0)=R_f(\mathbf{q},0)R_f(\mathbf{k},0)=\Vert \mathbf{q}\Vert \Vert \mathbf{k}\Vert$最后一个等号源于初始条件$f(q,0)=q和f(k,0)=k$，所以现在我们可以很简单地设 $Rf(q,m)=\Vert q\Vert ,Rf(k,m)=\Vert k\Vert$，即它不依赖于m。

至于第二个方程，同样代入m=n得到

${\Theta }_f(\mathbf{q},m)-{\Theta }_f(\mathbf{k},m)={\Theta }_g(\mathbf{q},\mathbf{k},0)={\Theta }_f(\mathbf{q},0)-{\Theta }_f(\mathbf{k},0)=\Theta (\mathbf{q})-\Theta (\mathbf{k})$

这里的 $\Theta (q),\Theta (k)$是$q,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(\mathbf{q},\mathbf{k},1)+\Theta (\mathbf{k})-\Theta (\mathbf{q})$

即$\phi (m)$为等差数列，设右端为 $\theta$，那么就解得 $\phi (m)=m\theta$。

3. 编码形式

综上，我们得到二维情况下用复数表示的RoPE：

$\mathbf{f}(\mathbf{q},m)=R_f(\mathbf{q},m)e^{\text{i}{\Theta }_f(\mathbf{q},m)}=\Vert q\Vert e^{\text{i}(\Theta (\mathbf{q})+m\theta )}=\mathbf{q}{e}^{\text{i}m\theta }$

根据复数乘法的几何意义，该变换实际上对应着向量的旋转，所以我们称之为“旋转式位置编码”，它还可以写成矩阵形式 （欧拉公式： $e^{i\Theta }=\cos \Theta +i\sin \Theta$）：

$\mathbf{f}(\mathbf{q},m)=\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{c}{q}_0\\ q_1\end{array}\right)$

由于内积满足线性叠加性，因此任意偶数维的RoPE，我们都可以表示为二维情形的拼接，即：

\begin{equation}\scriptsize{\underbrace{\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 & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{d/2-1} & -\sin m\theta_{d/2-1} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{d/2-1} & \cos m\theta_{d/2-1} \\ \end{pmatrix}}_{\boldsymbol{\mathcal{R}}_m} \begin{pmatrix}q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1}\end{pmatrix}}\end{equation}

也就是说，给位置为m的向量q乘上矩阵Rm、位置为n的向量k乘上矩阵Rn，用变换后的Q,K序列做Attention，那么Attention就自动包含相对位置信息了，因为成立恒等式：

$({\mathcal{R}}_m\mathbf{q}{)}^{\top }({\mathcal{R}}_n\mathbf{k})={\mathbf{q}}^{\top }{\mathcal{R}}_m^{\top }{\mathcal{R}}_n\mathbf{k}={\mathbf{q}}^{\top }{\mathcal{R}}_{n-m}\mathbf{k}$

值得指出的是，$Rm$是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性。

由于 $Rm$的稀疏性，所以直接用矩阵乘法来实现会很浪费算力，推荐通过下述方式来实现RoPE：

\begin{equation}\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_{d/2-1} \\ \cos m\theta_{d/2-1} \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_{d/2-1} \\ \sin m\theta_{d/2-1} \end{pmatrix}\end{equation}

其中⊗是逐位对应相乘，即Numpy、Tensorflow等计算框架中的∗运算。从这个实现也可以看到，RoPE可以视为是乘性位置编码的变体。