RoPE：旋转位置编码的实现

位置编码基本概念

在讲旋转矩阵位置编码之前让我们先深入理解一下为什么自注意力是“顺序失明”的。

自注意力的核心是计算一个句子中每个词与其他所有词的“注意力分数”或“相关性”。这个分数是通过查询向量（Query, Q）和键向量（Key, K）的点积来计算的。

想象一下这个句子：“猫追老鼠”

1. 模型为“猫”、“追”、“老鼠”生成了词嵌入向量。
2. 然后通过权重矩阵变换成Q, K, V向量。
3. 为了计算“猫”应该多关注“老鼠”，模型会计算 Q_猫 和 K_老鼠 的点积。

现在，我们把句子顺序打乱：“老鼠追猫”

1. 模型为这三个词生成的词嵌入、Q, K, V向量是完全一样的。
2. 为了计算“猫”应该多关注“老鼠”，模型依然是计算 Q_猫 和 K_老鼠 的点积。

你会发现，dot(Q_猫, K_老鼠) 的计算结果在两个句子中是完全相同的！ 自注意力机制无法从这个计算中得知“猫”是在“老鼠”之前还是之后。它只知道这两个词存在于同一个句子中。

这种特性被称为置换不变性（Permutation Invariance）。对于自注意力来说，[A, B, C] 和 [C, B, A] 是一回事。

但这在自然语言中是致命的，因为顺序决定了意义：

• “我打你” vs “你打我”
• “猫追老鼠” vs “老鼠追猫”

所以说我们必须引入位置编码，位置编码的核心目的，就是将序列中词元的绝对或相对位置信息注入到模型中，让模型能够理解单词的顺序关系。

在RoPE这一方法出现之前，主流的语言模型(BERT, GPT-1, GPT-2, GPT-3)均是采用的是可学习的绝对位置编码。其核心思想是创建一个位置编码矩阵 E，尺寸为 (max_sequence_length, d_model)。对于位置 i，其位置编码就是矩阵 E 的第 i 行。这个编码向量会与词元嵌入向量相加，然后输入到模型中，本质上就是一个Embbeding层。虽然按照这种方法模型可以灵活的根据具体任务和数据自行学习最优的位置表示，但其外推性极差，如果模型训练时的最大长度是 512，那么它完全不知道位置 513 的编码是什么。这是它在现代 LLM 中被逐渐取代的核心原因。

为什么需要外推性？假设一个模型的最长序列是500token，现在你告诉这个模型生成一个1000token长度的文章，如果模型没有外推性就无法完成这个任务。

随之改进的方法是一种在T5模型中被提到的方法：T5 相对位置偏置。其核心思想是不直接修改输入端的词元嵌入，而是在自注意力计算的最后一步，给注意力分数矩阵加上一个“偏置项”。这个偏置项只取决于查询（Query）和键（Key）之间的相对距离 i-j。

\[\mathrm{AttentionScore}(i, j) = \frac{Q_i K_j^\top}{\sqrt{d_k}} + \mathrm{Bias}_{i-j} \]

模型会学习一个小的查找表，存储不同相对距离的偏置值。为了节省参数，通常会使用“分桶”（Bucketing）技术，将较远的距离（如 > 128）视为同一个距离，共享一个偏置值。由于分桶机制，其外推能力仍然有限，但比可学习的绝对位置编码要好。

再然后就是旋转位置编码 (Rotary Position Embedding, RoPE)的提出，这是目前大语言模型领域最主流、最成功的位置编码方法之一。

算法解析

论文中提出为了能利用上 token 之间的相对位置信息，假设查询向量\(q_m\)和键向量\(k_m\)之间的内积操作可以被一个函数\(g\)表示，\(g\)的输入是词嵌入向量\(x_m\)和\(x_n\)以及它们的相对位置\(m-n\)

\[\langle \mathbf{f}_q(\mathbf{x}_m, m), \mathbf{f}_k(\mathbf{x}_n, n) \rangle = g(\mathbf{x}_m, \mathbf{x}_n, m-n) \]

假设现在词嵌入向量的维度是二维\(d=2\)，论文中提出了一个满足上述关系的\(f\)和\(g\)：

\[\begin{align*} f_q(\mathbf{x}_m, m) &= (\mathbf{W}_q \mathbf{x}_m) e^{im\theta} \\ f_k(\mathbf{x}_n, n) &= (\mathbf{W}_k \mathbf{x}_n) e^{in\theta} \\ g(\mathbf{x}_m, \mathbf{x}_n, m-n) &= \mathrm{Re}\left[ (\mathbf{W}_q \mathbf{x}_m) (\mathbf{W}_k \mathbf{x}_n)^* e^{i(m-n)\theta} \right] \end{align*} \]

这其中的\(Re\)表示复数中的实部。
通过欧拉公式的变换，\(f_q\)和\(f_k\)可以变成下面的式子

