ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING

ROFORMER：使用旋转位置嵌入的增强版TRANSFORMER

作者/机构:
Jianlin Su (Zhuiyi Technology Co., Ltd. Shenzhen)
Yu Lu (Zhuiyi Technology Co., Ltd. Shenzhen)
Shengfeng Pan (Zhuiyi Technology Co., Ltd. Shenzhen)
Ahmed Murtadha (Zhuiyi Technology Co., Ltd. Shenzhen)
Bo Wen (Zhuiyi Technology Co., Ltd. Shenzhen)
Yunfeng Liu (Zhuiyi Technology Co., Ltd. Shenzhen)

A1 主要贡献

本文研究了将位置信息集成到基于Transformer的语言模型中的多种方法，并提出了一种名为旋转位置嵌入（Rotary Position Embedding, RoPE）的新方法来有效利用位置信息。

• 核心问题: 现有Transformer架构中的自注意力机制是位置无关的。虽然已提出多种位置编码方法（绝对位置编码和相对位置编码），但它们通常将位置信息加到上下文表示中，这种“加性”方式不适用于线性的自注意力架构。

• 研究目标: 提出一种新的位置编码方法，该方法能自然地将相对位置信息融入自注意力公式中，并具备序列长度灵活性、随相对距离增加而衰减的依赖性，以及与线性自注意力兼容的能力。

• 核心创新点 (RoPE):

  1. 提出旋转位置嵌入 (RoPE): 本文提出了一种新的位置编码方法RoPE。其核心思想是通过一个旋转矩阵来编码绝对位置，同时在自注意力公式中自然地引入了显式的相对位置依赖。具体而言，Ro-PE通过将上下文表示乘以一个旋转矩阵来编码相对位置，并提供了清晰的理论解释。
  2. 研究RoPE的特性: 论文研究并证明了RoPE具有优良的性质。首先，它具有长程衰减性，即词元间的依赖性随着相对距离的增加而减弱，这符合自然语言编码的直觉。其次，RoPE能够与线性自注意力机制兼容，而此前基于相对位置编码的方法通常不具备此能力。
  3. 实验验证: 论文在多个长文本分类基准数据集上评估了使用RoPE增强的Transformer模型（称为RoFormer）。实验结果表明，RoFormer的性能始终优于其他替代方案，证明了RoPE的有效性。

A3 背景知识与相关工作

2.1 预备知识

自注意力机制中的位置信息表示。给定一个包含N个输入词元$S_N = \{w_i\}_{i=1}^N$的序列，其对应的词嵌入为$E_N = \{x_i\}_{i=1}^N$，其中$x_i \in R^d$是词元$w_i$的d维词嵌入向量，不包含位置信息。自注意力机制首先将位置信息融入词嵌入，然后将它们转换为查询（query）、键（key）和值（value）表示。

其中，$q_m$、$k_n$和$v_n$分别通过函数$f_q$、$f_k$和$f_v$融入了第$m$个和第$n$个位置的信息。然后，查询和键的值被用来计算注意力权重，而输出则是通过对值表示进行加权求和计算得到的。

现有的基于Transformer的位置编码方法主要集中在为公式(1)选择一个合适的函数。

2.2 绝对位置嵌入

将位置向量直接添加到词嵌入中。公式(1)的一个典型选择是：

其中$p_i \in R^d$是一个依赖于词元$x_i$位置的d维向量。先前的工作如Devlin等人【索引4，Bert: Pre-training of deep bidirectional transformers for language understanding，2019，NAACL-HLT】、Lan等人【索引7，Albert: A lite bert for self-supervised learning of language representations，2020，ICLR】、Clark等人【索引8，ELECTRA: Pre-training text encoders as discriminators rather than generators，2020，ICLR】、Radford等人【索引5，Language models are unsupervised multitask learners，2019】以及Radford和Narasimhan【索引9，Improving language understanding by generative pre-training，2018】引入了一组可训练的向量$p_i \in \{p_t\}_{t=1}^L$，其中L是最大序列长度。而Vaswani等人【索引3，Attention is all you need，2017，NeurIPS】则提出使用正弦函数来生成$p_i$。

