14 · RoFormer: Enhanced Transformer with Rotary Position Embedding (RoPE)

位置编码 旋转矩阵 多头注意力 Su et al. · Zhuiyi Technology · 2021
arXiv:2104.09864

上一代的问题:正弦位置编码是加法式 —— 将位置编码向量直接加到 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 等几乎所有主流模型采用。它也是目前唯一一种既支持相对位置表达又与线性注意力兼容的位置编码。

一、核心背景

1.1 Self-Attention 的位置无关性

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"。

1.2 绝对位置编码(Absolute PE)

早期方案将位置信息加(或拼接)到 token embedding 上,常见两种形式:

f(x_m, m) = W (x_m + p_m)   或   f(x_m, m) = W x_m + p_m
加法式绝对位置编码的一般形式:位置向量 p_m 直接加到 token 上

1.3 相对位置编码(Relative PE)

Shaw et al. (2018) 提出在注意力分数计算中引入相对位置偏置:

Attention(Q,K,V) = softmax( (QK^T)/sqrt(d) + B ) V

其中 B_{ij} 由相对距离 (i-j) 查表得到。Dai & Yang (2019) 进一步提出 Transformer-XL,将相对位置编码融入 QK 内积计算中。这些方案虽表达了相对位置,但每种模型都需自定义注意力计算逻辑,无法作为通用模块嵌入标准注意力。

1.4 与线性注意力的不兼容

Katharopoulos et al. (2020) 提出线性注意力(linear attention),通过核函数将 QK^T 替换为 φ(Q) φ(K)^T,使复杂度从 O(n²) 降到 O(n):

V'_i = ( φ(Q_i) Σ_j φ(K_j)^T V_j ) / ( φ(Q_i) Σ_j φ(K_j) )

加法式位置编码(无论是绝对还是相对)与线性注意力不兼容,因为加法式的 p_m 破坏了核函数的线性性质。这意味着高效长序列模型中无法使用有效的位置编码,成为制约高效 Transformer 发展的重要瓶颈。

二、核心思想

RoFormer 的洞察极为简洁优雅:不要将位置信息加到向量上,而是旋转向量。具体来说:

⟨ f_q(x_m, m), f_k(x_n, n) ⟩ = g(x_m, x_n, m - n)
核心性质:经过旋转后的 Q/K 内积只依赖于相对位置 (n-m),与绝对位置无关

这个特性使 RoPE 成为 Transformer 位置编码的理想解:乘法式、相对位置感知、保持范数、可兼容线性注意力。相比 Shaw 和 Transformer-XL 的重写注意力计算,RoPE 只需在 Q/K 映射后加一个旋转层,与标准注意力完全解耦。

三、公式推导

3.1 问题形式化

设位置 m 和 n 的 token 分别为 x_m, x_n,我们需要设计函数 f_q, f_k 使得:

⟨ f_q(x_m, m), f_k(x_n, n) ⟩ = g(x_m, x_n, m - n)

即内积仅依赖于语义 x_m, x_n 和相对位置 (n-m),不包含绝对位置 m 和 n 的独立信息。

3.2 二维情形的复数推导

从最简单的 d=2 维度开始。将二维向量看作复数,其中实部为第一维,虚部为第二维:

f_q(x_m, m) = (W_q x_m) · e^{i m θ}
f_k(x_n, n) = (W_k x_n) · e^{i n θ}

其中 e^{i m θ} 是复数旋转因子(欧拉公式 e^{iθ}=cosθ+i sinθ),θ 是预设的旋转频率。这里复数乘法等价于二维空间中的旋转:

q_m e^{i m θ} = (q_m cos(mθ) - q_{m+1} sin(mθ)) + i (q_m sin(mθ) + q_{m+1} cos(mθ))

现在计算旋转后的 Q 和 K 的内积(在复数域中为实部):

⟨ f_q(W_q x_m,m), f_k(W_k x_n,n) ⟩
= Re[ (W_q x_m)(W_k x_n)^* · e^{i(m-n)θ} ]

这个结果仅依赖于语义 (W_q x_m, W_k x_n) 和相对位置 (m-n)θ —— 这正是我们想要的!

3.3 一般形式:逐维旋转

对于维度 d > 2 的一般情况,将 d 维空间划分为 d/2 个 2 维子空间,每个子空间独立旋转且频率不同:

f_{q,k}(x_m, m) = R_{Θ, m} · W_{q,k} · x_m
RoPE 的一般形式:先做线性变换 W,再乘以旋转矩阵 R

