RoFormer

Enhanced Transformer with Rotary Position Embedding
ICLR / arXiv 2021-2023 Jianlin Su (追一科技) 等 NLP 位置编码机制 LLM 奠基之作

一分钟摘要

Transformer 模型本身不具备对序列位置信息的感知能力。为此,学术界提出了各类位置编码方案(如绝对位置编码和相对位置编码)。这篇论文提出了一种革命性的方法——旋转位置编码(Rotary Position Embedding, RoPE)。它通过将绝对位置编码转换为复数域上的旋转矩阵,并在计算自注意力时对向量进行相乘,巧妙地在自注意力机制中隐式融入了精确的相对位置信息。该方法具备高度的序列长度外推灵活性、符合直觉的相对距离衰减特性,并且完美兼容线性注意力机制。

研究动机与问题背景

在自然语言理解中,单词的前后序列顺序对语义建模至关重要。自注意力机制(Self-Attention)在捕获全局上下文方面极其强大,但它的架构天然是“位置无关”的。为了打破这种对称性,我们需要将位置信息注入模型。

现有方法的局限性

在 RoPE 提出之前,主流方案主要分为两类,但它们都存在明显妥协:

  • 绝对位置编码 (Absolute PE): 如原始 Transformer 使用的正弦波编码或 BERT 的可学习绝对编码。这些方法将位置信息与词向量相加。缺陷在于,在进行点积时,相对距离信息被掩盖,模型难以直接且一致地利用词元间的“距离感”。
  • 相对位置编码 (Relative PE): 如 T5、Transformer-XL 中使用的方法,通常是在注意力矩阵的点积结果中额外加上一个距离偏置项(Additive Position Encoding)。缺陷在于,由于这种加法分解打破了点积的纯粹性,它不仅计算变慢、显存占用增加,更致命的是,它无法适配线性自注意力架构

本文试图解决的核心问题

如何设计一种位置编码机制,既能在底层保留每个 token 的绝对位置,又能在注意力运算时自动显现出纯粹的相对位置关系,并且不改变点积乘法的结构?
RoPE 的解答:不相加,只旋转。

数学表示与算法建模

RoPE 的推导极其优美,核心思想在于使用复平面旋转来表达位置的变化。

核心符号定义

符号含义描述
$x_m, x_n$序列中第 $m$ 个和第 $n$ 个位置的词嵌入向量,维度为 $d$。
$q_m, k_n$融入了位置信息后的 Query 和 Key 向量。
$\langle q_m, k_n \rangle$Query 和 Key 在注意力机制中的内积运算。
$\theta$旋转的基准角度超参数(多维时有 $\theta_i$)。

形式化问题的提出

论文希望寻找一种特殊的函数 $f_q$ 和 $f_k$,将位置信息 $m, n$ 编码到向量后,两者的点积结果 $g$ 仅依赖于输入向量 $x_m, x_n$ 以及它们的相对距离 $m-n$:

$$ \langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, m-n) $$

2D 空间的优美推导

假设隐层维度 $d=2$。我们利用二维平面向量的复数几何特性。如果将向量表示为复数,要让两个复数的内积只取决于它们的角度差,最好的方法就是给它们分别乘以一个与绝对位置成正比的旋转因子 $e^{im\theta}$。推导后论文给出了以下解:

$$ f_q(x_m, m) = (W_q x_m) e^{im\theta}, \quad f_k(x_n, n) = (W_k x_n) e^{in\theta} $$

在这个解下,Query 和 Key 的内积天然满足相对位置关系:

$$ g(x_m, x_n, m-n) = \text{Re} \left[ (W_q x_m)(W_k x_n)^* e^{i(m-n)\theta} \right] $$

高维空间的推广与矩阵表示 (RoPE)

为了将 2D 推广到 $d$ 维,论文将 $d$ 维空间两两划分为 $d/2$ 个相互独立的 2D 子空间。对每个子空间施加不同旋转频率的旋转。这就等价于对原本的 $W_q x_m$ 乘上一个块对角旋转矩阵 $R_{\Theta, m}^d$:

$$ q_m = R_{\Theta, m}^d W_q x_m $$

