2026/5/27 · LLM 理论精读 · 第 3 / 3 讲

B · RoPE 旋转位置编码

如何将绝对位置转化为旋转矩阵，使 Query 和 Key 的点积自然包含相对位置关系，兼具良好的长度外推性。

B · RoPE 旋转位置编码

绝对位置编码把位置信息加到向量上；RoPE 把位置信息编进 Q·K 的点积里，使得注意力分数天然包含相对位置关系，且不需要修改模型结构。

一、问题背景

标准 Self-Attention 是排列不变的（permutation invariant）——打乱输入顺序，Attention 结果只是行列对应变化，语义完全相同。必须额外注入位置信息。

目标：设计一种位置编码方式，使得：

q_m \cdot k_n = f(q, k, m - n)

即第 $m$ 位置的 Query 和第 $n$ 位置的 Key 的点积，只取决于它们的相对位置 $m - n$ ，而不是绝对位置。

这样的注意力天然具有相对位置感知，且对未见过的更长序列有更好的外推能力。

二、从复数旋转说起

2.1 二维情形

考虑二维向量，把它看作复数 $z = x + iy$ 。复数乘以 $e^{i\theta}$ 等价于把向量旋转 $\theta$ 角：

z \cdot e^{i\theta} = (x + iy)(\cos\theta + i\sin\theta)

= (x\cos\theta - y\sin\theta) + i(x\sin\theta + y\cos\theta)

写成矩阵形式：

\begin{pmatrix} x' \\ y' \end{pmatrix} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x \\ y \end{pmatrix}

这就是旋转矩阵 $R(\theta)$ 。

2.2 旋转的关键性质

两个旋转后向量的点积：

(R(\theta_m) q) \cdot (R(\theta_n) k) = q^T R(\theta_m)^T R(\theta_n) k = q^T R(\theta_n - \theta_m) k

利用了旋转矩阵是正交矩阵（ $R^T = R^{-1}$ ），且 $R(\alpha)^T R(\beta) = R(\beta - \alpha)$ 。

结论：点积结果只依赖于 $\theta_n - \theta_m$ ，即相对位置。

令 $\theta_m = m\theta$ （位置 $m$ 对应旋转角 $m\theta$ ），则第 $m$ 和第 $n$ 个 token 之间的注意力分数只由相对位置 $(m - n)\theta$ 决定。

三、扩展到高维：RoPE 的完整定义

实际模型中 Q/K 的维度是 $d$ （通常 64~128）。RoPE 把 $d$ 维向量分成 $d/2$ 组，每组 2 个维度独立做旋转，使用不同的频率 $\theta_i$ ：

\Theta = \left\{ \theta_i = 10000^{-2(i-1)/d},\ i = 1, 2, \ldots, \frac{d}{2} \right\}

这个频率设计与原始 Transformer 的正弦位置编码完全相同，保证不同维度捕捉不同频率的位置信息（低频维度捕捉远距离依赖，高频维度捕捉局部依赖）。

对位置 $m$ 处的 $d$ 维向量 $q$ ，RoPE 的变换为：

f(q, m) = \begin{pmatrix} q_1 \\ q_2 \\ q_3 \\ q_4 \\ \vdots \\ q_{d-1} \\ q_d \end{pmatrix} \otimes \begin{pmatrix} \cos m\theta_1 \\ \cos m\theta_1 \\ \cos m\theta_2 \\ \cos m\theta_2 \\ \vdots \\ \cos m\theta_{d/2} \\ \cos m\theta_{d/2} \end{pmatrix} + \begin{pmatrix} -q_2 \\ q_1 \\ -q_4 \\ q_3 \\ \vdots \\ -q_d \\ q_{d-1} \end{pmatrix} \otimes \begin{pmatrix} \sin m\theta_1 \\ \sin m\theta_1 \\ \sin m\theta_2 \\ \sin m\theta_2 \\ \vdots \\ \sin m\theta_{d/2} \\ \sin m\theta_{d/2} \end{pmatrix}

即：每对相邻维度 $(q_{2i-1}, q_{2i})$ 独立旋转角度 $m\theta_i$ 。

四、代码实现

import torch
import torch.nn as nn

def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):
    """预计算所有位置的旋转频率（复数形式）"""
    # 频率：每个维度对（共 dim/2 个）的旋转速度
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    # 位置序列
    t = torch.arange(seq_len)
    # 外积：每个位置 × 每个频率 → (seq_len, dim/2)
    freqs = torch.outer(t, freqs)
    # 转成复数形式 e^{iθ} = cos θ + i sin θ
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor):
    """对 Q 或 K 应用旋转位置编码"""
    # x: (batch, seq_len, n_heads, head_dim)
    # 把最后一维按对拆成复数
    x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    # 旋转：复数乘法等价于向量旋转
    x_out = torch.view_as_real(x_ * freqs_cis).flatten(3)
    return x_out.type_as(x)

# 使用
freqs_cis = precompute_freqs_cis(dim=64, seq_len=2048)
q_rotated = apply_rotary_emb(q, freqs_cis[start_pos : start_pos + seq_len])
k_rotated = apply_rotary_emb(k, freqs_cis[start_pos : start_pos + seq_len])


## 五、与其他位置编码的对比

| 方案 | 相对位置感知 | 长度外推 | 参数量 | 实现复杂度 |
| --- | --- | --- | --- | --- |
| 正弦编码（原始 Transformer） | ❌ 绝对位置 | 一般 | 0 | 低 |
| 可学习位置嵌入（GPT-2） | ❌ 绝对位置 | 差 | $L \times d$ | 低 |
| ALiBi | ✅ | 好 | 0 | 中 |
| **RoPE** | ✅ | 好 | 0 | 中 |

RoPE 的优势：无额外参数、天然相对位置、与 KV Cache 兼容（只对当前 token 的位置做旋转）。
---
## 六、长度外推：YaRN 等扩展
原始 RoPE 在训练长度之外外推效果有限（高频维度的旋转周期与训练分布不匹配）。
**YaRN（Yet Another RoPE extensioN）** 的思路：对不同频率的维度采用不同的缩放策略——高频维度直接用（局部关系不变），低频维度线性插值（拉伸位置范围），中频维度平滑过渡。LLaMA 3 等模型的长上下文版本均采用类似方案。