0. 前言#

在正式推导之前, 我们先想清楚一件事: 我们到底希望位置编码做到什么?

回忆一下, Transformer 的 Self-Attention 计算的是 query 和 key 的内积:

问题是这个分数跟位置完全无关。把词序打乱, attention 输出一样。

所以我们需要给每个位置上的 query 和 key 加上位置信息, 让 attention 能感知到”谁在哪儿”。

但”加上位置信息”这个说法太笼统了。更精确地说, 我们希望:

编码位置后, 两个 token 的 attention score 只依赖于它们的相对位置差, 而不依赖于它们各自的绝对位置。

这是整篇文章推导的起点。搞清楚了这一点, 后面的所有公式就有了方向。

1. 设定目标#

把位置编码表示成两个函数:

• $f_q(\mathbf{q}, m)$: 对位置 $m$ 处的 query 向量 $\mathbf{q}$ 进行位置编码
• $f_k(\mathbf{k}, n)$: 对位置 $n$ 处的 key 向量 $\mathbf{k}$ 进行位置编码

我们期望:

即: 编码后的内积, 只跟词本身($\mathbf{q}$, $\mathbf{k}$)和相对位置($m-n$)有关, 跟绝对位置 $m$, $n$ 无关.

这就是 RoPE 的推导起点. 接下来我们要做的事是: 什么样的 $f_q$, $f_k$ 才能满足这个条件?

上图为 RoPE 的完整推导路线图. 从目标出发, 通过复数/三角函数, 最终推出旋转矩阵.

2. 从最简单的 2 维情况开始#

为了找到 $f_q$ 和 $f_k$ 的形式, 我们先考虑最简单的情况: $\mathbf{q}$ 和 $\mathbf{k}$ 是 2 维向量.

3. 验证: 旋转矩阵形式的内积#

有了旋转矩阵, 我们反过来验证一下内积.

在验证之前, 先确认旋转矩阵的两个重要性质:

性质 1: $R(\theta)$ 是正交矩阵, 且 $R(\theta)^T = R(-\theta)$

验算 $R(\theta)^T R(\theta)$:

同时 $R(-\theta) = \begin{pmatrix} \cos(-\theta) & -\sin(-\theta) \\ \sin(-\theta) & \cos(-\theta) \end{pmatrix} = \begin{pmatrix} \cos\theta & \sin\theta \\ -\sin\theta & \cos\theta \end{pmatrix} = R(\theta)^T$, 所以 $R(\theta)^T = R(-\theta) = R(\theta)^{-1}$.

性质 2: $R(\alpha)R(\beta) = R(\alpha+\beta)$ (旋转可加性)

把两个旋转矩阵相乘:

第二步用了三角函数的和角公式: $\cos(\alpha+\beta) = \cos\alpha\cos\beta - \sin\alpha\sin\beta$, $\sin(\alpha+\beta) = \sin\alpha\cos\beta + \cos\alpha\sin\beta$.

好, 有了这两条性质, 验证内积就很简单了. 不过先直观感受一下 RoPE 是怎么工作的:

左图: 原始向量 q 和 k. 中图: q 旋转 45°(编码位置 m=1), k 旋转 15°(编码位置 n=1). 右图: 同样两个向量, 但绝对位置更大(m=100, n=70), 因为相对差相同(都是 30°), 所以内积不变. 这就是”只依赖相对位置”的直观体现.

现在用数学验证:

展开 $R(-(m-n)\theta)$:

最后的结果只跟 $m-n$ 有关, 验证通过.

4. 为什么需要不同频率#

假设我们只有一个固定频率 $\theta$.

对于位置 $m$ 和 $m+1$, 旋转角度分别是 $m\theta$ 和 $(m+1)\theta$, 差值是 $\theta$. 但问题在于: 当 $m$ 很大时会发生什么?

来看一个具体例子. 假设 $\theta = 0.1$, 那么:

