更新于 2026年7月21日

在之前的系列文章当中,我们已经详细介绍了 Transformer 的原理及实现过程,其中我们就谈到自注意力机制(self-attention)的一个弊端之一就是无法捕捉到序列之间的位置关系,换句话说如过我们交换了序列中词元 token 之间的位置,self-attention 是无法感知到这一变化的,因此作者提出了通过位置编码来解决这一问题。

不过在后来有研究发现[1],即使去掉位置编码,基于 Transformer 架构的语言模型在各个任务上的效果仍旧不会有明显下降(影响不大),换句话说自注意力机制本身也能够感知到序列中 token 位置的变化,只是感知能力不够强,关于这一点大家可以自行去分析self-attention 的计算过程。

Causal transformer language models (LMs), such as GPT-3, typically require some form of positional encoding, such as positional embeddings. However, we show that LMs without any explicit positional encoding are still competitive with standard models, and that this phenomenon is robust across different datasets, model sizes, and sequence lengths. [1]

当然,我们自己也通过一个实际的实验来验证了这一点。在基于 Bert 模型的文本分类任务中,我们发现发现去掉 Positional embedding 和保留 Positional embedding 时,模型在 5 个 epoch 中 validation 上的准确率分别是(0.872, 0.877, 0.878 ) 和 (0.884, 0.885, 0.883),也就是并没有太大的影响。

不过虽说影响不是很大,但总归是有影响的,或者说现有的位置编码技术并没有很好的解决这一问题,因此也就陆续出现了各种不同的位置编码算法。

在接下来的内容中,将会介绍一种几乎已经成为了现在各个大模型标配的位置编码方法——旋转位置编码(Rotary Position Embedding ,RoPE)[2]。

关键字:旋转位置编码、RoPE、位置编码、LLM、Transformer、绝对位置编码、相对位置编码

1. 位置编码方法#

在正式介绍介绍旋转位置编码之前,我们先来回顾一下在 Transformer 模型中常见的绝对位置编码( Absolute Position Encoding)和相对位置编码(Relative Position Encoding) 的区别与含义。

1.1 绝对位置编码#

上面我们说到 Transformer 中自注意力机制本身结构是无序的(不像 RNN 有顺序),所以必须人为地加入位置信息来弥补这一点,因此便引入了绝对位置编码。绝对位置编码是指:我们为序列中每个 token 显式指定它在句子中的具体位置,例如第 1 个 token 的位置是什么、第 2 个 token 的位置是什么……并编码进输入中,可以看出这从全局的角度来考虑每个 token 的位置信息。

在原始的 Transformer 中,作者通过使用固定的三角函数编码的方式来计算每个 token 对应的位置信息,即

$$ PE_{(pos, 2i)} = \sin \left( \frac{pos}{10000^{2i/d_{\text{model}}}} \right), \quad PE_{(pos, 2i+1)} = \cos \left( \frac{pos}{10000^{2i/d_{\text{model}}}} \right)\tag{1} $$

其中 $pos$ 表示每个 token 的位置编号,$i$ 表示该位置上对应的向量维度。

最后,将计算出的位置信息以加法操作的方式与原始序列的 token embedding 结果进行融合,便得到了序列含有位置信息的向量化表示。

所以,绝对位置编码的好处是不同位置的编码是唯一的,并且可以外推到更长的序列中。

尽管在后来的 Bert 模型中,固定的三角函数编码方式被替换成了可训练的权重参数,但是本质上依旧是为每个位置分配了一个固定的位置向量。

1.2 相对位置编码#

那什么又是相对位置编码呢?

简单来说在相对位置编码中,与其告诉模型你在第几个位置,它更关心的是大家彼此间的相对位置。绝对位置编码是在词向量中告诉模型“这是第几个词”,而相对位置编码则是告诉模型“你和别的词相隔几个词”。

假设你是一个列车员,车厢里有编号 0 到 3 的座位。

  • 用绝对位置编码,你知道:“A” 坐在 1 号位,“B” 坐在 3 号位,”C“ 坐在 2 号位置;

  • 用相对位置编码,你知道:“B” 离 ”A“ 的座位距离为 2,”C“ 离 ”A“ 的座位距离为 1,而这些相对距离就会影响模型分配注意力的程度。

可以看出,相对位置信息对于位置的变化其实是更敏感的,更容易刻画其位置的变换。当然,某种程度上说绝对位置其实也是一种单一的相对位置,即仅以全局某个固定位置为参考点。

不过对于不同的相对位置编码方法其具体计算过程也存在着不小的差异,但是总体思想基本都是类似。

2. 旋转位置编码#

在简单介绍完绝对位置编码和相对位置编码的基本思想以后,我们继续来看为什么要提出旋转位置编码,动机到底是什么。

2.1 旋转编码动机#

作者在文章中谈到,尽管现在各种各样的位置编码方法都具有一定的成效,但是它们通常都是将位置信息以加的形式融入到上下文表示中,因此并不适合线性自注意力机制架构。同时,相对位置编码相对于绝对位置编码更有优势,因此作者提出了首个既采用相对位置编码又同时支持线性自注意力机制架构的旋转位置编码技术。

Despite the effectiveness of these approaches, they commonly add the position information to the context representation and thus render them unsuitable for the linear self-attention architecture.

那什么又是线性自注意力呢?

简单总结就是,尽管自注意力机制有着非常好的效果,计算过程如式(2)所示,但是随着模型参数量的不断增加,整个过程的时间复杂度变成了一个不容忽视的缺陷。

