RoPE / RoFormer: Enhanced Transformer with Rotary Position Embedding

本文是论文 RoFormer: Enhanced Transformer with Rotary Position Embedding 的阅读笔记和个人理解.

Recommended

与往常不同, 在本文开头先给出推荐.

首先苏神博客看RoPE整个系列的文章(RoPE的论文就是出自这些博客的集成):

• 让研究人员绞尽脑汁的Transformer位置编码 - 科学空间|Scientific Spaces.
• Transformer升级之路：2、博采众长的旋转式位置编码 - 科学空间|Scientific Spaces.
• 标签 rope 下的文章 - 科学空间|Scientific Spaces.

本文中穿插的内容也有少量由苏神博客中延伸.

如果想看比较简单一点的, 强烈推荐直接看抱抱脸推出的Blog设计位置编码, 从一个比较简单易懂的角度解释了RoPE的由来.

代码方面, 也推荐看看LLaMA对RoPE的实现:

• llama3/llama/model.py at main · meta-llama/llama3 · GitHub.

如果希望看形象一点的, 有图描述的, 推荐看:

• 避开复数推导，我们还可以怎么理解RoPE？ - 猛猿的文章 - 知乎.

本文是对RoPE论文的记录.

Background and Related Work

位置编码的主要作用是为了保证同样的Token输入到模型中是可以区分的, 从而可以编码它们的位置关系. 一个显然的例子, “我吃饭”和”饭吃我”明显因主宾语的位置不同而具有不同的含义.

熟悉的Position Embedding的本小节可以跳过.

Preliminary

对于长度为$N$ 的句子${\mathbb{S}}_N$ , 将其中所有Token组成的Embedding序列${\mathbb{E}}_N$, ${\mathit{x}}_i\in {\mathbb{R}}^d$ 为${\mathbb{E}}_N$ 中第$i$ 个位置上的$d$ 维Embedding.

当Transformer中的Self - Attention在忽略位置信息的情况下, 其Attention Score的计算方式如下:

$$\begin{array}{rl}{\mathit{q}}_m& =f_q({\mathit{x}}_m,m)\\ {\mathit{k}}_n& =f_k({\mathit{x}}_n,n)\\ {\mathit{v}}_n& =f_v({\mathit{x}}_n,n)\end{array}$$

其中${\mathit{q}}_m$, ${\mathit{k}}_n,{\mathit{v}}_n$ 分别为第$m,n$ 位置通过$f_q,f_k,f_v$ 获得的Query, Key, Value.

接着计算Scaled Dot-Product Attention:

