上一代的问题:正弦位置编码是加法式 —— 将位置编码向量直接加到 token embedding 上(Q/K+PE),导致位置信息与语义信息混合后无法解耦。加法式编码不能自然表达相对位置依赖,且与线性注意力(linear attention)完全不兼容。
这代改了什么:不再对 embedding 加位置向量,而是对 Q/K 向量分别乘以旋转矩阵(乘法而非加法),旋转角度正比于位置索引 m。Q 和 K 的内积 ⟨f_q(x_m,m), f_k(x_n,n)⟩ 自动包含相对位置 (n-m) 信息,无需任何额外计算。
效果:RoPE 成为 2023 年后开源大模型的标准位置编码方案,被 Llama 1/2/3、Mistral、Qwen、DeepSeek、Gemma、Yi、Baichuan 等几乎所有主流模型采用。它也是目前唯一一种既支持相对位置表达又与线性注意力兼容的位置编码。
Transformer 的核心是自注意力机制(Vaswani et al., 2017),它对输入序列进行集合级(set-level)处理。Yun et al. (2020) 从理论上证明:标准的自注意力是 排列等变(permutation equivariant) 的 —— 交换任意两个位置的输入,输出也对应交换,完全忽略序列顺序。这意味着不加位置编码的 Transformer 无法区分 "the cat sat on the mat" 和 "the mat sat on the cat"。
早期方案将位置信息加(或拼接)到 token embedding 上,常见两种形式:
Shaw et al. (2018) 提出在注意力分数计算中引入相对位置偏置:
其中 B_{ij} 由相对距离 (i-j) 查表得到。Dai & Yang (2019) 进一步提出 Transformer-XL,将相对位置编码融入 QK 内积计算中。这些方案虽表达了相对位置,但每种模型都需自定义注意力计算逻辑,无法作为通用模块嵌入标准注意力。
Katharopoulos et al. (2020) 提出线性注意力(linear attention),通过核函数将 QK^T 替换为 φ(Q) φ(K)^T,使复杂度从 O(n²) 降到 O(n):
加法式位置编码(无论是绝对还是相对)与线性注意力不兼容,因为加法式的 p_m 破坏了核函数的线性性质。这意味着高效长序列模型中无法使用有效的位置编码,成为制约高效 Transformer 发展的重要瓶颈。
RoFormer 的洞察极为简洁优雅:不要将位置信息加到向量上,而是旋转向量。具体来说:
这个特性使 RoPE 成为 Transformer 位置编码的理想解:乘法式、相对位置感知、保持范数、可兼容线性注意力。相比 Shaw 和 Transformer-XL 的重写注意力计算,RoPE 只需在 Q/K 映射后加一个旋转层,与标准注意力完全解耦。
设位置 m 和 n 的 token 分别为 x_m, x_n,我们需要设计函数 f_q, f_k 使得:
即内积仅依赖于语义 x_m, x_n 和相对位置 (n-m),不包含绝对位置 m 和 n 的独立信息。
从最简单的 d=2 维度开始。将二维向量看作复数,其中实部为第一维,虚部为第二维:
其中 e^{i m θ} 是复数旋转因子(欧拉公式 e^{iθ}=cosθ+i sinθ),θ 是预设的旋转频率。这里复数乘法等价于二维空间中的旋转:
现在计算旋转后的 Q 和 K 的内积(在复数域中为实部):
这个结果仅依赖于语义 (W_q x_m, W_k x_n) 和相对位置 (m-n)θ —— 这正是我们想要的!
对于维度 d > 2 的一般情况,将 d 维空间划分为 d/2 个 2 维子空间,每个子空间独立旋转且频率不同:
其中旋转矩阵 R_{Θ,m} 是块对角矩阵:
每个 2x2 块为标准的旋转矩阵:
频率 θ_i 采用与正弦位置编码相同的几何级数:
这意味着:
RoPE 最优雅的性质体现在旋转矩阵的群性质上。由于旋转矩阵组满足可交换性和可加性:
这个等式的含义极其深刻:
RoPE 的一个重要性质是:当两个 token 的距离增大时,它们注意力分数的上界以反比形式衰减。
考虑两个固定位置的向量 q 和 k,它们的 RoPE 内积为:
论文附录通过 Abel 变换证明:
这意味着:
直接计算 R_{Θ,m} x 需要 d×d 矩阵乘法(O(d²)),但因为 R 是稀疏分块矩阵,可以利用其结构将计算简化为 O(d)。
具体实现技巧:将向量 reorder 后做逐元素乘加操作:
def apply_rope(x, cos, sin):
"""x: (batch, seq_len, n_heads, d)
cos, sin: (batch, seq_len, d) 预计算
"""
# 将向量分成相邻对 (x0,x1), (x2,x3), ...
# 对每对应用旋转
# 第一步:构造旋转后的对向量
# 每对 (x_{2i}, x_{2i+1}) -> (x_{2i}·cos_i - x_{2i+1}·sin_i,
# x_{2i}·sin_i + x_{2i+1}·cos_i)
# 等价写法:利用交错 reorder
x_rot = torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1).reshape(x.shape)
# 注意:上述简化为实际需逐维度
out = x * cos + x_rot * sin
return out
将这个计算展开来看,对于每个 2 维子空间:
实际操作中,cos(mθ_i) 和 sin(mθ_i) 可以预计算为查找表(LUT),推理时只需查表并做逐元素乘加,效率极高。
cos, sin = precompute_freqs_cis(dim, max_seq_len, theta=10000.0)x = x.reshape(*x.shape[:-1], -1, 2) # 分成对x_rot = torch.stack([-x[..., 1], x[..., 0]], dim=-1)out = x * cos.unsqueeze(-2) + x_rot * sin.unsqueeze(-2)
RoPE 相较加法式编码的另一个关键优势是:旋转矩阵可以与线性注意力的核函数交换顺序。
对于线性注意力,我们有:
如果核函数 φ 对旋转满足某种交换性(如 φ(Rx) ≈ R φ(x)),则旋转可以移到核函数外部,使相对位置编码与线性注意力同时生效。实际上,RoPE 的设计使得在初始化附近,旋转矩阵可以穿过核函数近似成立,这是此前加法式编码做不到的。
论文第 2.2 节给出了更一般的推导框架。令 f_q(x_m,m) 和 f_k(x_n,n) 作用在 d 维向量上,我们从复数域开始找到所有满足条件的函数形式:
其中 R_Θ(m) 和 R_Θ(n) 需要满足:
R_Θ(n-m) = R_Θ(n) · R_Θ(-m) 且 R_Θ(-m) = R_Θ^T(m)
这等价于要求 R_Θ 构成一个一维旋转群的表示(即 SO(2) 群)。由于 SO(2) 的所有有限维不可约表示都是二维的(或复一维的),上述分块对角形式是最优解。
θ_i = 10000^{-2(i-1)/d} 的选择与 Vaswani 正弦编码的频率完全相同。这意味着:
这种多频率设计与正弦编码共享相同的频谱直觉,但 RoPE 改用了乘法而非加法,根本性地改变了表达方式。
论文附录给出以下证明主线。设相对距离 t = n-m > 0,定义:
其中 a_i = q_{2i} k_{2i} + q_{2i+1} k_{2i+1}, b_i = q_{2i+1} k_{2i} - q_{2i} k_{2i+1}。
利用 Abel 变换(分部求和):
这是因为对于 t 足够大,cos(tθ_i) 和 sin(tθ_i) 的振荡频率足以抵消大部分能量,剩余能量与 t 成反比。这个界在专业文献中称为 RoPE 的远程衰减界(long-term decay bound)。
def precompute_freqs_cis(dim: int, max_seq_len: int, theta: float = 10000.0):
"""预计算所有位置的 cos/sin 值"""
# 频率 (d/2)
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[:dim//2].float() / dim))
# 位置向量
t = torch.arange(max_seq_len, dtype=torch.float)
# 外积: (max_seq_len, d/2)
angles = torch.outer(t, freqs)
# 对每个角度计算 cos 和 sin
cos = angles.cos() # (max_seq_len, d/2) -> repeat_interleave(2) -> (max_seq_len, d)
sin = angles.sin()
return cos, sin
def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
"""
x: (batch, seq_len, n_heads, dim)
cos, sin: (seq_len, dim)
"""
# 分对: (..., dim) -> (..., dim//2, 2)
x_reshape = x.float().reshape(*x.shape[:-1], -1, 2)
x0, x1 = x_reshape[..., 0], x_reshape[..., 1]
# 余弦部分
cos = cos.unsqueeze(-2) # (seq_len, 1, dim)
sin = sin.unsqueeze(-2)
cos_reshape = cos.reshape(*cos.shape[:-1], -1, 2)
sin_reshape = sin.reshape(*sin.shape[:-1], -1, 2)
# 旋转: (x0cos - x1sin, x0sin + x1cos)
out0 = x0 * cos_reshape[..., 0] - x1 * sin_reshape[..., 0]
out1 = x0 * sin_reshape[..., 0] + x1 * cos_reshape[..., 1] # 实际上 cos/sin 每个子空间一样
# 恢复形状
return torch.stack([out0, out1], dim=-1).reshape(*x.shape)
class RoFormerAttention(nn.Module):
def forward(self, x):
B, L, D = x.shape
q = self.wq(x).view(B, L, self.n_heads, self.head_dim)
k = self.wk(x).view(B, L, self.n_heads, self.head_dim)
v = self.wv(x).view(B, L, self.n_heads, self.head_dim)
# 应用 RoPE 到 Q 和 K
cos, sin = self.freqs[:L].to(x.device)
q = apply_rotary_emb(q, cos, sin)
k = apply_rotary_emb(k, cos, sin)
# 标准注意力计算(QK^T 中自动包含相对位置)
score = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
attn = F.softmax(score, dim=-1)
out = attn @ v
return out.reshape(B, L, D)
| 指标 | 正弦/绝对 PE | RoPE(旋转 PE) | 影响 |
|---|---|---|---|
| 计算方式 | 加法(加到 embedding 上) | 乘法(旋转 Q/K 向量) | 位置与语义解耦,不受 embedding scale 影响 |
| 相对位置表达 | 隐式(靠注意力自学) | 显式(内积自动编码 n-m) | 模型不需要大量数据学习位置模式 |
| 计算开销 | 每 token 一次向量加法 | 每 Q/K 一次向量乘加(cos+sin) | 可忽略(O(d) vs O(d)) |
| 线性注意力兼容 | 不兼容 | 兼容 | 开启了高效长序列架构的方向 |
| 长度外推 | 较差(特别是 learned PE) | 中等(可借助 YaRN/NTK 大幅改进) | 优于 learned PE,但非完美 |
| 归一化不变量 | 不保持(加位置后 norm 改变) | 保持(旋转不改变 L2 范数) | 注意力 softmax 输入分布更稳定 |
| 模型采用 | GPT-2, BERT, T5 | Llama 1/2/3, Mistral, Qwen, DeepSeek, Gemma, Yi, Baichuan | RoPE 为 2023+ 标配 |
在 RoPE 之前,每个新模型几乎都会设计自己的位置编码方案。RoPE 出来后迅速成为共识,原因有三:
RoPE 虽然优于 learned PE,但直接外推到训练长度之外时性能仍会下降。后续工作提出了两大改进方向:
Transformer (Vaswani et al., 2017) — 首次提出正弦位置编码,采用几何级数频率。虽然用加法实现,但其频率设计被 RoPE 原样继承。RoPE 也可以看作是对 Vaswani 位置编码思想的"乘法化"重现。
Self-Attention with Relative Position (Shaw et al., 2018) — 在注意力分数上添加可学习的相对位置偏置。RoPE 的出发点相同(表达相对位置),但手段完全不同(乘法旋转 vs 加法偏置)。
Transformer-XL (Dai et al., 2019) — 将相对位置融入 QK 内积,重写注意力公式。RoPE 以更优雅的方式达到了相同目标,且不需修改注意力计算本身。
T5's Relative Bias (Raffel et al., 2020) — 将距离分桶(bucketing),每个桶对应一个可学习标量。T5 方案简单有效,但需要预设桶边界,且每个距离只能对应一个标量值,表达能力受限。
YaRN (Yet another RoPE extensioN, Peng et al., 2023) — 通过 NTK-by-parts 频率调整,将 RoPE 上下文从 4K 扩展到 128K。解决了 RoPE 原生外推能力不足的问题,是目前最重要的 RoPE 扩展。
NoPE (No Position Encoding, Haviv et al., 2022) / Causal LLM without PE — 挑战了位置编码的必要性,证明因果掩码本身可以编码一定位置信息。但实验显示,无位置编码的性能在大规模训练下仍然不佳,RoPE 仍是推荐方案。
CoPE (Contextual Position Encoding, Golovneva et al., 2024) — 用上下文相关的位置编码替代固定频率。CoPE 的定位不是取代 RoPE,而是补充:在 RoPE 基础上增加上下文感知性。
| 方案 | 类型 | 复杂度 | 外推能力 | 线性注意力 | 采用情况 |
|---|---|---|---|---|---|
| Sinusoidal PE | 绝对, 加法 | O(1) | 中等 | 不兼容 | GPT-2, 早期工作 |
| Learned PE | 绝对, 加法 | O(1) | 差 | 不兼容 | BERT, GPT |
| Relative Bias (T5) | 相对, 加法 | O(n²) | 需调桶 | 不兼容 | T5, PaLM |
| RoPE | 相对, 乘法 | O(d) | 中等(+YaRN优秀) | 兼容 | Llama/Mistral/Qwen... 标准 |
| ALiBi | 相对, 线性偏置 | O(1) | 好 | 兼容 | Bloom, MPT |
| NoPE | 无 | O(1) | N/A | 兼容 | 少数实验 |
RoPE 最本质的贡献,在于将位置编码从加法运算变成了乘法运算。这看似只是数学表达式的简单变化,却带来了三个连锁效应:
回看技术发展的历史,一个技术方案成为标准通常不是因为它能做到其他方案做不到的事,而是因为它把复杂的东西变简单了。RoPE 就是典型:
当数十亿参数的模型一次又一次地因为位置编码方案的选择而面临训练不稳定问题时,"换掉位置编码试试" 这种试错成本极高。RoPE 的 "just works" 反而成了最大的优势 —— 工程师不需要理解旋转矩阵的群性质,只需要知道 "加上 RoPE 训练更稳"。
RoPE 使用 Vaswani 的 10000^{-2i/d} 频率谱,这绝非偶然。正弦编码花了大量工程经验选出的频率分布(对数间隔、覆盖从快到慢的频谱),被 RoPE 复用并赋予新意。在 RoPE 中,这些频率不再只是加法编码器的正弦波,而是实际的旋转角速度,物理意义更清晰。
目前,RoPE 是事实标准,但并非完美。其长度外推能力在原生状态下有限(需配合 YaRN 或 PI),且在某些极端长上下文中仍会退化。但它的核心设计 —— 乘法旋转、相对位置、线性注意力兼容 —— 已经证明了自身的可持续发展潜力。即使未来出现更好的方案,它很可能也会保留 RoPE 的乘法旋转结构,仅在频率设计上做改进(如 CoPE 的自适应频率)。
可以说,RoPE 是 Transformer 位置编码问题的"最小作用量原理"解:在最少改动原架构的前提下,实现了最多的功能收益。这也是它最终战胜所有竞争对手的根本原因。