• 位置 $0$: 旋转 $0^\circ$, 向量不变
• 位置 $1$: 旋转 $5.7^\circ$
• …
• 位置 $30$: 旋转 $171.9^\circ$
• 位置 $31$: 旋转 $177.6^\circ$
• 位置 $63$: 旋转 $361.0^\circ$ — 转了一整圈多!

$R(63 \times 0.1) = R(6.3) \approx R(6.3-2\pi) = R(0.02)$, 所以 $m=63$ 和 $m=0$ 几乎有相同的旋转矩阵. 位置 $63$ 和位置 $0$ 在 attention 计算中无法区分!

也就是说, 由于 $\sin$ 和 $\cos$ 是周期函数, 旋转超过一周后, 位置信息就混叠了. 如果词序列很长, 后面的位置会周期性地”穿越”回前面.

所以我们需要多个频率, 让不同的维度对以不同的速度旋转:

• 低维度: 旋转速度快 ($\theta_i$ 大), 能区分精细位置
• 高维度: 旋转速度慢 ($\theta_i$ 小), 能感知大范围距离

频率设置和 Sinusoidal PE 一样:

这个公式的效果是: $i$ 越小, $\theta_i$ 越大(旋转越快), 反之亦然.

图中展示了三个不同频率的旋转速度. 红色转得最快(低维度), 蓝色次之, 绿色最慢(高维度).

5. 从 2D 扩展到高维#

既然 2D 情况下的位置编码是旋转, 那高维呢?

答案很自然: 把向量切成 $d/2$ 对, 每对独立旋转.

对于 8 维向量, 它的旋转矩阵是一个分块对角矩阵:

其中每个 $R(m\theta_i)$ 是 2×2 的旋转矩阵.

这个分块矩阵作用于 8 维向量时, 第 1-2 维用 $\theta_0$, 第 3-4 维用 $\theta_1$, 以此类推. 每对维度独立旋转, 互不干扰.

写成公式就是:

为什么可以这样做? 这里需要一点简单的推导。

假设我们把 $d$ 维向量 $\mathbf{q}$ 切分成 $d/2$ 个 2 维子向量 $\mathbf{q}^{(0)}, \mathbf{q}^{(1)}, ..., \mathbf{q}^{(d/2-1)}$, 其中 $\mathbf{q}^{(i)} = (q_{2i}, q_{2i+1})$. 那么内积可以写成:

这个式子成立, 仅仅是因为内积的定义就是对应位置相乘再求和, 我们可以自由地按任何顺序分组求和——交换律和结合律而已. 这个证明不需要任何额外的数学知识, 就是最基础的向量内积定义.

现在, 如果我们对每个子向量 $\mathbf{q}^{(i)}$ 独立施加旋转 $R(m\theta_i)$, 那么编码后的内积就是:

我们在 2D 情况下已经证明过, 每个子空间的内积都只依赖于 $m-n$:

所以整个内积也只依赖 $m-n$. 证毕.

这样我们就实现了 $d$ 维向量到 $d/2$ 个独立 2D 旋转的分解.

图中展示了 8 维向量被分成 4 对, 每对用不同的频率独立旋转, 整体构成块对角矩阵.

6. 高维验证小结#

上面的推导已经证明: 高维内积可以分解为：

$\langle f_q(\mathbf{q}, m), f_k(\mathbf{k}, n) \rangle = \sum_{i=0}^{d/2-1} \langle \mathbf{q}^{(i)}, R((n-m)\theta_i) \mathbf{k}^{(i)} \rangle$

其中 $\mathbf{q}^{(i)} = (q_{2i}, q_{2i+1})$ 是第 $i$ 个子空间的 2 维向量. 每个子空间的结果都只依赖于 $m-n$, 所以总和也只依赖 $m-n$.

这就是 RoPE 在高维下的完整形式.

7. 长距离衰减#