其中旋转矩阵由不同的频率 $\theta_i = 10000^{-2i/d}$ 构成,与绝对位置 $m$ 相乘。最终自注意力中的点积计算变为:

$$ q_m^\top k_n = (R_{\Theta, m}^d W_q x_m)^\top (R_{\Theta, n}^d W_k x_n) = x_m^\top W_q^\top R_{\Theta, n-m}^d W_k x_n $$

物理意义: 由于旋转矩阵的正交性,这表明经过 RoPE 处理后,位置 $m$ 的 Query 寻找位置 $n$ 的 Key 时,实际上等价于在计算没有位置信息的点积之前,先对向量进行了相对距离为 $n-m$ 的旋转。这完美地在绝对编码的形式下实现了相对编码的特性。

实现细节:计算优化的伪代码

由于 $R_{\Theta, m}^d$ 是极其稀疏的块对角矩阵,直接进行矩阵相乘非常低效。论文给出了一种利用元素级操作的高效实现方法:

将向量相邻两项配对计算,伪代码逻辑表达为:
$x_{out} = \begin{pmatrix} x_1 \cos m\theta_1 - x_2 \sin m\theta_1 \\ x_1 \sin m\theta_1 + x_2 \cos m\theta_1 \\ \vdots \end{pmatrix}$

在深度学习框架中,通常会利用切片操作重排特征维度(例如将偶数位和奇数位特征抽出),然后再进行 Hadamard 乘积和加减,时间复杂度降低至 $O(d)$。

实验设置与复现细节

为了验证 RoPE 的有效性,作者在各类 NLP 经典任务上进行了广泛实验。(注:本文主要探讨预训练底层架构,未使用 LLM 的 Prompt few-shot 评测体系,而是基于微调评估)

1. 机器翻译 (Machine Translation)

  • 数据集: WMT 2014 英语-德语 (En-De) 数据集,约 450 万句子对。采用 joint BPE 分词,词表大小 37k。
  • 模型设置: 以 Transformer-base 为 baseline,替换掉绝对位置编码。
  • 超参数: 优化器 Adam ($\beta_1=0.9, \beta_2=0.98$)。学习率从 1e-7 线性 warmup 到 5e-4。使用标签平滑 (Label smoothing 0.1)。解码采用 Beam Search (beam size 4, length penalty 0.6)。框架基于 Fairseq。

2. 语言模型预训练 (Pre-training)

  • 数据集: BookCorpus 和 Wikipedia 组合数据集(英文)。
  • 模型: 基于 bert-base-uncased。对比 vanilla BERT 与注入了 RoPE 的 RoFormer。
  • 训练配置: Batch size = 64,最大序列长度 Max Seq Length = 512。共训练 100k 步。
  • 优化器: AdamW,学习率 1e-5。指标为 MLM (Masked Language Modeling) Loss。

3. 中文长文本验证 (Chinese Data Pre-training)

这是体现 RoPE 对长序列外推优势的核心实验。

  • 预训练语料: 34GB 的中文维基百科、新闻和论坛数据。
  • 预训练策略: 作者采用了一种多阶段递进预训练法。例如 Stage 1(序列长度 512,batch 256,训练 200k 步),随后逐步增加最大序列长度到 1536(Stage 2 和 Stage 5),验证 RoPE 的长度伸缩能力。
  • 下游任务: CAIL2019-SCM(中国法研杯相似案例匹配),大部分案例文本长度超过 512 字符。将测试集以 6:2:2 划分,对比 RoFormer 与基于词的 WoBERT 及基于字符的 BERT。

实验结果与核心发现

1. 收敛速度超越 Baseline

预训练 MLM Loss 曲线(参见原论文 Figure 3 左侧)显示,使用 RoPE 替换了绝对位置编码的 RoFormer,在训练前期和后期的 Loss 均持续低于原生 BERT。证明旋转信息提供了比绝对加法更优越的依赖指导。

2. GLUE 评测全方位提升

微调 3 个 Epoch(Max len 512,学习率遍历 2,3,4,5e-5)。在多项关键指标上超越 Baseline。

模型 MRPC (F1) SST-2 (Acc) QNLI (Acc) STS-B (Corr) QQP (F1)
BERT (Devlin 等, 2019) 88.9 93.5 90.5 85.8 71.2
RoFormer (本文) 89.5 90.7 88.0 87.0 86.4