$$ \text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\in\mathbb{R}^{n\times d_v} \tag{2} $$

其中 $Q\in\mathbb{R}^{n\times d_{q}}$,$K\in\mathbb{R}^{n\times d_{k}}$,$V\in\mathbb{R}^{n\times d_{v}}$,且通常 $d_q=d_k=d_v=d$ 。

根据式(2)可知,在计算注意力权重矩阵 $QK^T$ 的时候,其时间复杂度为 $O(n^2d)$。但是,我们观察发现,如过式(2)中没有 $\text{softmax}$ 这个操作,那么整个自注意力计算过程便可以改写为 $QK^TV$ (暂时忽略scale)。同时,矩阵乘法满足结合律所以可以先计算 $K^TV$ 得到一个 $d\times d$ 的矩阵,然后再左乘一个 $Q$ 便得到了最后的结果,此时的时间复杂度为 $O(nd^2)$ 。

这,就是线性注意力,也是自注意力机制最终将要优化的方向。

例如,我们考虑矩阵链 $Q \cdot K^T \cdot V$ 的两种乘法顺序,其中 $Q \in \mathbb{R}^{100 \times 10}, K \in \mathbb{R}^{100 \times 10}, V \in \mathbb{R}^{100 \times 10}$

情况一: $(QK^T)V$ 顺序

在这样的情况下,计算 $QK^T$ 的时间复杂度为 $100\times10\times100=100000$ ;再计算 $(QK^T)V$,时间复杂度为 $100\times100\times10=100000$ ;总的复杂度为 $100000+100000=200000$ 。

情况二: $Q(K^TV)$ 顺序

在这样的情况下,计算 $K^TV$ 的时间复杂度为 $10 \times 100 \times 10 = 10000$ ;再计算 $Q(K^TV)$,时间复杂度为 $100 \times 10 \times 10 = 10000$;总的复杂度为 $10000 + 10000 = 20000$ 。

可以看到,最后两种计算顺序的结果虽然一样,但是时间复杂度差距却高达 10 倍!

至于为什么已有的位置编码不适合线性自注意力框架,我们后面再进行介绍。

基于这样的动机,Jianlin Su 等人提出了旋转位置编码,并且证明它适用于线性自注意力框架。

不过,可能是由于效果的原因,最后作者在实现的 Roformer 模型时使用的依旧是原始的自注意力计算过程,只是位置编码使用的是 RoPE 编码,但这并不影响 RoPE 是一种更优的相对位置编码方式。

因此,本文也将只会介绍 RoPE 的原理及实现过程!

2.2 旋转编码思想#

一般来说,绝对位置编码具有实现简单、计算速度快等优点,而相对位置编码则直接地体现了相对位置的变化,这也跟我们的直观理解吻合。

在这样的背景和动机下,作者便开始假设,既然自注意力的核心就是内积,也就是注意力的计算过程,那么我们能不能假设向量 $q^Tk$ 计算完成以后的注意力权重结果自然就带有了相对位置的信息呢?

这个大家听起来可能觉得有点奇怪,但其实一点也不怪!

因为注意力的分配其实就是和序列中每个 token 的顺序是有关系的,同一个 token 集合以不同的顺序组合成不同的序列语义不一样,分配到同一个 token 上的权重肯定也是不一样的,所以注意力值中自然也就包含了序列的位置信息。

因此,作者就假设存在函数 $f(x,pos)$,它能够给向量 $q$ 和 $k$ 分别施加上在位置 $m$ 和 $n$ 处的绝对位置信息,即 $\tilde{q}_m=f(q,m)$ 和 $\tilde{k}_n=f(k,n)$;最后,通过 $\tilde{q}_m^T\tilde{k}_n$ 在完成注意力权重计算以后就自然包含有了 $m$ 与 $n$ 之间的相对位置信息 $m-n$。

由此可见,作者想以一种表面上看似是绝对编码的计算过程,实则背后隐含着相对位置编码的本质,即通过绝对位置编码的计算过程,来实现相对位置编码的结果。

最后,之所以整个编码过程被取名为“旋转位置编码”是因为作者按照上述构想,求解得到的计算公式带有矩阵旋转的过程,所以取名为旋转位置编码。

2.3 旋转编码原理#

「第2讲 Transformer 位置编码教程:图解正弦 Positional Encoding 与编解码过程」中,我们已经详细介绍了多头注意力机制的原理和过程。不过为了后面能够一步步得出 RoPE 的计算原理,这里我们先再来以一种更加广义的方式来回顾一下自注意力机制的计算过程。

2.3.1 自注意力计算过程#

设 $\mathbb{S}_N=\{w_i\}^N_{i=1}$ 表示长度为 $N$ 的原始输入序列, $w_i$ 则表示第 $i$ 个 token,同时 $\mathbb{S}_N$ 对应的词嵌入表示为 $\mathbb{E}_N=\{x_i\}^N_{i=1}$ ,$x_i\in \mathbb{R}^d$ 则表示第 $i$ 个 token $w_i$ 对应的词嵌入形式(不包含位置信息)。

此时,自注意力机制会首先对序列的词嵌入表示融入位置信息,然后计算得到 queries、keys 和 values

$$ \begin{aligned} q_m&=f_q(x_m,m)\\[2ex] k_n&=f_k(x_n,n)\\[2ex] v_n&=f_v(x_n,n) \end{aligned}\tag{3} $$