其中$p_{i,2t}$是d维向量$p_i$的第2t个元素。本文提出的RoPE从正弦函数的角度看与此直觉相关，但RoPE不是直接将位置信息加到上下文表示上，而是通过与正弦函数相乘来融入相对位置信息。

2.3 相对位置嵌入

将相对位置信息编码到注意力机制中。Shaw等人【索引11，Self-attention with relative position representations，2018，NAACL-HLT】应用了公式(1)的不同设置如下：

其中$p˜_r^k, p˜_r^v \in R^d$是可训练的相对位置嵌入。$r = \text{clip}(m - n, r_{\min}, r_{\max})$表示位置m和n之间的相对距离。他们裁剪相对距离的假设是，精确的相对位置信息在超过一定距离后不再有用。保持公式(3)的形式，Dai等人【索引13，Transformer-xl: Attentive language models beyond a fixed-length context，2019，ACL】提出将公式(2)中的$q_m^T k_n$分解为：

其核心思想是将绝对位置嵌入$p_n$替换为其正弦编码的相对对应物$p˜_{m-n}$，同时将第三项和第四项中的绝对位置$p_m$替换为两个与查询位置无关的可训练向量u和v。此外，用于基于内容的键向量$x_n$和基于位置的键向量$p_n$的权重矩阵$W_k$被区分为$W_k$和$W_k$，从而得到：

值得注意的是，值项中的位置信息通过设置$f_v(x_j) := W_v x_j$被移除了。后来的工作，如Raffel等人【索引15，Exploring the limits of transfer learning with a unified text-to-text transformer，2020，J. Mach. Learn. Res.】、He等人【索引17，Deberta: Decoding-enhanced bert with disentangled attention，2020，ArXiv】、Ke等人【索引16，Rethinking positional encoding in language pre-training，2020，ArXiv】、Huang等人【索引18，Improve transformer models with better relative position embeddings，2020，EMNLP】都遵循了这些设置，仅将相对位置信息编码到注意力权重中。然而，Raffel等人【索引15，Exploring the limits of transfer learning with a unified text-to-text transformer，2020，J. Mach. Learn. Res.】将公式(6)改写为：

其中$b_{i,j}$是一个可训练的偏置。Ke等人【索引16，Rethinking positional encoding in language pre-training，2020，ArXiv】研究了公式(6)的中间两项，发现绝对位置和单词之间几乎没有相关性。Raffel等人【索引15，Exploring the limits of transfer learning with a unified text-to-text transformer，2020，J. Mach. Learn. Res.】提出使用不同的投影矩阵来建模一对单词或位置。

He等人【索引17，Deberta: Decoding-enhanced bert with disentangled attention，2020，ArXiv】认为，两个词元之间的相对位置只能通过公式(6)的中间两项来完全建模。因此，他们简单地将绝对位置嵌入$p_m$和$p_n$替换为相对位置嵌入$p˜_{m-n}$：

对四种相对位置嵌入变体的一项比较【索引9，Improving language understanding by generative pre-training，2018】表明，与公式(10)相似的变体是其中最有效的。总的来说，所有这些方法都试图在Vaswani等人【索引3，Attention is all you need，2017，NeurIPS】最初提出的自注意力设置下，基于公式(3)的分解来修改公式(6)。它们通常都是将位置信息直接添加到上下文表示中。与此不同，本文的方法旨在在某些约束下从公式(1)推导出相对位置编码，这种推导出的方法通过旋转上下文表示来融入相对位置信息，从而更具可解释性。

A2 方法细节

3.1 公式化

将相对位置信息融入内积计算。基于Transformer的语言建模通常通过自注意力机制利用单个词元的位置信息。如公式(2)所示，$q_m^T k_n$通常在不同位置的词元之间传递知识。为了融入相对位置信息，我们要求查询$q_m$和键$k_n$的内积能由一个函数g来表达，该函数仅以词嵌入$x_m, x_n$及其相对位置$m-n$作为输入变量。换句话说，我们希望内积仅以相对形式编码位置信息：