💡 解读:在 QQP 问答匹配任务上,RoFormer 相对原版 BERT 带来了惊人的 15.2% 的跃升,MRPC 与 STS-B 也有稳定提升。虽然部分单句分类任务轻微下降,但长依赖匹配任务优势明显。

3. 中文长文本上的卓越外推性

针对 CAIL2019-SCM,论文分别对比了在最大长度 512 截断和 1024 截断下的表现。

模型配置 验证集 Acc 测试集 Acc
BERT-512 64.13% 67.77%
WoBERT-512 64.07% 68.10%
RoFormer-512 68.29% 64.13%
RoFormer-1024 66.07% 69.79%

结论机制:当最大输入长度放宽至 1024 时,RoFormer 依然能稳健工作并取得最佳的 69.79% 测试准确率,证明 RoPE 虽然是在较短长度训练的,但具备天然的相对距离约束,使得模型可以较好地泛化到比预训练时更长的序列上(即长度外推性)。

审稿人视角的犀利评论

✅ 优点与创新 (Strengths)

  1. 极其优美的理论闭环: 在众多玄学炼丹和堆砌结构的研究中,RoPE 基于复数几何的推导极为严谨。它漂亮地统一了绝对编码(每个 Token 保存自身位置的旋转角)和相对编码(点积时自然相减生成相对角度差),堪称现代 NLP 中最为优雅的设计之一。
  2. 计算范式的胜利: 由于保持了内积的形式,没有引入加法偏置,RoPE 天然兼容具有线性复杂度的注意力机制(如本文测试的 PerFormer),这打通了长文本线性注意力落地的最后一道障碍。
  3. 零额外参数: 基准角度 $\theta$ 是固定的,这让模型在实现强效位置感知的同时,完全没有任何新增的权重参数。

⚠️ 局限与疑问 (Weaknesses & Doubts)

  • 实验规模的局限: 受限于当时的计算资源,文中的预训练规模(Base 级别)和数据集偏小。在小模型上,各种位置编码的差距可能被数据的噪音淹没(从 GLUE 有升有降就能看出)。论文未能全面展示在大规模自回归生成语言模型(如 GPT 架构)上的碾压级威力。
  • 长程衰减特性的数学证明偏弱: 论文在 3.4.3 节试图证明内积随相对距离增加而衰减,使用了 Abel 变换进行放缩(建立上界)。但上界收敛并不代表函数本身稳定衰减(上界太松)。这是一个定性观察而非严格的下界保证。

🚀 未来拓展 (Future Directions)

作为审稿人,我会建议:既然复数域旋转赋予了外推性,那么当序列长度达到数十万(100K+)时,固定的 $\theta$ 基底可能会因为高频震荡失去局部精度。是否可以通过插值(Interpolation)而非外推来解决超长文本问题?(剧透:这正是后来基于 RoPE 的 NTK-aware 缩放和 YaRN 所做的事。

One More Thing ...

属于 RoPE 的真实荣光在论文之外

如果你只看这篇论文的实验表格,会觉得它只是一次“涨了几个点的微创新”。
但将时间轴拉到今天,你会发现 RoPE 构成了整个大模型时代的基石

从 Meta 的 LLaMA 系列,到 Google 的 PaLM,再到国内的 Qwen、ChatGLM 等几乎所有的主流开源大模型,无一例外抛弃了最初的绝对位置编码,全盘采用了 RoPE。为什么?

  • 无限上下文的催化剂: 2023年,研究人员发现只需将 RoPE 中的基底频率(Base)从 10000 调整到更大(如 1M),或者基于 NTK 插值法微调,就能在不重新预训练模型的前提下,让仅具有 4K 上下文的大模型直接支持 32K 甚至 128K 的长文本!这一切,全靠 RoPE 旋转矩阵严谨的数学可解释性。
  • 真正的“化繁为简”: 论文第一作者苏剑林最初在个人博客(科学空间)中发表了该思路。从草根博客走向顶级会议,再到影响几千亿参数的全球大模型,这种从极简数学直觉出发产生的爆炸性影响力,正是学术界最迷人的地方。