旋转式位置编码(Rotary Position Embedding,RoPE),这是一种配合Attention机制能达到“绝对位置编码的方式实现相对位置编码”的设计。

建议先阅读《大模型基础之Sinusoidal位置编码》,RoPE的思路与之很类似,就是寻找\boldsymbol{f},满足

\begin{equation}\langle\boldsymbol{f}(\boldsymbol{q}, m), \boldsymbol{f}(\boldsymbol{k}, n)\rangle = g(\boldsymbol{q},\boldsymbol{k},m-n)\end{equation}

我们直接把《大模型基础之Sinusoidal位置编码》的\boldsymbol{p}_m拿过来,令\boldsymbol{f}(\boldsymbol{q}, m)=\boldsymbol{q} \cdot \boldsymbol{p}_m即可。

我们证明一下:

\begin{align} \boldsymbol{f}(\boldsymbol{q}, m) &=\boldsymbol{q} \cdot \boldsymbol{p}_m\\ &=(q_0+q_1i) \cdot (\cos m\theta + i\sin m\theta )\\ &=q_0\cos(m\theta) - q_1\sin(m\theta)+i(q_0\sin(m\theta) + q_1\cos(m\theta)) \end{align}

实部: q_0\cos(m\theta) - q_1\sin(m\theta)
虚部: q_0\sin(m\theta) + q_1\cos(m\theta)

这与二维旋转矩阵的应用相似。具体地说,给定一个二维向量 \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} ,我们可以使用以下的旋转矩阵来表示它关于原点旋转 m\theta 度后的位置:

R(m\theta) = \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{pmatrix}

将向量 \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} 与旋转矩阵相乘,我们得到:

R(m\theta) \cdot \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} = \begin{pmatrix} q_0\cos(m\theta) - q_1\sin(m\theta) \\ q_0\sin(m\theta) + q_1\cos(m\theta) \end{pmatrix}

因此,化简后的矩阵形式为:

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

我们现在讨论一个问题,两个向量旋转后的内积与它们旋转的角度差有关吗?

具体来说,如果两个向量 \mathbf{a} \mathbf{b} 在旋转前的内积是 \mathbf{a} \cdot \mathbf{b} ,并且 \mathbf{a} 旋转了 \alpha 角度,而 \mathbf{b} 旋转了 \beta 角度,那么旋转后的内积与原来的内积之间的关系可以通过以下方式来描述:

设向量 \mathbf{a} \mathbf{b} 的夹角为 \theta ,那么 \mathbf{a} \cdot \mathbf{b} = |\mathbf{a}| |\mathbf{b}| \cos(\theta)

\mathbf{a} 旋转 \alpha 角度后,与 \mathbf{b} 的夹角变为 \theta + \alpha 。同理,当 \mathbf{b} 旋转 \beta 角度后,与 \mathbf{a} 的夹角变为 \theta + \beta

因此,当两向量分别旋转后,它们之间的夹角变为 \theta + \alpha - \beta

旋转后的内积为:
\mathbf{a}' \cdot \mathbf{b}' = |\mathbf{a}| |\mathbf{b}| \cos(\theta + \alpha - \beta)

从上面的表达式可以看出,两向量旋转后的内积与它们旋转前的内积和旋转的角度差 \alpha - \beta 有关。

具体地,当 \alpha - \beta 为0时,即两向量旋转的角度相同,它们的内积不变。当 \alpha - \beta 增大或减小时,夹角 \theta 会相应地增大或减小,从而影响内积的值。

因此当\boldsymbol{f}(\boldsymbol{q}, m)=\boldsymbol{q} \cdot \boldsymbol{p}_m\langle\boldsymbol{f}(\boldsymbol{q}, m), \boldsymbol{f}(\boldsymbol{k}, n)\rangle = g(\boldsymbol{q},\boldsymbol{k},m-n)是成立的。

相应的高维度形式为:

\begin{equation}\scriptsize{\underbrace{\begin{bmatrix} \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{bmatrix}}_{\boldsymbol{\mathcal{R}}_m} \begin{bmatrix}q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1}\end{bmatrix}}\end{equation}

也就是说,给位置为m的向量\boldsymbol{q}乘上矩阵\boldsymbol{\mathcal{R}}_m、位置为n的向量\boldsymbol{k}乘上矩阵\boldsymbol{\mathcal{R}}_n,用变换后的\boldsymbol{Q}\boldsymbol{K}序列做Attention,那么Attention就自动包含相对位置信息了,因为成立恒等式:
\begin{equation}(\boldsymbol{\mathcal{R}}_m \boldsymbol{q})^{\top}(\boldsymbol{\mathcal{R}}_n \boldsymbol{k}) = \boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n \boldsymbol{k} = \boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{k}\end{equation}
值得指出的是,\boldsymbol{\mathcal{R}}_m是一个正交矩阵,它不会改变向量的模长,因此通常来说它不会改变原模型的稳定性。

由于\boldsymbol{\mathcal{R}}_m的稀疏性,所以直接用矩阵乘法来实现会很浪费算力,推荐通过下述方式来实现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}
其中\otimes是逐位对应相乘,即Numpy、Tensorflow等计算框架中的*运算。从这个实现也可以看到,RoPE可以视为是乘性位置编码的变体。