上一篇文章中，我们对原始的Sinusoidal位置编码做了较为详细的推导和理解，总的感觉是Sinusoidal位置编码是一种"想要成为相对位置编码的绝对位置编码"。
本文将会介绍追一科技自研Rotary Transformer（RoFormer）模型，它的主要改动是引用了苏剑林大佬构思的"旋转式位置编码（Rotary Position Embedding，RoPE）"，这是一种配合Attention机制能达到"绝对位置编码的方式实现相对位置编码的设计"。而也正因为这种设计，它还是目前唯一一种可用于线性Attention的相对位置编码
基本思路
我们假设通过下述运算来给boldsymbol{q},boldsymbol{k} 添加绝对位置信息：
也就是说，我们分别为boldsymbol{q},boldsymbol{k} 设计操作boldsymbol{f}(cdot, m),boldsymbol{f}(cdot, n) ，使得经过该操作后，tilde{boldsymbol{q}}_m,tilde{boldsymbol{k}}_n 就带有了位置m,n 的绝对位置信息。Attention的核心运算是内积，所以我们希望内积的结果带有相对位置信息，因此假设存在恒等关系：
所以我们要求出该恒等式的一个（尽可能简单的）解。求解过程还需要一些初始条件，显然我们可以合理地设f(boldsymbol{q},0)=boldsymbol{q} 和f(boldsymbol{k},0)=boldsymbol{k}
求解过程
同上一篇思路一样，我们先考虑二维情形，然后借助复数来求解。在复数中有langleboldsymbol{q},boldsymbol{k}rangle=text{Re}[boldsymbol{q}boldsymbol{k}^*] ，text{Re}[] 代表复数的实部，我们希望存在某个函数g ，使其能够包含相对位置信息，即：
简单起见，我们假设存在复数g(boldsymbol{q},boldsymbol{k},m-n) ，使得boldsymbol{f}(boldsymbol{q}, m)boldsymbol{f}^*(boldsymbol{k}, n) = boldsymbol{g}(boldsymbol{q},boldsymbol{k},m-n) ，然后我们用复数的指数形式，设
begin{equation}begin{aligned} boldsymbol{f}(boldsymbol{q}, m) =&, R_f (boldsymbol{q}, m)e^{text{i}Theta_f(boldsymbol{q}, m)} \ boldsymbol{f}(boldsymbol{k}, n) =&, R_f (boldsymbol{k}, n)e^{text{i}Theta_f(boldsymbol{k}, n)} \ boldsymbol{g}(boldsymbol{q}, boldsymbol{k}, m-n) =&, R_g (boldsymbol{q}, boldsymbol{k}, m-n)e^{text{i}Theta_g(boldsymbol{q}, boldsymbol{k}, m-n)} \ end{aligned}tag{4}end{equation}
那么带入方程后就得到方程组
begin{equation}begin{aligned} R_f (boldsymbol{q}, m) R_f (boldsymbol{k}, n) =&, R_g (boldsymbol{q}, boldsymbol{k}, m-n) \ Theta_f (boldsymbol{q}, m) - Theta_f (boldsymbol{k}, n) =&, Theta_g (boldsymbol{q}, boldsymbol{k}, m-n) end{aligned}tag{5}end{equation}
对于第一个方程，带入m=n=0 得到
最后一个等号源于初始条件boldsymbol{f}(boldsymbol{q}, 0)=boldsymbol{q} 和boldsymbol{f}(boldsymbol{k}, 0)=boldsymbol{k} 。所以现在我们可以很简单地设R_f (boldsymbol{q}, m)=Vert boldsymbol{q}Vert, R_f (boldsymbol{k}, m)=Vert boldsymbol{k}Vert ，即它不依赖于m 。至于第二个方程，同样带入m=n=0 得到
这里的Theta(boldsymbol{q}),Theta(boldsymbol{k}) 是boldsymbol{q},boldsymbol{k} 本身的幅角，最后一个等号同样源于初始条件。根据上式得到Theta_f (boldsymbol{q}, m) - Theta (boldsymbol{q}) = Theta_f (boldsymbol{k}, m) - Theta (boldsymbol{k}) ，所以Theta_f (boldsymbol{q}, m) - Theta (boldsymbol{q}) 应该是一个只与m 相关，跟boldsymbol{q} 无关的函数，记为varphi(m) ，即Theta_f (boldsymbol{q}, m) = Theta (boldsymbol{q}) varphi(m) 。接着带入n=m-1 ，整理得到
首先将n=m-1 带入式(5)的第二个式子，得到 begin{aligned} Theta_f (boldsymbol{q}, m) - Theta_f (boldsymbol{k}, m-1) =, Theta_g (boldsymbol{q}, boldsymbol{k}, 1) end{aligned} 上式两边同减Theta(boldsymbol{q}) 得 begin{aligned} Theta_f (boldsymbol{q}, m) - Theta(boldsymbol{q}) - Theta_f (boldsymbol{k}, m-1) =, Theta_g (boldsymbol{q}, boldsymbol{k}, 1) - Theta(boldsymbol{q}) end{aligned} 因为Theta_f (boldsymbol{q}, m) - Theta(boldsymbol{q})=varphi(m) ，所以 begin{aligned} varphi(m) - Theta_f (boldsymbol{k}, m-1) =, Theta_g (boldsymbol{q}, boldsymbol{k}, 1) - Theta(boldsymbol{q}) end{aligned} 上式两边同加Theta(boldsymbol{k}) 得 begin{aligned} varphi(m) Theta(boldsymbol{k}) - Theta_f (boldsymbol{k}, m-1) =, Theta_g (boldsymbol{q}, boldsymbol{k}, 1) Theta(boldsymbol{k}) - Theta(boldsymbol{q}) end{aligned} 因为Theta (boldsymbol{q}) - Theta_f (boldsymbol{q}, m) = Theta (boldsymbol{k}) - Theta_f (boldsymbol{k}, m) ，所以 begin{aligned} varphi(m) Theta (boldsymbol{q}) - Theta_f (boldsymbol{q}, m-1) =, Theta_g (boldsymbol{q}, boldsymbol{k}, 1) Theta(boldsymbol{k}) - Theta(boldsymbol{q}) end{aligned} 最后因为Theta_f (boldsymbol{q}, m-1) - Theta(boldsymbol{q})=varphi(m-1) ，所以
说了这么多，就是为了证明{varphi(m)} 是等差数列，设右端为theta ，那么就解得varphi(m)=mtheta
编码形式
综上，我们得到二维情况下用复数表示的RoPE：
begin{aligned} Vert qVert e^{text{i}(Theta(q) mtheta)}&=Vert qVert (e^{text{i}Theta(boldsymbol{q})}cdot e^{text{i}mtheta}) \ &= (Vert qVert e^{text{i}Theta(boldsymbol{q})} )e^{text{i}mtheta}\ &=boldsymbol{q}e^{text{i}mtheta} end{aligned}
根据复数乘法的几何意义，该变换实际上对应着向量的旋转，所以我们称之为"旋转式位置编码"，它还可以写成矩阵形式：
begin{equation} boldsymbol{f}(boldsymbol{q}, m) =begin{pmatrix}cos mtheta & -sin mtheta\ sin mtheta & cos mthetaend{pmatrix} begin{pmatrix}q_0 \ q_1end{pmatrix}tag{10}end{equation}
为什么旋转对应矩阵相乘，可以看这篇文章：旋转之一 - 复数与2D旋转，或者大家直接搜复数乘法与向量旋转
由于内积满足线性叠加性，因此任意偶数维的RoPE，我们都可以表示为二维情形的拼接，即
begin{equation}scriptsize{underbrace{begin{pmatrix} cos mtheta_0 & -sin mtheta_0 & 0 & 0 & cdots & 0 & 0 \ sin mtheta_0 & cos mtheta_0 & 0 & 0 & cdots & 0 & 0 \ 0 & 0 & cos mtheta_1 & -sin mtheta_1 & cdots & 0 & 0 \ 0 & 0 & sin mtheta_1 & cos mtheta_1 & cdots & 0 & 0 \ vdots & vdots & vdots & vdots & ddots & vdots & vdots \ 0 & 0 & 0 & 0 & cdots & cos mtheta_{d/2-1} & -sin mtheta_{d/2-1} \ 0 & 0 & 0 & 0 & cdots & sin mtheta_{d/2-1} & cos mtheta_{d/2-1} \ end{pmatrix}}_{boldsymbol{mathcal{R}}_m} begin{pmatrix}q_0 \ q_1 \ q_2 \ q_3 \ vdots \ q_{d-2} \ q_{d-1}end{pmatrix}}tag{11}end{equation}
也就是说，给位置为m 的向量boldsymbol{q} 乘上矩阵boldsymbol{R}_m ，位置为n 的向量boldsymbol{k} 乘上矩阵boldsymbol{R}_n ，用变换后的boldsymbol{Q},boldsymbol{K} 序列做Attention，那么Attention就自动包含相对位置信息了，因为下列等式恒成立：
总可以利用积化和差公式，使得cos 或者sin 中带有(m-n) 项
值得指出的是，boldsymbol{R}_m 是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性
由于boldsymbol{R}_m 的稀疏性，所以直接用矩阵乘法来实现会很浪费算力，推荐通过使用下述方式来实现RoPE：
begin{equation}begin{pmatrix}q_0 \ q_1 \ q_2 \ q_3 \ vdots \ q_{d-2} \ q_{d-1} end{pmatrix}otimesbegin{pmatrix}cos mtheta_0 \ cos mtheta_0 \ cos mtheta_1 \ cos mtheta_1 \ vdots \ cos mtheta_{d/2-1} \ cos mtheta_{d/2-1} end{pmatrix} begin{pmatrix}-q_1 \ q_0 \ -q_3 \ q_2 \ vdots \ -q_{d-1} \ q_{d-2} end{pmatrix}otimesbegin{pmatrix}sin mtheta_0 \ sin mtheta_0 \ sin mtheta_1 \ sin mtheta_1 \ vdots \ sin mtheta_{d/2-1} \ sin mtheta_{d/2-1} end{pmatrix}tag{13}end{equation}
其中otimes 是逐位对应相乘，即Numpy、Tensorflow、PyTorch等计算框架中的* 运算。从这个实现角度也可以看到，RoPE可以视为是乘性位置编码的变体
远程衰减
可以看到，RoPE形式上和Sinusoidal位置编码有点相似，只不过Sinusoidal位置编码是加性的，而RoPE可以视为乘性的。在theta_i 的选择上，我们同样沿用了Sinusoidal位置编码的方案，即theta_i=10000^{−2i/d} ，它可以带来一定的远程衰减性
具体证明如下：将boldsymbol{q},boldsymbol{k} 两两分组后，它们加上RoPE后的内积可以用复数乘法表示为
begin{equation} (boldsymbol{mathcal{R}}_m boldsymbol{q})^{top}(boldsymbol{mathcal{R}}_n boldsymbol{k}) = text{Re}left[sum_{i=0}^{d/2-1}boldsymbol{q}_{[2i:2i 1]}boldsymbol{k}_{[2i:2i 1]}^* e^{text{i}(m-n)theta_i}right]tag{14}end{equation}
式(14)主要是由式(12)变形得到
记h_i = boldsymbol{q}_{[2i:2i 1]}boldsymbol{k}_{[2i:2i 1]}^*, S_j = sumlimits_{i=0}^{j-1} e^{text{i}(m-n)theta_i} ，并约定h_{d/2}=0,S_0=0 ，那么有：
begin{equation}sum_{i=0}^{d/2-1}boldsymbol{q}_{[2i:2i 1]}boldsymbol{k}_{[2i:2i 1]}^* e^{text{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)tag{15}end{equation}
式(15)的第二个等式主要由Abel分部求和法得到（令d/2=n ）： begin{aligned} sum_{i=0}^{n-1} h_i (S_{i 1} - S_i)&=-0 h_{n-1}S_n-h_{n-1}S_{n-1} h_{n-2}S_{n-1}-h_{n-2}S_{n-2}cdots h_0S_1-h_0S_0\ &=-h_{n}S_n h_{n-1}S_n cdots h_0S_1-h_0S_0\ &=sum_{i=0}^{n-1}S_{i 1}(h_i-h_{i 1})\ &=-sum_{i=0}^{n-1}S_{i 1}(h_{i 1}-h_i) end{aligned}
所以
begin{equation}begin{aligned} left|sum_{i=0}^{d/2-1}boldsymbol{q}_{[2i:2i 1]}boldsymbol{k}_{[2i:2i 1]}^* e^{text{i}(m-n)theta_i}right| =&, left|sum_{i=0}^{d/2-1} S_{i 1}(h_{i 1} - h_i)right| \ leq&, sum_{i=0}^{d/2-1} |S_{i 1}| |h_{i 1} - h_i| \ leq&, left(max_i |h_{i 1} - h_i|right)sum_{i=0}^{d/2-1} |S_{i 1}| end{aligned}tag{16}end{equation}
因此我们可以考察frac{1}{d/2}sumlimits_{i=1}^{d/2} |S_i| 随着相对距离变化的情况作为衰减性的体现，Mathematica代码如下
代码语言：javascript复制d = 128;
[Theta][t_] = 10000^(-2*t/d);
f[m_] = Sum[
Norm[Sum[Exp[I*m*[Theta][i]], {i, 0, j}]], {j, 0, d/2 - 1}]/(d/2);
Plot[f[m], {m, 0, 256}, AxesLabel -> {相对距离, 相对大小}]结果如下图：
从途中我们可以看到，随着相对距离的变大，内积结果有衰减趋势的出现。因此，选择theta_i=10000^{-2i/d} ，确实能带来一定的远程衰减性。当然，同上一篇文章说的一样，能带来远程衰减性的不止这个选择，几乎任意的光滑单调函数都可以，这里只是沿用了已有的选择而已。苏剑林大佬还试过以theta_i=10000^{-2i/d} 为初始化，将theta_i 视为可训练参数，然后训练一段时间后发现theta_i 并没有显著更新，因此干脆就直接固定theta_i=10000^{-2i/d} 了
线性场景
最后，我们指出，RoPE是目前唯一一种可以用于线性Attention的相对位置编码。这是因为其他的相对位置编码，都是直接基于Attention矩阵进行操作的，但是线性Attention并没有事先算出Attention矩阵，因此也就不存在操作Attention矩阵的做法，所以其他的方案无法应用到线性Attention中。而对于RoPE来说，它是用绝对位置编码的方式来实现相关位置编码，不需要操作Attention矩阵，因此有了应用到线性Attention的可能性
关于线性Attention的介绍，这里不再重复，有需要的读者请参考《去掉 Attention 的 Softmax，复杂度降为 O (n)》。线性Attention的常见形式是：
其中phi,varphi 是值域非负的激活函数。可以看到，线性Attention也是基于内积的，所以很自然的想法是将RoPE插入到内积中：
但这样存在的问题是，[boldsymbol{mathcal{R}}_iphi(boldsymbol{q}_i)]^{top} [boldsymbol{mathcal{R}}_jvarphi(boldsymbol{k}_j)] 可能为负数，因此它不再是常规的概率注意力，而且分母有0的风险，可能会带来优化上的不稳定，考虑到boldsymbol{R}_i,boldsymbol{R}_j 都是正交矩阵，它不改变向量的模长，因此我们可以抛弃常规的概率归一化要求，使用如下运算作为一种新的线性Attention：
也就是说，RoPE只插入分子中，而分母则不改变，这样的注意力不再是基于概率的（注意力矩阵不再满足归一性），但它某种意义上来说也是一个归一化方案，而且也没有证据表明非概率式的注意力就不好（比如Nyströmformer也算是没有严格依据概率分布的方式构筑注意力），所以我们将它作为候选方案之一进行实验，而我们初步的实验结果显示这样的线性Attention也是有效的
此外，苏剑林大佬在《去掉 Attention 的 Softmax，复杂度降为 O (n)》中还提出过另外一种线性Attention方案：
它不依赖于值域的非负性，而RoPE也不改变模长，因此RoPE可以直接应用于此类线性Attention，并且不改变它的概率意义
模型开源
关于模型开源部分请查看文章开头部分的Rotary Transformer（RoFormer）
Reference
- Transformer升级之路：2、博采众长的旋转式位置编码