前面我们从目标出发推导了 RoPE 的形式. 现在来看 RoPE 自带的一个优雅性质: 长距离衰减.

假设 $\mathbf{q}$ 和 $\mathbf{k}$ 是来自同一分布的随机向量, 各分量均值为 0, 方差为 1, 且不同分量之间相互独立. 这意味着我们希望相关系数: 当 $j = l$ 时 $\mathbb{E}[q_j k_l] = 1$, 当 $j \neq l$ 时 $\mathbb{E}[q_j k_l] = 0$. 用 Kronecker delta 符号 $\delta_{jl}$ 统一表示就是 $\mathbb{E}[q_j k_l] = \delta_{jl}$ (即 $j=l$ 时为 1, 否则为 0).

$\delta_{jl}$ 是 Kronecker delta: $\delta_{jl}=1$ 当 $j=l$, 否则 $\delta_{jl}=0$.

现在计算 RoPE 编码后内积的期望:

展开每个子空间的内积:

展开后有四项含 $q_{2i}k_{2i}$, 四项含 $q_{2i+1}k_{2i+1}$, 以及交叉项 $q_{2i}k_{2i+1}$ 和 $q_{2i+1}k_{2i}$.

由于 $q$ 和 $k$ 的不同分量独立且均值为 0, 交叉项的期望为 0. 只有 $q_{2i}k_{2i}$ 和 $q_{2i+1}k_{2i+1}$ 的期望为 1:

其中用了三角恒等式 $\cos m\theta_i\cos n\theta_i + \sin m\theta_i\sin n\theta_i = \cos((m-n)\theta_i)$.

所以每个子空间贡献 $2\cos((m-n)\theta_i)$. 把所有 $d/2$ 个子空间加起来:

这个求和函数在 $m-n$ 增大时会呈现震荡衰减的趋势——距离越远, 预期的注意力分数越低. 这符合我们在自然语言中的直觉:相邻词通常比远距离的词关联更紧密, 而且”带震荡”的性质意味着某些特定距离的 token 也能获得较强注意力, 这与现实中周期性短语(如每第 n 个词)的匹配模式一致.

图中展示了不同维度下 RoPE 内积随距离的衰减趋势. 注意衰减不是单调的, 而是带有震荡的.

8. 30 行代码实现#

1
import torch
2

3
def precompute_rope_frequencies(dim: int, max_len: int, base: int = 10000):
4
    """预计算所有位置的 sin/cos 值"""
5
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
6
    positions = torch.arange(max_len).float()
7
    angles = positions[:, None] * inv_freq[None, :]         # (max_len, dim/2)
8
    angles = torch.cat([angles, angles], dim=-1)            # (max_len, dim)
9
    return angles.cos(), angles.sin()
10

11
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
12
    """对 x 施加 RoPE 旋转"""
13
    # x: (batch, seq_len, head, dim)
14
    # cos, sin: (seq_len, dim)
15
    cos = cos[None, :, None, :]    # (1, seq_len, 1, dim)
16
    sin = sin[None, :, None, :]    # (1, seq_len, 1, dim)
17
    # 每对 (x_{2i}, x_{2i+1}) 交换并取反 = 旋转
18
    x_rot = torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1).reshape(x.shape)
19
    return x * cos + x_rot * sin

使用:

1
cos, sin = precompute_rope_frequencies(dim=128, max_len=4096)
2
q_rotated = apply_rope(q, cos, sin)
3
k_rotated = apply_rope(k, cos, sin)
4
attn = torch.matmul(q_rotated, k_rotated.transpose(-2, -1))

9. 总结#

回头看整个推导过程, 最优雅的地方在于: 我们没有”发明”旋转矩阵, 而是从目标出发”发现”了它.

RoPE 现在已经是 LLaMA、Mistral、Gemma 等主流大模型的标配位置编码. 理解它的推导过程, 对理解后面长上下文扩展(PI、NTK、YaRN)也大有帮助.