其中旋转矩阵 R_{Θ,m} 是块对角矩阵:

R_{Θ,m} = diag( R_1(m), R_2(m), ..., R_{d/2}(m) )

每个 2x2 块为标准的旋转矩阵:

R_i(m) = [cos(mθ_i)   -sin(mθ_i)]
          [sin(mθ_i)   cos(mθ_i)]

频率 θ_i 采用与正弦位置编码相同的几何级数:

θ_i = 10000^{-2(i-1)/d}   其中 i=1,2,...,d/2

这意味着:

3.4 核心性质:相对位置

RoPE 最优雅的性质体现在旋转矩阵的群性质上。由于旋转矩阵组满足可交换性和可加性:

R_m^T R_n = R_{-m} R_n = R_{n-m}
(R_m q)^T (R_n k) = q^T R_m^T R_n k = q^T R_{n-m} k
核心等式:旋转后 Q/K 的内积只依赖于相对位置 (n-m)

这个等式的含义极其深刻:

3.5 长期衰减性质(Long-term Decay)

RoPE 的一个重要性质是:当两个 token 的距离增大时,它们注意力分数的上界以反比形式衰减

考虑两个固定位置的向量 q 和 k,它们的 RoPE 内积为:

S(m-n) = (R_m q)^T (R_n k) = Σ_{i=1}^{d/2} q_i^T R_{n-m}^{(i)} k_i

论文附录通过 Abel 变换证明:

|S(t)| ≤ C · (d/2) / t   对于 t = n-m 足够大
长期衰减不等式:内积绝对值被 O(1/t) 上界约束

这意味着:

3.6 高效数值实现

直接计算 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 维子空间:

R_i(m) · [x_{2i}, x_{2i+1}]^T =
[x_{2i} cos(mθ_i) - x_{2i+1} sin(mθ_i),
 x_{2i} sin(mθ_i) + x_{2i+1} cos(mθ_i)]^T

实际操作中,cos(mθ_i) 和 sin(mθ_i) 可以预计算为查找表(LUT),推理时只需查表并做逐元素乘加,效率极高。

实际代码示例(PyTorch):
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)

3.7 与线性注意力的兼容性

RoPE 相较加法式编码的另一个关键优势是:旋转矩阵可以与线性注意力的核函数交换顺序

对于线性注意力,我们有:

V' = φ(Q R) (φ(K R)^T V) / (φ(Q R) Σ φ(K R))

如果核函数 φ 对旋转满足某种交换性(如 φ(Rx) ≈ R φ(x)),则旋转可以移到核函数外部,使相对位置编码与线性注意力同时生效。实际上,RoPE 的设计使得在初始化附近,旋转矩阵可以穿过核函数近似成立,这是此前加法式编码做不到的。

论文核心论点:Rotational position encoding preserves the norm and the linearity of the original vectors under non-negative activation, making RoPE uniquely compatible with linear attention.

四、理论推导:从复数域到长期衰减界

4.1 最一般形式的推导

论文第 2.2 节给出了更一般的推导框架。令 f_q(x_m,m) 和 f_k(x_n,n) 作用在 d 维向量上,我们从复数域开始找到所有满足条件的函数形式:

f_q(x_m, m) = R_Θ(m) · W_q · x_m
f_k(x_n, n) = R_Θ(n) · W_k · x_n

其中 R_Θ(m) 和 R_Θ(n) 需要满足:
R_Θ(n-m) = R_Θ(n) · R_Θ(-m) 且 R_Θ(-m) = R_Θ^T(m)

这等价于要求 R_Θ 构成一个一维旋转群的表示(即 SO(2) 群)。由于 SO(2) 的所有有限维不可约表示都是二维的(或复一维的),上述分块对角形式是最优解。

4.2 频率选择的理论依据

θ_i = 10000^{-2(i-1)/d} 的选择与 Vaswani 正弦编码的频率完全相同。这意味着:

这种多频率设计与正弦编码共享相同的频谱直觉,但 RoPE 改用了乘法而非加法,根本性地改变了表达方式。

4.3 长期衰减上界的完整证明

论文附录给出以下证明主线。设相对距离 t = n-m > 0,定义:

S(t) = Σ_{i=1}^{d/2} s_i(t),   s_i(t) = a_i · cos(tθ_i) + b_i · sin(tθ_i)

其中 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 变换(分部求和):

|S(t)| = | Σ_{i=1}^{d/2} s_i(t) | ≤ C · Σ_{i=1}^{d/2} 1/tθ_i ≤ C · (d/2) / t