其中 $q_m$ ,$k_n$ 和 $v_n$ 分别表示通过 $f_q$ , $f_k$ 和 $f_v$ 给 $x_i$ 的第 $m$ 和 $n$ 个位置融入位置信息后的结果。

进一步,开始计算注意力权重及输出

$$ \begin{aligned} \alpha_{m,n}&=\frac{\exp\left(\frac{q^T_mk_n}{\sqrt{d}}\right)}{\sum_{j=1}^N\exp\left(\frac{q^T_mk_j}{\sqrt{d}}\right)}\\[3ex] o_m&=\sum^N_{n=1}a_{m,n}v_n \end{aligned}\tag{4} $$

下面,我们继续来看在基于上述计算框架下的不同编码方式。

2.3.2 绝对位置编码计算过程#

在基于上述计算框架下,一种典型的绝对位置编码便是 Transformer 论文中所提出的位置编码方式,此时有

$$ f_{t:t\in\{q,k,v\}}(x_i,i):=W_{t:t\in\{q,k,v\}}(x_i+p_i)\tag{5} $$

其中 $W_t$ 表示线性变化中的权重, $p_i\in\mathbb{R}^d$ 便是对应的绝对位置编码信息,即

$$ \begin{cases} p_{i,2t}&=\sin(k/10000^{2t/d})\\[3ex] p_{i,2t+1}&=\cos(k/10000^{2t/d}) \end{cases}\tag{6} $$

此时可以看出,式(4)~式(6)便是 Transformer 中自注意力机制的计算过程。

从式(5)中我们可以看出,因为有 $x_i+p_i$ 中加操作的存在,所以在计算注意力权重矩阵的时候会拆开得到多项,因此并不适合线性自注意力框架。

2.3.3 相对位置编码计算过程#

然而,RoPE 并不直接将位置信息通过相加的方式融入到上下文表示中,而是直接通过乘以正弦函数来合并相对位置信息。

However, instead of directly adding the position to the context representation, RoPE proposes to incorporate the relative position information by multiplying with the sinusoidal functions.

为了在自注意力机制中融入相对位置信息,我们可以设想有一个函数 $g$,并以 $x_m,x_n$ 以及它们的相对位置信息 $m-n$ 作为 $g$ 的输入,最终计算得到输出 $q_m$ 和 $k_n$ 内积的结果。换句话说我们希望存在这样一个函数 $g$,使得它满足

In order to incorporate relative position information, we require the inner product of query $q_m$ and key $k_n$ to be formulated by a function $g$, which takes only the word embeddings $x_m, x_n$ and their relative position $m-n$ as input variables.

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

上面式(7)的含义是,我们假定存在 $g(x_m,x_n,m-n)$ 这样的相对位置信息编码方式,它能够等价地通过左侧绝对位置编码内积的计算方式来计算得到,也就是论文中描述的“通过绝对位置编码的计算过程,来实现相对位置编码的结果”的含义。

RoPE encodes the absolute position with a rotation matrix and meanwhile incorporates the explicit relative position dependency in self-attention formulation.

这里需要注意的一点是,RoPE 考量的是 queries 中的向量 $q_m$ 和 keys 中的向量 $k_n$ 之间的相对位置位置关系。

因此,对于式(7)中的假设,其最终目标便是在满足上面猜想的情况下,寻找得到 $f_q(x_m,m)$ 和 $f_k(x_n,n)$ 这两个函数。

The ultimate goal is to find an equivalent encoding mechanism to solve the functions $f_q(x_m,m)$ and $f_k(x_n,n)$ to conform the aforementioned relation.

那最终找到的 $f_q(x_m,m)$ 和 $f_k(x_n,n)$ 长什么样呢?

以二维平面为例,作者最终推导得到的 $f_q(x_m,m)$ 和 $f_k(x_n,n)$ 为

$$ f_{\{q,k\}}(x_m,m)= \begin{bmatrix}\cos m\theta&-\sin m\theta\\[2ex] \sin m\theta&\cos m\theta \end{bmatrix} \begin{bmatrix} W^{(11)}_{\{q,k\}}&W^{(12)}_{\{q,k\}}\\[2ex] W^{(21)}_{\{q,k\}}&W^{(22)}_{\{q,k\}} \end{bmatrix} \begin{bmatrix} x^{(1)}_m\\[2ex]x^{(2)}_m \end{bmatrix}\tag{8} $$

其中 $\theta\in\mathbb{R}$ 是一个非零常数,$(x^{(1)}_m,x^{(2)}_m)$ 是 $x_m$ 在二维平面下的表示。

从式(8)可以看出,最右侧两个矩阵相乘其实就是原始词嵌入表示经过一次线性变换后的结果;然后其左侧的这个矩阵被称为旋转矩阵,其作用是让其右侧的向量逆时针旋转 $m\theta$ 的角度,其中 $\theta$ 计算方式为

$$ \Theta=\{\theta_i=10000^{-2(i-1)/d},i\in[1,2,...,d/2]\}\tag{9} $$

当然,对于多维情况下的计算过程,可以把整个维度分成两组来完成计算,如图1所示。

图 1. 旋转编码计算过程图
图 1. 旋转编码计算过程图

在图1中,对于左下角中的每一个词向量,每两个维度为一组使用同一个 $\theta$ 由 $x_1,x_2$ 旋转为 $x_1^{\prime},x_2^{\prime}$ 得到新的结果,然后再进行后续同样的注意力计算。

2.3.4 图示计算示例#