最终目标是找到一个等效的编码机制来求解函数$f_q(x_m, m)$和$f_k(x_n, n)$以满足上述关系。

3.2 旋转位置嵌入

3.2.1 二维情况

从二维情况出发寻找解决方案。我们从维度d=2的简单情况开始。在这种设置下，我们利用二维平面上向量的几何性质及其复数形式来证明（详见3.4.1节），我们的公式(11)的一个解是：

其中$Re[·]$表示复数的实部，$(W_k x_n)^*$表示$(W_k x_n)$的共轭复数。$θ \in R$是一个预设的非零常数。我们可以进一步将$f_{\{q,k\}}$写成矩阵乘法形式：

其中$(x_m^{(1)}, x_m^{(2)})$是$x_m$在二维坐标系中的表示。类似地，g也可以被看作一个矩阵，从而在二维情况下为3.1节的公式化提供了解决方案。具体来说，融入相对位置嵌入变得非常直接：只需将经过仿射变换的词嵌入向量旋转其位置索引的角度倍数即可，这解释了旋转位置嵌入背后的直觉。

3.2.2 通用形式

将二维解推广到高维空间。为了将我们在二维中的结果推广到任意$x_i \in R^d$（其中d为偶数），我们将d维空间划分为d/2个子空间，并利用内积的线性特性将它们组合起来，使得$f_{\{q,k\}}$变为：

其中

是旋转矩阵，其预定义参数为$\Theta = \{\theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, ..., d/2]\}$。RoPE的图形化说明如图1所示。将我们的RoPE应用于公式(2)中的自注意力，我们得到：

其中$R^d_{\Theta,n-m} = (R^d_{\Theta,m})^T R^d_{\Theta,n}$。需要注意的是，$R^d_{\Theta}$是一个正交矩阵，这确保了在编码位置信息过程中的稳定性。此外，由于$R^d_{\Theta}$的稀疏性，直接如公式(16)中那样进行矩阵乘法计算效率不高；我们在理论解释部分提供了另一种实现方式。与之前工作中采用的加性位置嵌入方法（即公式(3)到(10)）相反，我们的方法是乘性的。而且，RoPE通过旋转矩阵的乘积自然地融入了相对位置信息，而不是在加性位置编码的展开式中修改项。

3.3 RoPE的性质

长程衰减特性。遵循Vaswani等人【索引3，Attention is all you need，2017，NeurIPS】的做法，我们设置$\theta_i = 10000^{-2i/d}$。可以证明，这种设置提供了一个长程衰减的特性（详见3.4.3节），这意味着当相对位置增加时，内积会衰减。这个特性与一对相距很远的词元之间应该有较少联系的直觉相符。

RoPE与线性注意力的结合。自注意力可以被重写为一个更通用的形式。

原始的自注意力选择$sim(q_m, k_n) = \exp(q_m^T k_n / \sqrt{d})$。注意，原始自注意力需要计算每对词元的查询和键的内积，这具有$O(N^2)$的二次复杂度。遵循Katharopoulos等人【索引22，Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】的方法，线性注意力将公式(17)重新表述为

其中$\phi(·)$和$\varphi(·)$通常是非负函数。Katharopoulos等人【索引22，Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】提出$\phi(x) = \varphi(x) = \text{elu}(x) + 1$，并利用矩阵乘法的结合律首先计算键和值之间的乘积。Shen等人【索引23，Efficient attention: Attention with linear complexities，2021，WACV】在内积之前使用softmax函数分别对查询和键进行归一化，这等价于$\phi(q_i) = \text{softmax}(q_i)$和$\phi(k_j) = \exp(k_j)$。由于RoPE通过旋转注入位置信息，这保持了隐藏表示的范数不变，我们可以通过将旋转矩阵与非负函数的输出相乘，来将RoPE与线性注意力结合起来。