\[\begin{align*} f_q(\mathbf{x}_m, m) &= \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} W_q^{(1,1)} & W_q^{(1,2)} \\ W_q^{(2,1)} & W_q^{(2,2)} \end{pmatrix} \begin{pmatrix} x_m^{(1)} \\ x_m^{(2)} \end{pmatrix} \\ &= \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} q_m^{(1)} \\ q_m^{(2)} \end{pmatrix} \end{align*} \]

\[\begin{align*} f_k(\mathbf{x}_m, m) &= \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} W_k^{(1,1)} & W_k^{(1,2)} \\ W_k^{(2,1)} & W_k^{(2,2)} \end{pmatrix} \begin{pmatrix} x_m^{(1)} \\ x_m^{(2)} \end{pmatrix} \\ &= \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} k_m^{(1)} \\ k_m^{(2)} \end{pmatrix} \end{align*} \]

这相当于一个q或向量乘以一个旋转矩阵，这也就是为什么叫做旋转位置编码

最终\(g(\mathbf{x}_m, \mathbf{x}_n, m-n)\)可以表示如下

\[g(\mathbf{x}_m, \mathbf{x}_n, m-n) = \begin{pmatrix} \mathbf{q}_m^{(1)} & \mathbf{q}_m^{(2)} \end{pmatrix} \begin{pmatrix} \cos((m-n)\theta) & -\sin((m-n)\theta) \\ \sin((m-n)\theta) & \cos((m-n)\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \\ k_n^{(2)} \end{pmatrix} \]

从二维拓展到多维：

\[f_{\{q,k\}}(\mathbf{x}_m, m) = \mathbf{R}^d_{\Theta,m} \mathbf{W}_{\{q,k\}} \mathbf{x}_m \]

\[\mathbf{R}^d_{\Theta,m} = \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} \]

\[\Theta = \{ \theta_i = 10000^{-2(i-1)/d}, \; i \in [1, 2, \dots, d/2] \} \]

这里的超参数设置为10000可以带来良好的远程衰减性

![[Pasted image 20251031145438.png]]

代码实现

原文中提到由于\(\mathbf{R}^d_{\Theta,m}\)的稀疏性，直接用矩阵乘法来实现会很浪费算力，推荐通过下述方式来实现 RoPE：

\[\mathbf{R}^d_{\Theta,m} \mathbf{x} = \begin{pmatrix} x_0 \\ x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_{d-2} \\ x_{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} -x_1 \\ x_0 \\ -x_3 \\ x_2 \\ \vdots \\ -x_{d-1} \\ x_{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} \]

其中\(\otimes\)是逐位对应相乘，即计算框架中的\(*\)运算。
结合原文中的Figure 1可以发现，原文作者是把相邻的两个维度作为一组进行旋转变换，这样会带来一个问题：q,k矩阵都是存储在一个连续的内存上，如果照着原论文的处理方法(x[..., 0::2])，就需要产生非常多的切片，意味着计算时需要在内存中来回“跳跃”着读取数据，而不是顺序读取。这对CPU缓存和GPU内存读取来说是非常不友好的，会大大降低访存效率。
dot(RoPE(q_pair, m), RoPE(k_pair, n)) 的结果是一个只依赖于 (m-n) 和原始 q_pair, k_pair 的函数。无论取哪两个维度，最终的和都具有我们需要的那个核心性质：只依赖于相对位置 m-n。所以我们可以认为取任意两个维度的值进行旋转变换都是可行的。
所以我们决定只对q,k在中间进行一次切分，将前半部分的所有维度与后半部分的所有维度一一对应起来。\(d_0 与 d_{n/2}\) 配对，\(d_1 与 d_{n/2 + 1}\) 配对......

def apply_rope(x, cos, sin):assert x.ndim == 4n = x.shape[3] // 2x1 = x[..., n:]x2 = x[..., :n]y1 = x1 * cos + x2 * siny2 = x1 * -sin + x2 *cosout = torch.cat([y1, y2], 3)out = out.to(x.dtype)return out

apply_rope会在注意力层中被使用
接下来就是构建cos和sin矩阵

def _precompute_rope_(self, rope_len, head_dim, base=10000):channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32)theta = 1 / (base ** (channel_range / head_dim))m = torch.arange(rope_len, dtype=torch.float32)freqs = torch.outer(m, theta)cos, sin = freqs.cos(), freqs.sin()cos, sin = cos.bfloat16, sin.bfloat16cos, sin = cos[None, :, None, :], sin[None, :, None, :] return cos, sin

在这里我们对每个头单独进行RoPE，因为注意力分数的计算是在每个头内部发生的。因为 RoPE 的基本操作单位是“维度对”，而不是单个维度。每一对维度共享同一个旋转频率，所以我们只需要计算 head_dim / 2 个频率就足够了。
.outer()是外积