假设现在样本序列长度为 5 ,那么首先需要根据式(9)给每个位置构建一个旋转角度即 $m\theta$ , 且因为这是二维情况所以只需要计算 $\theta_1=1$ 即可,如果是4维则还需要计算 $\theta_2$;然后便可计算其对应的正弦和余弦值,整个过程如图2所示。

图 2. 旋转矩阵计算图
图 2. 旋转矩阵计算图

在图2中,左侧的 $m$ 表示的便是序列中每个 token 的位置序号,再将其与不同维度上的 $\theta$ 相乘便得到了不同位置不同维度上所需要旋转的最终角度。

最后,再对旋转矩阵右乘 $q$ 或 $k$ 即可得到其融入绝对位置信息后的编码结果,如图3所示。

图 3. 旋转编码计算图
图 3. 旋转编码计算图

注意,此时得到结果只是包含绝对位置信息编码的结果。

在计算得到旋转后的 queries 和 keys 后,再按照式(2)完成后续计算即可。

以上便是旋转编码的原理及计算示例过程,当然在实际编码实现时还可以表示成图4所示的更加高效的形式。

图 4. 旋转编码展开计算图
图 4. 旋转编码展开计算图

2.4 从零实现旋转编码#

在清楚 RoPE 的原理以后,我们再来看如何从零实现这一过程。根据上面的计算示例可以知道,整个过程一共需要两步:① 计算每个位置在不同维度上的旋转角度;② 将其作用于 $q$ 或 $k$ 得到绝对位置信息编码后的输出。

首先,编码实现计算每个位置在不同维度上的旋转角度,示例代码如下:

1 def precompute_freqs_cos_sin(dim: int, seq_len: int, theta: float = 10000.0):
2     freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) #  [dim/2,]
3     position = torch.arange(seq_len).unsqueeze(1)  # [seq_len,1]
4     rot_theta = position * freqs  # (seq_len, dim/2)
5     cos = rot_theta.cos()  # [seq_len, dim/2]
6     sin = rot_theta.sin()  # [seq_len, dim/2]
7     return cos, sin

在上述代码中,第1行 dim 表示每个 token 对应的维度且必须要为偶数,因为整体是将维度分为两部分来进行计算。第2~4行便是实现式(8)中 $m\theta$ 的计算过程,其中第2行是计算每个维度对应的旋转频率,维度越高频率越低旋转角度也就越小。第5~6行则是计算对应的正弦和余弦值。

如过大家对 Transformer 中的编码方式还有印象的话,会发现上述计算步骤两者相差无几,如下便是 Transformer 中对应的计算过程。

1 pe = torch.zeros(max_len, d_model) 
2 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  
3 div_term = torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))
4 pe[:, 0::2] = torch.sin(position * div_term)  
5 pe[:, 1::2] = torch.cos(position * div_term)