值得注意的是，我们保持分母不变以避免除以零的风险，并且分子中的求和可能包含负项。尽管公式(19)中每个值$v_i$的权重没有严格进行概率归一化，但我们认为这种计算仍然可以建模值的重要性。

3.4 理论解释

3.4.1 二维情况下RoPE的推导

在二维空间中通过复数形式进行推导。在d=2的情况下，我们考虑对应于查询和键的两个词嵌入向量$x_q, x_k$及其位置m和n。根据公式(1)，它们编码了位置信息的对应向量是：

其中$q_m$和$k_n$的下标表示编码的位置信息。假设存在一个函数g定义了由$f_{\{q,k\}}$产生的向量之间的内积：

我们还要求满足以下初始条件：

这可以理解为编码了空位置信息的向量。在这些设定下，我们尝试找到$f_q, f_k$的解。首先，我们利用向量在二维中的几何意义及其复数对应形式，将公式(20)和(21)中的函数分解为：

其中$R_f, R_g$和$\Theta_f, \Theta_g$分别是$f_{\{q,k\}}$和g的径向和角向分量。将它们代入公式(21)，我们得到关系：

相应的初始条件为：

其中$\|q\|, \|k\|$和$\theta_q, \theta_k$是q和k在二维平面上的径向和角向部分。接下来，我们在公式(24)中设置m=n并考虑公式(25)中的初始条件：

一方面，从公式(26a)可以直接得到$R_f$的一个解：

这解释了径向函数$R_q, R_k$和$R_g$与位置信息无关。另一方面，从公式(26b)中可以注意到，$\Theta_q(x_q, m) - \theta_q = \Theta_k(x_k, m) - \theta_k$表明角向函数不依赖于查询和键，我们将其设为$\Theta_f := \Theta_q = \Theta_k$，且项$\Theta_f(x_{\{q,k\}}, m) - \theta_{\{q,k\}}$是位置m的函数，与词嵌入$x_{\{q,k\}}$无关，我们记为$\phi(m)$，得到：

进一步，将$n=m+1$代入公式(24)并考虑上述方程，我们可以得到：

由于右侧是一个与m无关的常数，具有连续整数输入的$\phi(m)$产生一个等差数列：

其中$\theta, \gamma \in R$是常数且$\theta$非零。总结我们从公式(27)到(30)的解：

注意，我们没有对公式(22)中的$f_q$和$f_k$施加任何约束，因此$f_q(x_m, 0)$和$f_k(x_n, 0)$可以自由选择。为了使我们的结果与公式(3)具有可比性，我们定义：

然后，我们在最终解的公式(31)中简单地设置$\gamma = 0$：

3.4.2 旋转矩阵乘法的计算高效实现

利用旋转矩阵的稀疏性进行高效计算。利用公式(15)中$R^d_{\Theta,m}$的稀疏性，$R^d$与$x \in R^d$的乘法可以更高效地实现为：

3.4.3 RoPE的长程衰减

通过阿贝尔变换证明长程衰减特性。我们可以将向量$q = W_q x_m$和$k = W_k x_n$的元素成对分组，公式(16)中RoPE的内积可以写成复数乘法。

其中$q_{[2i:2i+1]}$表示q的第2i到(2i+1)个元素。记$h_i = q_{[2i:2i+1]} k^*_{[2i:2i+1]}$和$S_j = \sum_{i=0}^{j-1} e^{i(m-n)\theta_i}$，并设$h_{d/2} = 0$和$S_0 = 0$，我们可以使用阿贝尔变换重写求和：

因此，

通过设置$\theta_i = 10000^{-2i/d}$，可以观察到$\frac{1}{d/2} \sum_{i=1}^{d/2} |S_i|$的值随着相对距离$m-n$的增加而衰减，如图2所示。

A4 实验环境