$$\begin{array}{rl}{a}_{m,n}& =\frac{\exp \left(\frac{{\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n}{\sqrt{d}}\right)}{\sum _{j=1}^N\exp \left(\frac{{\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_j}{\sqrt{d}}\right)}\\ {\mathbf{o}}_m& =\sum _{n=1}^N{a}_{m,n}{\mathit{v}}_n\end{array}$$

$a_{m,n}$ 为$m,n$ Token之间的Attention Score, ${\mathbf{o}}_m$为第$m$ 个Token的计算完Self - Attention后的表示.

Absolute Position Embedding

最经典的$f_{t:t\in \{q,k,v\}}(\cdot )$ 的方法是直接将Embedding ${\mathit{x}}_i$ 与位置信息${\mathit{p}}_i$ 直接相加, 再通过一个线性投影${\mathit{W}}_{t:t\in \{q,k,v\}}$ 直接得到:

$$f_{t:t\in \{q,k,v\}}({\mathit{x}}_i,i):={\mathit{W}}_{t:t\in \{q,k,v\}}({\mathit{x}}_i+{\mathit{p}}_i)$$

而在Transformer的原论文中, 作者通过正余弦函数来生成${\mathit{p}}_i$:

$$\{\begin{array}{ll}{\mathit{p}}_{i,2t}& =\sin \left(k/{10000}^{2t/d}\right)\\ {\mathit{p}}_{i,2t+1}& =\cos \left(k/{10000}^{2t/d}\right)\end{array}$$

${\mathit{p}}_{i,2t},{\mathit{p}}_{i,2t+1}$ 分别为${\mathit{p}}_i\in {\mathbb{R}}^d$ 的奇数维和偶数维.

当然, 在BERT, ViT等论文中也有直接使用Learnable的绝对位置编码, 直接将绝对位置设定为一个可学习的参数让网络自己学.

Relative Position Embedding

由于绝对位置编码不能编码两个Token之间的相对位置关系, 一些工作中也通过相对位置的建模, 将相对位置信息作为位置编码:

$$\begin{array}{rl}{f}_q\left({\mathit{x}}_m\right)& :={\mathit{W}}_q{\mathit{x}}_m\\ f_k({\mathit{x}}_n,n)& :={\mathit{W}}_k({\mathit{x}}_n+{\tilde{\mathit{p}}}_r^k)\\ f_v({\mathit{x}}_n,n)& :={\mathit{W}}_v({\mathit{x}}_n+{\tilde{\mathit{p}}}_r^v)\end{array}$$

${\tilde{\mathit{p}}}_r^k,{\tilde{\mathit{p}}}_r^v\in {\mathbb{R}}^d$ 为可学习的相对位置编码.

其中, $r$ 为$m,n$ 之间的相对位置, 用$r=\text{clip}(m-n,r_{min},r_{max})$ 得到. 这假设了超出一定距离限制的位置信息是不太重要的. 此时直接采用最远距离的相对位置编码即可.

一些工作(TransformerXL / XLNet)直接沿用这个形式将它们按项拆分:

$${\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n={\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+{\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{p}}_n+{\mathit{p}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+{\mathit{p}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{p}}_n$$

进一步的, 由于投影矩阵 ${\mathit{W}}_k$ 有时候会对Content-based Key ${\mathit{x}}_n$ 做投影, 有时候会对Position-based Key ${\mathit{p}}_n$ 做投影, 这显然不合理对吧? 所以可以额外加一个矩阵${\tilde{\mathit{W}}}_k$ 让它来编码${\mathit{p}}_n$:

$${\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n={\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+{\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\tilde{\mathit{W}}}_k{\tilde{\mathit{p}}}_{m-n}+{\mathbf{u}}^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+{\mathbf{v}}^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\tilde{\mathit{W}}}_k{\tilde{\mathit{p}}}_{m-n}$$

这玩意看起来太麻烦了, 我们要的其实不就是一项位置信息和一项内容信息吗? 于是一些工作(T5)直接将位置信息和内容信息解耦, 直接将相对位置变为一个Learnable bias term $b_{i,j}$:

$${\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n={\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+b_{i,j}$$

用两个Projection进一步分别编码${\mathit{p}}_m,{\mathit{p}}_n$, 即相对位置, 似乎是更合理的:

$${\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n={\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+{\mathit{p}}_m^{\mathrm{\top }}{\mathbf{U}}_q^{\mathrm{\top }}{\mathbf{U}}_k{\mathit{p}}_n+b_{i,j}$$

也有一些工作(DeBERTa)基于绝对位置的形式, 认为简单的将绝对位置${\mathit{p}}_n,{\mathit{p}}_m$ 替换为相对位置${\tilde{\mathit{p}}}_{m-n}$就可以:

$${\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n={\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n+{\mathit{x}}_m^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\tilde{\mathit{p}}}_{m-n}+{\tilde{\mathit{p}}}_{m-n}^{\mathrm{\top }}{\mathit{W}}_q^{\mathrm{\top }}{\mathit{W}}_k{\mathit{x}}_n$$

RoPE

用下图可以概括RoPE的核心思想:

为什么要叫Rotary Position Embedding呢? 正如它的名字一样, 它是通过旋转来区分不同位置的. 上图举了一个$d=2$ 时的简单例子, 对于绝对位置$m$, RoPE将向量在极坐标系中通过旋转能够代表$m$ 的角度来表征它们的位置.

但RoPE是以表征相对位置闻名的, 因为它可以借助一些巧妙的性质, 使得计算Self-Attention的内积时能将相对位置信息注入到其中. 下面就来看看RoPE是怎么做到这一点的.

Formulation

结合前面的分析可知, 想要融入Token之间的相对位置关系, 只需要将${\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n$ 建模为能注入相对位置$m-n$ 的函数$g$:

$$⟨f_q({\mathit{x}}_m,m),f_k({\mathit{x}}_n,n)⟩=g({\mathit{x}}_m,{\mathit{x}}_n,m-n)$$

所以目标就是找到满足$g$ 性质的$f_q({\mathit{x}}_m,m),f_k({\mathit{x}}_n,n)$.

Rotary Position Embedding

Derivation of RoPE under 2D

在本节中先讨论简单的二维情况.

SDPA(Scaled Dot-Product Attention)是两向量的内积.

如果实数运算规则不好找到函数$g$, 有没有其他的运算规则能帮助我们求解? 比如思考实数内积与复数内积有什么关系吗? 刚刚好, 两向量的内积等于一个复数与另一个复数的共轭的乘积的实部.

例如, 对于任意两个实数向量$\mathbf{a}=(a_1,a_2),\mathbf{b}=(b_1,b_2)$, 它们的内积定义为:

$$⟨\mathbf{a}\cdot \mathbf{b}⟩=a_1{b}_1+a_2{b}_2$$

当我们将$\mathbf{a},\mathbf{b}$ 看成复数时, 有复数$z_a=a_1+a_{2}i$, 其共轭为$z_a^{\ast }=a_1-a_{2}i$, 同理对于复数$z_b=b_1+b_{2}i$ 有共轭$z_b^{\ast }=b_1-b_{2}i$. 什么叫”求一个复数与另一个复数的共轭的乘积的实部”? 计算$z_a,z_b^{\ast }$ 之间的内积:

$$\begin{array}{rl}{z}_a\cdot z_b^{\ast }& =(a_1+a_{2}i)(b_1-b_{2}i)\\ & =a_1{b}_1+a_1(-b_{2}i)+a_{2}i{b}_1+a_{2}i(-b_{2}i)\\ & =a_1{b}_1+a_2{b}_2+i(a_2{b}_1-a_1{b}_2)\end{array}$$

实部$a_1{b}_1+a_2{b}_2$ 恰好就是$\mathbf{a},\mathbf{b}$ 在实数域内积的值.

好, 有了上面的性质, 下面开始今天的正题. 先往上小节说的Formulation上靠拢呗.

如果想要表征绝对位置信息, 需要使得${\mathit{q}}_m,{\mathit{k}}_n$ 仅依赖于它们的绝对位置$m,n$:

$$\begin{array}{rl}{\mathit{q}}_m& =f_q({\mathit{x}}_q,m)\\ {\mathit{k}}_n& =f_k({\mathit{x}}_k,n)\end{array}$$

利用实数内积与复数内积的关联, 我们假设有利用复数运算的函数$g$ 能够在二者做内积的时候编码它们的相对位置关系, 使得其内积依赖于相对位置$m-n$, 以此在内积上表征二者的相对位置信息:

$${\mathit{q}}_m^{⊺}{\mathit{k}}_n=⟨f_q({\mathit{x}}_m,m),f_k({\mathit{x}}_n,n)⟩=g({\mathit{x}}_m,{\mathit{x}}_n,n-m)$$

同时, 这种编码方式还应该保证在初始状态下不编码位置信息:

$$\begin{array}{rl}\mathit{q}& =f_q({\mathit{x}}_q,0)\\ \mathit{k}& =f_k({\mathit{x}}_k,0)\end{array}$$

首先, 用复数的指数形式(都学过极坐标系中表示复数吧), 将$f_q({\mathit{x}}_m,m),f_k({\mathit{x}}_n,n)$ 表征为:

$$\begin{array}{rl}{f}_q({\mathit{x}}_q,m)& =R_q({\mathit{x}}_q,m)e^{i{\mathrm{\Theta }}_q({\mathit{x}}_q,m)}\\ f_k({\mathit{x}}_k,n)& =R_k({\mathit{x}}_k,n)e^{i{\mathrm{\Theta }}_k({\mathit{x}}_k,n)}\\ g({\mathit{x}}_q,{\mathit{x}}_k,n-m)& =R_g({\mathit{x}}_q,{\mathit{x}}_k,n-m)e^{i{\mathrm{\Theta }}_g({\mathit{x}}_q,{\mathit{x}}_k,n-m)}\end{array}$$

其中, $R_f,R_g,{\mathrm{\Theta }}_f,{\mathrm{\Theta }}_g$ 分别为$f_{\{q,k\}}$ 和$g$ 的幅度(模)和角度. $f$ 的意思是$f_{\{q,k\}}$, 代表$q,k$.

不难发现, 代入到${\mathit{q}}_m,{\mathit{k}}_n$ 中, 当二者做内积时, 得到幅度$R$ 和角度$\mathrm{\Theta }$ 之间的关系:

$$\begin{array}{rl}{R}_q({\mathit{x}}_q,m)R_k({\mathit{x}}_k,n)& =R_g({\mathit{x}}_q,{\mathit{x}}_k,n-m)\\ {\mathrm{\Theta }}_k({\mathit{x}}_k,n)-{\mathrm{\Theta }}_q({\mathit{x}}_q,m)& ={\mathrm{\Theta }}_g({\mathit{x}}_q,{\mathit{x}}_k,n-m)\end{array}$$

此时, 它在初始条件(看成是不表征位置信息)的情况下, 应满足:

$$\begin{array}{rl}\mathit{q}& =\Vert \mathit{q}\Vert e^{i{\theta }_q}=R_q({\mathit{x}}_q,0)e^{i{\mathrm{\Theta }}_q({\mathit{x}}_q,0)}\\ \mathit{k}& =\Vert \mathit{k}\Vert e^{i{\theta }_k}=R_k({\mathit{x}}_k,0)e^{i{\mathrm{\Theta }}_k({\mathit{x}}_k,0)}\end{array}$$

其中, $\Vert q\Vert ,\Vert k\Vert$ 和${\theta }_q,{\theta }_k$ 分别为$\mathit{q},\mathit{k}$ 在二维平面上的幅值(模)和角度.

接着将相同的位置$m=n$ 带入到幅度$R$ 和角度$\mathrm{\Theta }$ 的关系式中, 分别得到:

$$\begin{array}{rl}{R}_q({\mathit{x}}_q,m)R_k({\mathit{x}}_k,m)& =R_g({\mathit{x}}_q,{\mathit{x}}_k,0)=R_q({\mathit{x}}_q,0)R_k({\mathit{x}}_k,0)=\Vert \mathit{q}\Vert \Vert \mathit{k}\Vert \\ {\mathrm{\Theta }}_k({\mathit{x}}_k,m)-{\mathrm{\Theta }}_q({\mathit{x}}_q,m)& ={\mathrm{\Theta }}_g({\mathit{x}}_q,{\mathit{x}}_k,0)={\mathrm{\Theta }}_k({\mathit{x}}_k,0)-{\mathrm{\Theta }}_q({\mathit{x}}_q,0)={\theta }_k-{\theta }_q\end{array}$$

此时, 幅值$R_f$ 的一个显然解为:

$$\begin{array}{rl}{R}_q({\mathit{x}}_q,m)& =R_q({\mathit{x}}_q,0)=\Vert \mathit{q}\Vert \\ R_k({\mathit{x}}_k,n)& =R_k({\mathit{x}}_k,0)=\Vert \mathit{k}\Vert \\ R_g({\mathit{x}}_q,{\mathit{x}}_k,n-m)& =R_g({\mathit{x}}_q,{\mathit{x}}_k,0)=\Vert \mathit{q}\Vert \Vert \mathit{k}\Vert \end{array}$$

它显然不依赖位置信息, 即满足$\mathit{q}=f_q({\mathit{x}}_q,0),\mathit{k}=f_k({\mathit{x}}_k,0)$.

那么角度$\mathrm{\Theta }$ 的解呢?

当位置相同时, 不应该编码任何位置信息. 令${\mathrm{\Theta }}_f:={\mathrm{\Theta }}_q={\mathrm{\Theta }}_k$, 根据${\mathrm{\Theta }}_q({\mathit{x}}_q,m)-{\theta }_q={\mathrm{\Theta }}_k({\mathit{x}}_k,m)-{\theta }_k$, 此时明显两侧是相同的函数. 由此, 可以推断出${\mathrm{\Theta }}_f({\mathit{x}}_{\{q,k\}},m)-{\theta }_{\{q,k\}}$ 是与位置$m$ 相关, 与Word Embedding ${\mathit{x}}_{\{q,k\}}$ 无关的函数, 可设${\mathrm{\Theta }}_f({\mathit{x}}_{\{q,k\}},m)-{\theta }_{\{q,k\}}=\varphi (m)$, 得到:

$${\mathrm{\Theta }}_f({\mathit{x}}_{\{q,k\}},m)=\varphi (m)+{\theta }_{\{q,k\}}$$

$n=m+1$, 将上式代入到角度关系式当中, 可以得到:

$$\varphi (m+1)-\varphi (m)={\mathrm{\Theta }}_g({\mathit{x}}_q,{\mathit{x}}_k,1)+{\theta }_q-{\theta }_k$$

上式中右侧的常数项与绝对位置$m$ 是无关的, 所以令上式右端整体为$\theta$, $\varphi (m)$ 就是一个等差数列:

$$\varphi (m)=m\theta +\gamma$$

其中$\gamma \in \mathbb{R}$ 为常数(首项), $\theta$ 非零, 绝对位置$m$ 在这个式子中也是公差.

通过对幅值$R_f$ 和角度${\mathrm{\Theta }}_f$ 的求解, 整理一下上面的式子, 可以得到:

$$\begin{array}{rl}{f}_q({\mathit{x}}_q,m)& =\Vert \mathit{q}\Vert e^{i({\theta }_q+m\theta +\gamma )}=\mathit{q}{e}^{i(m\theta +\gamma )}\\ f_k({\mathit{x}}_k,n)& =\Vert \mathit{k}\Vert e^{i({\theta }_k+n\theta +\gamma )}=\mathit{k}{e}^{i(n\theta +\gamma )}\end{array}$$

通常$\mathit{q}=f_q({\mathit{x}}_m,0),\mathit{k}=f_k({\mathit{x}}_n,0)$ 是直接由一层投影得到的:

$$\begin{array}{rl}\mathit{q}=f_q({\mathit{x}}_m,0)& ={\mathit{W}}_q{\mathit{x}}_n\\ \mathit{k}=f_k({\mathit{x}}_n,0)& ={\mathit{W}}_k{\mathit{x}}_n\end{array}$$

方便起见, 直接设定$\gamma =0$, 因此得到最终的RoPE形式:

$$\begin{array}{rl}{f}_q({\mathit{x}}_m,m)& =({\mathit{W}}_q{\mathit{x}}_m)e^{im\theta }\\ f_k({\mathit{x}}_n,n)& =({\mathit{W}}_k{\mathit{x}}_n)e^{in\theta }\end{array}$$

所以, 我们在二维情况下, 最后找到的函数$g$ 为:

$$g({\mathit{x}}_m,{\mathit{x}}_n,m-n)=\mathrm{Re}[({\mathit{W}}_q{\mathit{x}}_m)({\mathit{W}}_k{\mathit{x}}_n{)}^{\ast }{e}^{i(m-n)\theta }]$$

其中$\text{Re}(\cdot )$ 为取复数的实部, $({\mathit{W}}_k{\mathit{x}}_n{)}^{\ast }$ 为$({\mathit{W}}_k{\mathit{x}}_n)$ 的共轭复数, $\theta \in \mathbb{R}$ 为预设好的非零常数(基波长, 沿用Transformer的原论文, 经常设为10000, 底数的选择可以参考这里).

进一步写出二维情况下$f_{q,k}$ 的表达式:

$$\begin{array}{rl}{f}_{\{q,k\}}({\mathit{x}}_m,m)& =\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{cc}{W}_{\{q,k\}}^{(11)}& W_{\{q,k\}}^{(12)}\\ W_{\{q,k\}}^{(21)}& W_{\{q,k\}}^{(22)}\end{array}\right)\left(\begin{array}{c}{x}_m^{(1)}\\ x_m^{(2)}\end{array}\right)\\ & =\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{c}\{q,k{\}}_m^{(1)}\\ \{q,k{\}}_m^{(2)}\end{array}\right)\end{array}$$

它就是$\{\mathit{q}\mathbf{,}\mathit{k}{\}}_m$ 乘了一个旋转矩阵, 对应着空间中的向量旋转操作, 这才是RoPE(Rotary Position Embedding)名字的由来.

General form

更为普遍的, 将上节得到的二维情况扩展到任意偶数维度$d$. 将RoPE应用于${\mathit{x}}_i\in {\mathbb{R}}^d$:

$$f_{\{q,k\}}({\mathit{x}}_m,m)={\mathit{R}}_{\mathrm{\Theta },m}^d{\mathit{W}}_{\{q,k\}}{\mathit{x}}_m$$

非常自然的思路就是将$d$ 维拆成$d/2$ 个的二维子空间, 所以${\mathit{R}}_{\mathrm{\Theta },m}^d$ 为:

$${\mathit{R}}_{\mathrm{\Theta },m}^d=\left(\begin{array}{ccccccc}\cos m{\theta }_1& -\sin m{\theta }_1& 0& 0& \cdots & 0& 0\\ \sin m{\theta }_1& \cos m{\theta }_1& 0& 0& \cdots & 0& 0\\ 0& 0& \cos m{\theta }_2& -\sin m{\theta }_2& \cdots & 0& 0\\ 0& 0& \sin m{\theta }_2& \cos m{\theta }_2& \cdots & 0& 0\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ 0& 0& 0& 0& \cdots & \cos m{\theta }_{d/2}& -\sin m{\theta }_{d/2}\\ 0& 0& 0& 0& \cdots & \sin m{\theta }_{d/2}& \cos m{\theta }_{d/2}\end{array}\right)$$

其中, $\mathrm{\Theta }=\{{\theta }_i={10000}^{-2(i-1)/d},i\in [1,2,\dots ,d/2]\}$. 所以$i$ 越大, ${\theta }_i$ 越小, 即下标越大的维度旋转速度越慢.

这篇文章提供了一个非常有意思的角度, 不同角速度的维度就像时针分针秒针一样, 通过不同的快慢角速度来组合表征一个位置信息.

将RoPE直接作用于Self-Attention, 可得:

$${\mathit{q}}_m^{⊺}{\mathit{k}}_n=({\mathit{R}}_{\mathrm{\Theta },m}^d{\mathit{W}}_q{\mathit{x}}_m{)}^{⊺}({\mathit{R}}_{\mathrm{\Theta },n}^d{\mathit{W}}_k{\mathit{x}}_n)={\mathit{x}}^{⊺}{\mathit{W}}_q{R}_{\mathrm{\Theta },n-m}^d{\mathit{W}}_k{\mathit{x}}_n$$

其中, ${\mathit{R}}_{\mathrm{\Theta },n-m}^d={\left({\mathit{R}}_{\mathrm{\Theta },m}^d\right)}^{\mathrm{\top }}{\mathit{R}}_{\mathrm{\Theta },n}^d$. 由于${\mathit{R}}_{\mathrm{\Theta }}^d$ 是一个正交矩阵, 所以它在编码位置信息的时候是相对稳定的.

RoPE相较于之前位置编码的最大不同, 在于它是乘性的, 而不是加性的, 因此它在编码的相对位置信息可以天然的融入到Self-Attention的内积中.

Properties of RoPE

RoPE有一些优良的性质.

首先, RoPE具有长程衰减, 和最早版本的Transformer一样, RoPE设定${\theta }_i={10000}^{-2i/d}$, 底数$\theta$ 越大长程衰减越慢.

其次, RoPE在Linear Attention当中可以使得Self-Attention被更一般的形式重写.

例如, 常规的SDPA可以写成:

$$\mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}{)}_m=\frac{\sum _{n=1}^N\mathrm{sim}({\mathit{q}}_m,{\mathit{k}}_n){\mathit{v}}_n}{\sum _{n=1}^N\mathrm{sim}({\mathit{q}}_m,{\mathit{k}}_n)}$$

Self-Attention中采用的是$\mathrm{sim}({\mathit{q}}_m,{\mathit{k}}_n)=\exp \left({\mathit{q}}_m^{\mathrm{\top }}{\mathit{k}}_n/\sqrt{d}\right)$, 它具有$\mathbb{O}(N^2)$ 的复杂度.

Linear Attention的复杂度是比SDPA更低的线性复杂度, 形式如下:

$$\mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}{)}_m=\frac{\sum _{n=1}^N\varphi ({\mathit{q}}_m{)}^{⊺}\phi ({\mathit{k}}_n){\mathit{v}}_n}{\sum _{n=1}^N\varphi ({\mathit{q}}_m{)}^{⊺}\phi ({\mathit{k}}_n)}$$

其中$\varphi (\cdot ),\phi (\cdot )$ 通常是非负函数, 例如$\varphi (x)=\phi (x)=\text{elu}(x)+1$, 或者$\varphi \left({\mathit{q}}_i\right)=\mathrm{softmax}\left({\mathit{q}}_i\right),\phi \left({\mathit{k}}_j\right)=\exp \left({\mathit{k}}_j\right)$.

由于RoPE的存在, 可以使得Hidden State的范数不变, 因此可以将RoPE直接乘, 从而无缝集成到Linear Attention中, 从而不增加Linear Attention的复杂度:

$$\mathrm{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}{)}_m=\frac{\sum _{n=1}^N({\mathit{R}}_{\mathrm{\Theta },m}^d\varphi ({\mathit{q}}_m){)}^{⊺}({\mathit{R}}_{\mathrm{\Theta },n}^d\phi ({\mathit{k}}_n)){\mathit{v}}_n}{\sum _{n=1}^N\varphi ({\mathit{q}}_m{)}^{⊺}\phi ({\mathit{k}}_n)}$$

Theoretical Explanation

Computational efficient realization of rotary matrix multiplication

由于${\mathit{R}}_{\mathrm{\Theta },m}^d$ 是一个稀疏矩阵, 所以它的计算效率很低. 下式可以化简计算:

$${\mathit{R}}_{\mathrm{\Theta },m}^d\mathit{x}=\left(\begin{array}{c}{x}_1\\ x_2\\ x_3\\ x_4\\ ⋮\\ x_{d-1}\\ x_d\end{array}\right)\otimes \left(\begin{array}{c}\cos m{\theta }_1\\ \cos m{\theta }_1\\ \cos m{\theta }_2\\ \cos m{\theta }_2\\ ⋮\\ \cos m{\theta }_{d/2}\\ \cos m{\theta }_{d/2}\end{array}\right)+\left(\begin{array}{c}-x_2\\ x_1\\ -x_4\\ x_3\\ ⋮\\ -x_d\\ x_{d-1}\end{array}\right)\otimes \left(\begin{array}{c}\sin m{\theta }_1\\ \sin m{\theta }_1\\ \sin m{\theta }_2\\ \sin m{\theta }_2\\ ⋮\\ \sin m{\theta }_{d/2}\\ \sin m{\theta }_{d/2}\end{array}\right)$$

其中$\otimes$ 为逐元素点乘.

Long-term decay of RoPE

在RoPE的作用下, 将向量$\mathit{q}={\mathit{W}}_q{\mathit{x}}_m,\mathit{k}={\mathit{W}}_k{\mathit{x}}_n$ 两维两维的分组成对, 其内积可以写成:

$$({\mathit{R}}_{\mathrm{\Theta },m}^d{\mathit{W}}_q{\mathit{x}}_m{)}^{⊺}({\mathit{R}}_{\mathrm{\Theta },n}^d{\mathit{W}}_k{\mathit{x}}_n)=\mathrm{Re}[\sum _{i=0}^{d/2-1}{\mathit{q}}_{[2i:2i+1]}{\mathit{k}}_{[2i:2i+1]}^{\ast }{e}^{i(m-n){\theta }_i}]$$

其中${\mathit{q}}_{[2i:2i+1]}$ 代表$\mathit{q}$ 中的第$2i$ 到第$2i+1$ 个元素, $\mathit{k}$ 同理.

令$h_i={\mathit{q}}_{[2i:2i+1]}{\mathit{k}}_{[2i:2i+1]}^{\ast },S_j=\sum _{i=0}^{j-1}{e}^{\text{i}(m-n){\theta }_i}$, 同时设$h_{d/2}=0,S_0=0$, 由Abel变换可以得到:

$$\sum _{i=0}^{d/2-1}{\mathit{q}}_{[2i:2i+1]}{\mathit{k}}_{[2i:2i+1]}^{\ast }{e}^{i(m-n){\theta }_i}=\sum _{i=0}^{d/2-1}{h}_i(S_{i+1}-S_i)=-\sum _{i=0}^{d/2-1}{S}_{i+1}(h_{i+1}-h_i)$$

因此有:

$$\begin{array}{rl}|\sum _{i=0}^{d/2-1}{\mathit{q}}_{[2i:2i+1]}{\mathit{k}}_{[2i:2i+1]}^{\ast }{e}^{i(m-n){\theta }_i}|& =|\sum _{i=0}^{d/2-1}{S}_{i+1}(h_{i+1}-h_i)|\\ & \le \sum _{i=0}^{d/2-1}|S_{i+1}||(h_{i+1}-h_i)|\\ & \le (\underset{i}{max}|h_{i+1}-h_i|)\sum _{i=0}^{d/2-1}|S_{i+1}|\end{array}$$

即当相对位置$m-n$ 增大时, 在${\theta }_i={10000}^{-2i/d}$的设定下, $\frac{1}{d/2}\sum _{i=1}^{d/2}\left|S_i\right|$ 的值是逐渐减小的, 如下图所示:

Experiments and Evaluations

详细的实验参数设置和模型参数设置请参考原论文.

Machine Translation

作者将RoFormer与Vallina Transformer在WMT 2014 English-German(450w平行语料)上对比:

RoFormer略高于Vallina Transformer.

Pre-training Language Modeling

接着, 在BookCorpus和Wikipedia这两个常用的语料上做预训练, 与BERT做对比, 结果如下:

RoFormer明显的比BERT要收敛更快一些, 而且从曲线上来看它的训练看起来更加稳定.

Fine-tuning on GLUE tasks

在GLUE上做Finetune, 继续与BERT对比:

RoFormer在MRPC, STS-B, QQP上有显著优势, 它们都是偏语义类的任务, 在SST-2和QNLI, MNLI中表现弱于BERT.

Performer with RoPE

上文中提到了RoPE在Linear Attention中的优势, 因此测试与使用Linear Attention的Performer做个结合:

RoPE能够在维持原复杂度不变的条件下加速收敛.

Evaluation on Chinese Data

除了在英文数据上进行评估, 还要在中文数据上进行评估.

在本实验中, 作者将自己提出的WoBERT的绝对位置编码替换为RoPE:

并在34G的语料上进行不同参数的多阶段训练, 以适应不同的输入长度:

随着序列长度的增加, RoFormer总是能取得更好的性能, 说明RoFormer在长度外推设置下表现较好.

在CAIL2019-SCM(一个语义相似任务)上来测试RoFormer的长文本能力, 表现如下:

• 当RoFormer输入长为512时, 能够略优于BERT和WoBERT.
• 当RoFormer的最大输入长度扩展到1024的时候, 显著优于BERT和WoBERT.

此外, 还能观察到在中文场景下, 词建模的WoBERT要优于字建模的BERT. 我以前也有类似的观点, 更大的优势是词建模如果用到生成类任务中能带来更快的推理速度.

Summary

与其说RoPE简洁旦优雅, 不如说RoPE简洁甚至优雅. RoPE用绝对位置编码实现了相对位置编码, 并且能够丝滑无缝的融入到Self-Attention当中.

由于我硕士期间是做信息抽取的, 所以我印象中RoPE最早在信息抽取里面运用的非常广泛, 其中可能原因有两个:

• 其一是因为相对位置对信息抽取的Span抽取影响比较大, 因为每个Span的都是通过起止边界位置确定的, 在Universal Information发展后期, 基本都是要带上RoPE.
• 另外可能是因为信息抽取抽取Span的大多都是通过一个乘积计算得到得分, 与内积形式一致, 这与RoPE给Self-Attention设计的形式恰恰好是一致的, 所以RoPE表现好也不奇怪.

RoPE早期在其他领域似乎没有得到特别多的关注. 但是随着PaLM, LLaMA等LLM采用了RoPE, RoPE便渐渐的在LLM里面成为一项标准的配置, 所以RoPE甚至也被人誉为LLM时代的ResNet.

在更长的上下文场景下, RoPE的续作是ReRoPE, 超过了NTK-RoPE, 感兴趣的可以继续深入.

虽然论文本身不是苏神自己写的, 但最后在文末也不得不感叹, 苏神真强啊.