进一步,我们在完成 cos, sin 值的计算以后,便可以完成旋转编码过程,示例代码如下:

 1 def apply_rope_matrix(x, cos, sin):
 2     bsz, seq_len, dim = x.shape
 3     x = x.view(bsz, seq_len, dim // 2, 2)
 4     x1 = x[..., 0]  # [bsz, seq_len, dim // 2]
 5     x2 = x[..., 1]  # 
 6     cos = cos.unsqueeze(0)  # [1, seq_len, dim/2]
 7     sin = sin.unsqueeze(0)
 8     x_rotated_x = x1 * cos - x2 * sin
 9     x_rotated_y = x1 * sin + x2 * cos
10     x_rotated = torch.stack([x_rotated_x, x_rotated_y], dim=-1)
11     return x_rotated.view(bsz, seq_len, dim)

在上述代码中,第3~5行便是将输入在最后一个维度(也就是词向了维度)分成两部分,其中 x[..., 0] 里的 ... 表示前面的所有维度,此处等价于 x[:,:,:,0]。第6~7行是对cos, sin 进行维度扩充。第8~9行便是完成旋转计算过程。第10~11行则是完成拼接,并重塑成 [bsz, seq_len, dim] 的形状。

在这里我们可以看出,对于 Transformer 中的位置编码方法来说,它是将计算得到的正弦和余弦值直接作为绝对位置信息并相加到输入向量中;而对于 RoPE 来说它是将计算得到的正弦和余弦值作为旋转因子,以旋转的方式来给输入融入绝对位置信息。

最后,我们可以通过如下方式来进行使用,示例代码如下:

 1 def rope_matrix(xq, xk):
 2     bsz, seq_len, dim = xq.shape
 3     cos, sin = precompute_freqs_cos_sin(dim, seq_len)
 4     q_rot = apply_rope_matrix(xq, cos, sin)
 5     k_rot = apply_rope_matrix(xk, cos, sin)
 6     print(q_rot)
 7 
 8 if __name__ == '__main__':
 9     dim = 8
10     seq_len = 5
11     xq = torch.randn(2, seq_len, dim)
12     xk = torch.randn(2, seq_len, dim)
13     rope_matrix(xq, xk)

上述代码运行结束后便会得到类似如下结果:

tensor([[[-0.6874,  0.3616,  1.5172, -0.8326, -0.6149,  1.2178,  1.3626, -0.8120],
         [-0.2720,  0.6102,  0.2150, -0.2891, -0.8054, -0.5313,  1.0437, 0.3801],
         [ 2.0354, -0.2829,  0.1590, -2.0516, -0.0071, -0.1854, -0.8610, -0.6257],
         [ 1.0256, -0.6891,  0.8885, -0.5522, -0.3403,  0.6830, -0.9162, -0.1615],
         [ 0.3622,  0.4205, -0.7930,  0.6242,  0.0189, -0.1573,  0.6351, -0.7674]],

        [[ 0.6201, -0.7417, -0.8371,  0.0958, -0.1005, -1.8306,  0.3344, 0.5319],
         [ 1.5699,  0.6796, -1.5662,  0.2982, -0.3243,  1.7357, -0.5256, -1.7252],
         [-1.2461,  1.1995,  0.0210,  0.5950, -0.0178,  0.8088,  0.0184, -0.3366],
         [ 2.2306, -0.4862, -0.5815, -0.8906,  0.6232, -2.3946, -0.3258, 1.1161],
         [-2.4300, -0.6247,  0.1077,  1.9647, -1.3716, -0.7104, -1.0277, -0.2207]]])

上述完整实现代码,可参见

2.5 旋转编码的运用#

对于旋转编码在模型中的实际运用方式与传统的绝对位置编码方式有着很大的差异,主要集中在两个方面:

① 旋转编码在每个注意力层中都会用到

在之前介绍的 Transformer 或 Bert 模型中,对于绝对位置编码来说整个过程都可以看作是输入模型前的 Embedding 操作,是一次性的;而对于 RoPE 来说,它则是在每一个注意力层中都会用到。因为相对位置编码是在计算注意力权重矩阵之前根据 $Q$ 和 $K$ 计算得到的,所以在每一次计算注意力权重矩阵之前都应该使用一次相对位置信息编码,然后再作用 $V$ 得到输出。

图 5. 百川大模型网络结构图
图 5. 百川大模型网络结构图

如图5所示便是百川大模型的网络结构示意图,可以看到在每一个注意力层当中都使用到了 RoPE 位置编码(即图中的 RotaryEmb)。

② 旋转编码只对 $Q$ 和 $K$ 进行编码

因为旋转编码的初始动机便是想让注意力权重 $\text{softmax}(QK^T)$ 具备相对位置的敏感性,而注意力分数 $\text{softmax}(QK^T)$ 只依赖 $Q$ 和 $K$。换句话说,$V$ 是作为一系列信息的集合可以是没有位置信息的,而注意力权重是通过 $(Q, K)$ 计算得到,不同的位置关系计算得到的权重矩阵肯定是不一样的,这些权重乘到 $V$ 上已经可以反映出一个 token 该不该被关注、关注多少,得到的结果也自然是考虑了位置信息后的,所以没必要再额外对 $V$ 做旋转或者改动。

当然,其实在深度学习中很多做法并没有严格理论,如过一定要回答为什么要这样做,那就是:实验表明,然后记住用就行。

若果你只想学会如何使用,那么接下来的内容可以直接略过!

3. 推导过程#

先学会用,再探究为什么。我们在知晓 RoPE 的一个基本原理以及使用方法以后,我们再来大致看一下它到底是怎么来的,作者到底是怎么想的。

3.1 合理假设#

现在,我们假设二维平面中有两个词嵌入表示向量(或再经过一次线性变换后) $x_q$ 和 $x_k$,并且根据式(3)可以得到它们分别在位置 $m$ 和 $n$ 处的编码结果为

$$ \begin{aligned} q_m&=f_q(x_q,m)\\[2ex] k_n&=f_k(x_k,n) \end{aligned}\tag{10} $$

此时, $q_m$ 和 $k_n$ 就带有了位置 $m$ 和 $n$ 处的绝对位置信息

并且,在这样一个体系下我们给出一个合理的初始条件

$$ \begin{aligned} q_0&=f_q(x_q,0)=x_q\\[2ex] k_0&=f_k(x_k,0)=x_k \end{aligned}\tag{11} $$

也就是说,在位置 0 处的向量添加位置信息后依旧是其本身。

因为自注意力机制的核心运算是内积,所以我们希望内积后的结果自然就带有了相对位置信息。我们假设存在一个函数 $g$ ,它定义了上面两个分别由 $f_q$ 和 $f_k$ 计算得到的向量之间的一种运算,且最后的结果就是 $q_m$ 和 $k_n$ 之间的内积,即

$$ q^T_mk_n=\langle f_q(x_m,m),f_k(x_n,n)\rangle= g(x_m,x_n,m-n)\tag{12} $$

式(12)到底什么意思呢?怎么看起来有点怪。

需要特别解释一下的是,式(11)中 $f_q(x_m,m)$ 指代向量 $x_q$ 在 $m$ 处进行绝对位置编码后的结果 $q_m$; $f_k(x_n,m)$ 指代向量 $x_k$ 在 $n$ 处进行绝对位置编码后的结果 $k_n$;$g(x_m,x_n,m-n)$ 中 $x_m$ 指向量 $x_q$ 在位置 $m$ 处时的情况,$x_n$ 指向量 $k_n$ 在 $n$ 处时的情况。

各个向量之间的关系可以通过图6来进行示意。

图 6. 式(12)解释图

因此,式(12)的含义就是说,两个向量 $x_m$ 和 $x_n$ 经过绝对位置信息编码后的结果 $q_m$ 和 $k_n$ ,再经过一次内积运算 $q^T_mk_n$ 计算得到的注意力值便带有相对位置的编码信息,而这两步计算过程可以通过一个函数 $g$ 来进行实现。

所以,只要证明了 $f_q$ 和 $f_k$ 是的确存在的,那么便可以通过像绝对位置编码计算过程一样来实现相对位置信息的编码,即:两个向量分别先被 $f_q$ 和 $f_k$ 作用,然后再进行内积运算。

这里我们需要再次明确作者提出 RoPE 的核心思想:内积后的结果能够带有相对位置编码信息

3.2 证明过程#

在有了上述的背景设定以后,我们使用反证法来进行求解,即先假定等式(12)是成立的,并求解对应的 $f_q$ 和 $f_k$ 。如g果 $f_q$ 和 $f_k$ 有解,也就是说明假设成立,如果无解自然也就不成立。

提示:以下过程会涉及到复数和欧拉公式的相关内容,预备知识可以参见文章 XXXXXXXX

在式(10)中,由于 $q_m$ 和 $k_n$ 均为一个二维向量,根据复数的指数形式表示法可得

$$ \begin{aligned} f_q(x_q,m)&=R_q(x_q,m)e^{i\Theta_q(x_q,m)}\\[2ex] f_k(x_k,n)&=R_k(x_k,n)e^{i\Theta_k(x_k,n)} \end{aligned}\tag{13} $$

其中 $R_q(x_q,m)$ 和 $R_k(x_k,n)$ 分别表示在极坐标系下 $q_m$ 和 $k_n$ 到极点的距离,$\Theta_q(x_q,m)$ 和 $\Theta_k(x_k,n)$ 表示 $q_m$ 和 $k_n$ 分别与极轴之间的夹角。

进一步,根据两个向量的内积可以写成复数相乘取实部的形式,由式(13)可得

$$ \langle f_q(x_q,m),f_k(x_k,n)\rangle=\text{Re}[f_q(x_q,m)f_k^*(x_k,n)]=g(x_q,x_k,m-n)\tag{14} $$

其中 $f_k^*(x_k,n)$ 表示 $f_k(x_k,n)$ 的共轭,$\text{Re}[\cdot]$ 表示取该复数对应的实部 。

为方便化简,我们假设存在一个复数 $\boldsymbol{g}(x_q,x_k,m-n)$ ,使得 $f_q(x_q,m)f_k^*(x_k,n)= \boldsymbol{g}(x_q,x_k,m-n)$ 。

为什么我们需要这样假设?以及能这样假设?

因为我们需要将式(14)中的 $f_q(x_q,m)f_k^*(x_k,n)$ 和 $g(x_q,x_k,m-n)$ 关联起来,所以需要这样假设,并且构造那样的形式。又 $f_q(x_q,m)f_k^*(x_k,n)$ 的结果的确是一个复数,所以自然能这样假设。

由此,可以得到

$$ \boldsymbol{g}(x_q,x_k,m-n)=R_g(x_q,x_k,m-n)e^{i\Theta_g(x_q,x_k,m-n)}\tag{15} $$

其中 $R_g(x_q,x_k,m-n)$ 表示该复数到极点的距离, $\Theta_g(x_q,x_k,m-n)$ 表示复数与极轴之间的夹角。

进一步,根据式(13)有

$$ f_q(x_q,m)f_k^*(x_k,n)=R_q(x_q,m)e^{i\Theta_q(x_q,m)}R_k(x_k,n)e^{-i\Theta_k(x_k,n)}\tag{16} $$

根据式(15)和式(16),由上面的假设可得

$$ R_q(x_q,m)R_k(x_k,n)e^{i\Theta_q(x_q,m)-i\Theta_k(x_k,n)}=R_g(x_q,x_k,m-n)e^{i\Theta_g(x_q,x_k,m-n)}\tag{17} $$

所以根据式(17)中各项相等有

$$ \begin{aligned} R_q(x_q,m)R_k(x_k,n)&=R_g(x_q,x_k,m-n)\\[2ex] \Theta_q(x_q,m)-\Theta_k(x_k,n)&=\Theta_g(x_q,x_k,m-n) \end{aligned}\tag{18} $$

此时,根据式(11) 中的初始条件有

$$ \begin{aligned} q_0=\|q_0\|e^{i\theta_{q_0}}=R_{q_0}(x_q,0)e^{i\Theta_{q_0}(x_q,0)}\\[2ex] k_0=\|k_0\|e^{i\theta_{k_0}}=R_{k_0}(x_k,0)e^{i\Theta_{k_0}(x_k,0)} \end{aligned}\tag{19} $$

其中 $\|q_0\|,\|k_0\|$ 分别表示 $q_0,k_0$ 到极点的距离,$\theta_{q_0},\theta_{k_0}$ 分别表示 $q_0,k_0$ 与极轴的夹角。

因为 $g(x_m,x_n,m-n)$ 考量的是 $x_q$ 中第 $m$ 个位置与 $x_k$ 中第 $n$ 位置之间的相对位置关系(注意看图6),那自然是 $m,n$ 取任何值函数 $g$ 都能得到包含有相对位置信息的结果,那自然也就包括 $m=n$ 的情况。

所以,现在我们为了求解,不妨以该情况为切入点,即当 $m=n$ 时根据式(18)和式(19)有

$$ \begin{aligned} R_q(x_q,m)R_k(x_k,m)&=R_g(x_q,x_k,0)\\[2ex] \Theta_q(x_q,m)-\Theta_k(x_k,m)&=\Theta_g(x_q,x_k,0) \end{aligned}\tag{20} $$

因为式(20)是当 $m=n$ 时,恒成立,所以当 $m=n=0$ 时同样能成立,此时有

$$ \begin{aligned} R_q(x_q,0)R_k(x_k,0)&=R_g(x_q,x_k,0)\\[2ex] \Theta_q(x_q,0)-\Theta_k(x_k,0)&=\Theta_g(x_q,x_k,0) \end{aligned}\tag{21} $$

由式(19)中第1个等式 $q_0$ 乘以 $k_0$ 的共轭,等式左右两边可得

$$ \|q_0\|\|k_0\|e^{(i\theta_{q_0}-i\theta_{k_0})}=R_{q_0}(x_q,0)R_{k_0}(x_k,0)e^{(i\Theta_{q_0}(x_q,0)-i\Theta_{k_0}(x_k,0))}\tag{22} $$

进一步根据式(22)可得

$$ \begin{aligned} \|q_0\|\|k_0\|&=R_{q_0}(x_q,0)R_{k_0}(x_k,0)\\[2ex] \theta_{q_0}-\theta_{k_0}&=\Theta_{q_0}(x_q,0)-\Theta_{k_0}(x_k,0) \end{aligned}\tag{23} $$

这里需要注意的一点是,式(23)中的下标 $q_0,k_0$ 其实就是式(21)中的 $q,k$ ,因为都表示在位置 0 处的相关信息,含义相同只是标识不一样。

所以由式(21)和式(23)可得

$$ \begin{aligned} R_q(x_q,0)R_k(x_k,0)&=R_g(x_q,x_k,0)=\|q\|\|k\|\\[2ex] \Theta_q(x_q,0)-\Theta_k(x_k,0)&=\Theta_g(x_q,x_k,0)=\theta_{q}-\theta_{k} \end{aligned}\tag{24} $$

根据式(24)可得

$$ R_q(x_q,0)=\|q\|;\quad R_k(x_k,0)=\|k\|\tag{25} $$

根据式(20)和式(24)可得

$$ \begin{aligned} \Theta_q(x_q,m)-\Theta_k(x_k,m)=\theta_q-\theta_k\\[2ex] \Theta_q(x_q,m)-\theta_q=\Theta_k(x_k,m)-\theta_k \end{aligned}\tag{26} $$

这里我们再次回顾一下,$\Theta_q(x_q,m)$ 表示向量 $x_q$ 融入位置 $m$ 处的信息变为 $q_m$ 后,$q_m$ 与极轴之间的夹角;$\theta_q$ 在这里的含义是向量 $x_q$ 融入位置 $0$ 处的信息变为 $q_0$ 后,$q_0$(也即 $x_q$ )与极轴之间的夹角 。所以,$\Theta_q(x_q,m)-\theta_q$ 应该是一个只与 $m$ 相关,与 $x_q$ 无关的函数,故可以记为 $\phi(m)$,即

$$ \Theta_q(x_q,m)=\theta_q+\phi(m)\\[2ex] \Theta_k(x_k,m)=\theta_k+\phi(m)\tag{27} $$

进一步,令 $n=m-1$ ,并代入式(18)有

$$ \begin{aligned} \Theta_q(x_q,m)-\Theta_k(x_k,m-1)=\Theta_g(x_q,x_k,1) \end{aligned} \tag{28} $$

将式(27)代入式(28)有

$$ \phi(m)-\phi(m-1)=\Theta_g(x_q,x_k,1)+\theta_k-\theta_q \tag{29} $$

可以看出 $\{\phi(m)\}$ 是一个以 $\theta=\Theta_g(x_q,x_k,1)+\theta_k-\theta_q$ 为公差,$\gamma$ 为首项的等差数列,即

$$ \phi(m)=m\theta+\gamma\tag{30} $$

根据式(13)、式(25)、式(27)和式(30)有

$$ \begin{aligned} f_q(x_q,m)&=\|q\|e^{i\Theta_q(x_q,m)}=\|q\|e^{i(\theta_q+m\theta+\gamma)}\\[2ex] f_k(x_k,n)&=\|k\|e^{i\Theta_k(x_k,n)}=\|k\|e^{i(\theta_k+n\theta+\gamma)}\end{aligned}\tag{31} $$

这里需要注意的是,式(31)中的 $q$ 和 $k$ 依旧指的是式(11)中的 $q_0$ 和 $k_0$ ,并且此时依旧是利用 $m=n=0$ 这样的边界条件在进行求解。下面,我们换成更一般的形式,定义

$$ \begin{aligned} q&=f_q(x_m,0)=W_qx_m\\[2ex] k&=f_k(x_n,0)=W_kx_n \end{aligned}\tag{32} $$

最后,我们简单地令式(30)中 $\gamma=0$ 可得

$$ \begin{aligned} f_q(x_m,m)=(W_qx_m)e^{im\theta}\\[2ex] f_k(x_n,n)=(W_kx_n)e^{in\theta} \end{aligned}\tag{33} $$

根据式(13)便可得到式(8)中的表达式。

由于内积满足线性叠加(即对应维度相乘再求和)的关系,因此任意偶数维度的 RoPE,都可以表示为二维形式的拼接,即

$$ \begin{bmatrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \cdots & 0 & 0 \\ 0 & 0 & \sin m\theta_1 & \cos m\theta_1 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{(d/2)-1} & -\sin m\theta_{(d/2)-1} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{(d/2)-1} & \cos m\theta_{(d/2)-1} \\ \end{bmatrix} \begin{bmatrix} W^{(00)}_{\{q,k\}}&W^{(01)}_{\{q,k\}}&\cdots &W^{(0[d-1])}_{\{q,k\}}\\ W^{(10)}_{\{q,k\}}&W^{(11)}_{\{q,k\}}&\cdots &W^{(1[d-1])}_{\{q,k\}}\\ W^{(20)}_{\{q,k\}}&W^{(21)}_{\{q,k\}}&\cdots &W^{(2[d-1])}_{\{q,k\}}\\ \vdots&\vdots&\vdots&\vdots\\ W^{([d-1]0)}_{\{q,k\}}&W^{([d-1]1)}_{\{q,k\}}&\cdots &W^{([d-1][d-1])}_{\{q,k\}} \end{bmatrix} \begin{bmatrix} x^{(0)}_m\\x^{(1)}_m\\x^{(2)}_m\\x^{(3)}_m\\ \vdots\\ x^{(d-1)}_m \end{bmatrix}\tag{34} $$

当然,在代码实现的时候为了更加高效,我们可以采用如图4所示的方式进行。

3.3 线性注意力#

在介绍完旋转位置编码的原理以后我们再来看 RoPE 为什么能够满足线性注意力计算。

根据式(2)和式(4)可知,我们可以将自注意力机制的计算过程通过一共更加泛化的方式进行表示,如下所示

$$ \text{Attention}(Q,K,V)_m=\sum_{n=1}^N\frac{\text{sim}(q_m,k_n)}{\sum_{n=1}^N\text{sim}(q_m,k_n)}v_n\tag{35} $$

其中 $\text{sim}(q_m,k_n)=\exp(q^Tk_n/\sqrt{d})$.

在这样的情境下,式(35)完成整个自注意力的计算过程,其时间复杂度便是 $O(n^2d)$。

在论文[6]中,Katharopoulos 等人提出了线性注意力计算框架

$$ \text{Attention}(Q,K,V)_m=\sum_{n=1}^N\frac{\phi(q_m)^T\varphi (k_n)}{\sum_{n=1}^N\phi(q_m)^T\ \varphi (k_n)}v_n\tag{36} $$

其中 $\phi(\cdot), \varphi (\cdot)$ 通常为非负,因为我们要确保注意力权重为非负。

此时可以看出,在式(36)的计算框架下,我们便可以个实现 $Q(K^TV)$ 的计算顺序,以达到线性复杂度的目的。

进一步,在[6]中,作者令 $\phi(x)= \varphi (x)=\text{elu}(x)+1$ 来完成了自注意力机制的计算过程,其中 $\text{elu}(x)$ 的定义如下

$$ \text{ELU}(x) = \begin{cases} x, & \text{if } x > 0 \\[2ex] \alpha \left( e^{x} - 1 \right), & \text{if } x \le 0 \end{cases} \tag{37} $$

其中 $\alpha > 0$ 是一个超参数,通常 $\alpha = 1.0$ 。

对于 RoPE 来说,因为它只是通过旋转改变了向量的向量的角度(并没有改变模长),因此我们可以将其与式(36)进行结合,如下

$$ \text{Attention}(Q,K,V)_m=\sum_{n=1}^N\frac{\left(R^d_{\Theta,m}\phi(q_m)\right)^T\left(R^d_{\Theta,n}\varphi (k_n)\right)}{\sum_{n=1}^N\phi(q_m)^T\ \varphi (k_n)}v_n\tag{38} $$

此时可以看出,使用 RoPE 的自注意力计算过程同样是可以满足线性时间复杂度的。

不过此时需要注意的是,因为向量 $\phi(q_m),\varphi (k_n)$ 进行了旋转,所以并不能保证计算得到的权重是非负的。可能也正式因为这样的原因,作者在运用 RoPE 的过程中还是通过 Softmax 进行了归一化。

4. 总结#

在这篇论文中,作者以相对位置编码和线性注意力框架为动机,通过合理的假设得到了旋转注意力编码 RoPE,并将其运用在了 Transfromer 中,取名 RoFormer。从现在的实际情况来看,各个大模型中仅仅只是将 RoPE 作为了一种新的位置编码方式来进行使用,并且依旧是进行了 Softmax 操作,即时间复杂度仍然为 $O(n^2d)$ 。

总结,本次内容就到此结束,更多内容可参见作者原文。总的来说,对于 RoPE 我们只需其思想和用法即可,至于怎么来的过程也无需太多关注,比较枯燥。

引用#

[1] Haviv A, Ram O, Press O, et al. Transformer language models without positional encodings still learn positional information[J]. arXiv preprint arXiv:2203.16634, 2022.

[2] Jianlin Su, Yu Lu, Shengfeng Pan, et al. RoFormer: Enhanced Transformer with Rotary Position Embedding[J]. arXiv preprint arXiv:2104.09864

[3] 苏剑林,线性Attention的探索:Attention必须有个Softmax吗? 2020,7. https://spaces.ac.cn/archives/7546

[4] 苏剑林,Transformer升级之路:2、博采众长的旋转式位置编码. 2021, 3. https://spaces.ac.cn/archives/8265

[5] https://github.com/meta-llama/llama/blob/llama_v1/llama/model.py

[6] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156–5165. PMLR, 2020.

阅读 --

复数及欧拉公式介绍

本文是旋转位置编码 RoPE 系列的数学前置篇,系统讲清复数与欧拉公式这两个 RoPE 推导必备工具。从复数 z=a+bi 的实部虚部、复平面坐标表示、共轭复数 inner product 写法讲起,逐步推导出欧拉公式 …