• 数据集:

  • 机器翻译: WMT 2014 English-German 数据集【索引26】，包含约450万个句子对。
  • 英语预训练: BookCorpus【索引30】和 Wikipedia Corpus【索引31】，来源于Huggingface Datasets库。
  • GLUE下游任务: MRPC【索引33】，SST-2【索引34】，QNLI【索引35】，STS-B【索引36】，QQP【索引37】，MNLI【索引38】。
  • 线性注意力实验: Enwik8 数据集【索引41】，源自英文维基百科。
  • 中文预训练: 约34GB数据，收集自中文维基百科、新闻和论坛。
  • 中文长文本下游任务: CAIL2019-SCM【索引44】，包含8964个三元组案例，用于相似案例匹配。

• 模型架构:

  • 机器翻译: 基于Transformer【索引3】的基线模型进行修改。
  • 英语预训练: bert-base-uncased【索引4】作为基线。
  • 线性注意力: 12层、768维、12头的字符级PerFormer【索引25】。
  • 中文任务: 基于WoBERT【索引42】进行修改，与其他中文预训练模型（BERT, NEZHA【索引43】）进行比较。

• 硬件配置:

  • 在两个云服务器上运行，每个服务器配备4块V100 GPU。

• 软件配置:

  • 代码实现: 使用PyTorch。
  • 依赖库: fairseq toolkit【索引28】 (MIT License), Huggingface Datasets library (Apache License 2.0), Huggingface Transformers library【索引40】 (Apache License 2.0)。
  • 优化器: Adam，AdamW。

A4 实验结果

4.1 机器翻译

• 实验内容: 在WMT 2014英德翻译任务上，将RoFormer与基线Transformer模型进行比较。
• 实验结果: 如表1所示，RoFormer在BLEU得分上优于基线Transformer模型。
  

  

4.2 预训练语言模型

• 实验内容: 在BookCorpus和Wikipedia语料库上进行预训练，比较RoFormer和BERT的MLM损失。
• 实验结果: 如图3左图所示，与原始BERT相比，RoFormer收敛更快。
  

  

4.3 GLUE任务微调

• 实验内容: 将预训练好的RoFormer在多个GLUE下游任务上进行微调，评估其泛化能力。
• 实验结果: 如表2所示，RoFormer在六个数据集中有三个（MRPC, STS-B, QQP）显著优于BERT，且提升幅度可观。
  

  

4.4 Performer与RoPE结合

• 实验内容: 将RoPE集成到线性注意力的PerFormer模型中，并在Enwik8数据集上进行预训练，比较有无RoPE的损失曲线。
• 实验结果: 如图3右图所示，将RoPE加入Performer后，在相同训练步数下实现了更快的收敛和更低的损失。

4.5 中文数据评估

• 预训练:
  • 实验内容: 在大规模中文语料上分阶段预训练RoFormer，逐步增加最大序列长度。
  • 实验结果: 如表4所示，随着序列长度上限的增加，RoFormer的准确率也随之提升，证明了其处理长文本的能力。
    

    

    
    

    

• 下游任务 (CAIL2019-SCM):
  • 实验内容: 在长文本相似案例匹配任务CAIL2019-SCM上，比较RoFormer与BERT、WoBERT在不同输入长度下的表现。
  • 实验结果: 如表5所示，在短文本截断（512）时，RoFormer与WoBERT相当；当最大输入长度增加到1024时，RoFormer的性能比WoBERT高出1.5个百分点。
    

    

4.5.5 工作局限性

• 尽管在数学上将相对位置关系格式化为二维子空间下的旋转，但缺乏对为何它比其他位置编码策略收敛更快的深入解释。
• 尽管证明了模型具有有利的长程衰减特性，这与其他位置编码机制相似，但模型在长文本上表现更优的原因尚未得到可靠的解释。
• RoFormer基于Transformer架构，需要大量硬件资源进行预训练。

A5 结论

本文提出了一种新的位置嵌入方法RoPE，它在自注意力机制中融入了显式的相对位置依赖，以提升Transformer架构的性能。理论分析表明，相对位置可以自然地通过自注意力中的向量乘积来表达，而绝对位置信息则通过一个旋转矩阵进行编码。此外，本文从数学上阐述了该方法应用于Transformer时的有利特性。最后，在英文和中文基准数据集上的实验表明，该方法在预训练中能促进更快的收敛，并且所提出的RoFormer在长文本任务上能取得更好的性能。