这是因为对于 t 足够大,cos(tθ_i) 和 sin(tθ_i) 的振荡频率足以抵消大部分能量,剩余能量与 t 成反比。这个界在专业文献中称为 RoPE 的远程衰减界(long-term decay bound)

五、算法与实现

5.1 预计算频率

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

5.2 RoPE 前向

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)
实现要点:
1. cos/sin 只需预计算一次,长度可外推(RoPE 本身不限制推理长度)
2. 旋转发生在 Q/K 的头上(每个 head 独立旋转),频率对所有 head 共享
3. 实际实现中,cos 和 sin 的形状为 (max_seq_len, dim),前向时按需截断
4. Llama 等模型在 Q 和 K 上独立应用 RoPE,V 不做旋转

5.3 在 RoFormer 中的整体使用

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+ 标配

6.1 为什么 RoPE 胜出了

在 RoPE 之前,每个新模型几乎都会设计自己的位置编码方案。RoPE 出来后迅速成为共识,原因有三:

6.2 RoPE 的长度外推问题与改进

RoPE 虽然优于 learned PE,但直接外推到训练长度之外时性能仍会下降。后续工作提出了两大改进方向:

后续发展:YaRN (2023) 证明 RoPE 的频率谱可以理解为一种神经正切核(NTK),不同的旋转频率段对位置编码的贡献不同,"by-parts" 策略为每一段设计最优的缩放因子。这使 RoPE + YaRN 成为目前最有效的免训练上下文扩展方案。

七、上下游关联

上一篇:位置编码的前期工作

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 方案简单有效,但需要预设桶边界,且每个距离只能对应一个标量值,表达能力受限。

下一篇:RoPE 的演进

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 基础上增加上下文感知性。

采用 RoPE 的著名模型

Llama 系列: Meta 的 Llama 1 (2023)、Llama 2 (2023)、Llama 3 (2024) 全部使用 RoPE。Llama 的广泛采用是 RoPE 成为标准的关键推动力。

Mistral 7B (2023): 使用 RoPE + sliding window attention。

Qwen 系列 (2023-2024): 阿里 Qwen 和 Qwen2 全部使用 RoPE,在长上下文(128K)场景下表现优异。

DeepSeek (2024): DeepSeek V2/V3 使用 RoPE 并进一步改进了频率设计。

其他: Gemma (Google, 2024)、Yi (零一万物)、Baichuan (百川智能)、Phi (Microsoft) 等几乎所有 2023 后发布的开源 LLM 都采用 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 兼容 少数实验

八、个人思考

8.1 乘法 vs 加法:范式转换的威力

RoPE 最本质的贡献,在于将位置编码从加法运算变成了乘法运算。这看似只是数学表达式的简单变化,却带来了三个连锁效应:

8.2 为什么 RoPE 能成为事实标准?

回看技术发展的历史,一个技术方案成为标准通常不是因为它能做到其他方案做不到的事,而是因为它把复杂的东西变简单了。RoPE 就是典型:

当数十亿参数的模型一次又一次地因为位置编码方案的选择而面临训练不稳定问题时,"换掉位置编码试试" 这种试错成本极高。RoPE 的 "just works" 反而成了最大的优势 —— 工程师不需要理解旋转矩阵的群性质,只需要知道 "加上 RoPE 训练更稳"。

8.3 频率设计的传承与创新

RoPE 使用 Vaswani 的 10000^{-2i/d} 频率谱,这绝非偶然。正弦编码花了大量工程经验选出的频率分布(对数间隔、覆盖从快到慢的频谱),被 RoPE 复用并赋予新意。在 RoPE 中,这些频率不再只是加法编码器的正弦波,而是实际的旋转角速度,物理意义更清晰。

8.4 位置编码的终局?

目前,RoPE 是事实标准,但并非完美。其长度外推能力在原生状态下有限(需配合 YaRN 或 PI),且在某些极端长上下文中仍会退化。但它的核心设计 —— 乘法旋转、相对位置、线性注意力兼容 —— 已经证明了自身的可持续发展潜力。即使未来出现更好的方案,它很可能也会保留 RoPE 的乘法旋转结构,仅在频率设计上做改进(如 CoPE 的自适应频率)。

可以说,RoPE 是 Transformer 位置编码问题的"最小作用量原理"解:在最少改动原架构的前提下,实现了最多的功能收益。这也是它最终战胜所有竞争对手的根本原因。