参考文献引用总结

• 【索引3，Attention is all you need, 2017, NeurIPS】
  • 引用段落：A3-2.2, A3-2.3, A2-3.3
  • 引用描述：在介绍绝对位置嵌入时，提到Vaswani等人提出使用正弦函数生成位置向量；在回顾相对位置嵌入时，指出所有相关方法都是基于其最初提出的自注意力设置；在讨论RoPE的长程衰减特性时，提到遵循其设置$\theta_i$的值。
• 【索引4，Bert: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】
  • 引用段落：A3-2.2
  • 引用描述：在介绍绝对位置嵌入时，作为使用可训练位置向量的代表性工作之一被引用。
• 【索引5，Language models are unsupervised multitask learners, 2019】
  • 引用段落：A3-2.2
  • 引用描述：在介绍绝对位置嵌入时，作为使用可训练位置向量的代表性工作之一被引用。
• 【索引7，Albert: A lite bert for self-supervised learning of language representations, 2020, ICLR】
  • 引用段落：A3-2.2
  • 引用描述：在介绍绝对位置嵌入时，作为使用可训练位置向量的代表性工作之一被引用。
• 【索引8，ELECTRA: Pre-training text encoders as discriminators rather than generators, 2020, ICLR】
  • 引用段落：A3-2.2
  • 引用描述：在介绍绝对位置嵌入时，作为使用可训练位置向量的代表性工作之一被引用。
• 【索引9，Improving language understanding by generative pre-training, 2018】
  • 引用段落：A3-2.2, A3-2.3
  • 引用描述：在介绍绝对位置嵌入时，作为使用可训练位置向量的代表性工作之一被引用；在介绍相对位置嵌入时，提到其工作对四种相对位置嵌入变体进行了比较。
• 【索引11，Self-attention with relative position representations, 2018, NAACL-HLT】
  • 引用段落：A3-2.3
  • 引用描述：作为相对位置嵌入的开创性工作之一，介绍了其通过可训练嵌入来表示相对距离的方法。
• 【索引13，Transformer-xl: Attentive language models beyond a fixed-length context, 2019, ACL】
  • 引用段落：A3-2.3
  • 引用描述：介绍了其对注意力内积计算的分解方法，将相对位置信息引入到键向量的计算中。
• 【索引15，Exploring the limits of transfer learning with a unified text-to-text transformer, 2020, J. Mach. Learn. Res.】
  • 引用段落：A3-2.3
  • 引用描述：介绍了其对相对位置编码的简化，即使用一个可训练的偏置项$b_{i,j}$，以及提出使用不同投影矩阵建模词对或位置对。
• 【索引16，Rethinking positional encoding in language pre-training, 2020, ArXiv】
  • 引用段落：A3-2.3
  • 引用描述：介绍了其研究发现绝对位置与单词之间相关性不大，为简化相对位置编码提供了依据。
• 【索引17，Deberta: Decoding-enhanced bert with disentangled attention, 2020, ArXiv】
  • 引用段落：A3-2.3
  • 引用描述：介绍了其提出的观点，即仅用中间两项即可完全建模相对位置，并用相对位置嵌入替代了绝对位置嵌入。
• 【索引18，Improve transformer models with better relative position embeddings, 2020, EMNLP】
  • 引用段落：A3-2.3
  • 引用描述：作为遵循仅将相对位置信息编码到注意力权重中的后续工作之一被引用。
• 【索引22，Transformers are rnns: Fast autoregressive transformers with linear attention, 2020, ICML】
  • 引用段落：A2-3.3
  • 引用描述：在讨论线性注意力时，引用其工作将自注意力重构为线性复杂度的形式，并提出了具体的非负核函数。
• 【索引23，Efficient attention: Attention with linear complexities, 2021, WACV】
  • 引用段落：A2-3.3
  • 引用描述：作为线性注意力的另一项工作，引用了其使用softmax分别归一化query和key的方法。