旋转位置编码详细介绍_旋转位置编码复数形式

旋转矩阵

1. 复数$z=a+bi$可以看成一个向量$\left[\begin{array}{c}a\\ b\end{array}\right]$，反过来，一个二维向量也可以看成一个复数。
2. 复数$z=a+bi$也可以看成复平面上的一个点，在极坐标系下，
   $(a=r\cos \theta ,b=r\cos \theta )$，其中$\theta$为幅角，$r$为模长，等于$\sqrt{(a^2+b^2)}$
3. 现在我们考虑将一个向量$\left[\begin{array}{c}x\\ y\end{array}\right]$旋转$\theta$，怎么计算旋转之后的向量?
   旋转之前的向量，用极坐标可以这么表示，$\alpha$为起始的幅角
   $$\{\begin{array}{l}x=r\cos \alpha \\ y=r\sin \alpha \end{array}$$
   旋转之后的向量，同样可以这么表示
   $$\{\begin{array}{l}{x}^{\prime }=r\cos (\alpha +\theta )\\ y^{\prime }=r\sin (\alpha +\theta )\end{array}$$
   把上述公式用和差化积展开
   $$\begin{array}{c}\begin{array}{rl}{x}^{\prime }& =r\cos (\alpha +\theta )\\ & =r(\cos \alpha \cos \theta -\sin \theta \sin \alpha )\\ & =x\cos \theta -y\sin \theta \end{array}\end{array}$$
   $$\begin{array}{c}\begin{array}{rl}{y}^{\prime }& =r\sin (\alpha +\theta )\\ & =r(\sin \alpha \cos \theta +\cos \alpha \sin \theta )\\ & =x\sin \theta +y\cos \theta \end{array}\end{array}$$
   整理一下，写成矩阵形式，
   $$\left[\begin{array}{c}{x}^{\prime }\\ y^{\prime }\end{array}\right]=\left[\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]\left[\begin{array}{c}x\\ y\end{array}\right]$$
   上述公式的$\left[\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]$就是我们所说的旋转矩阵

复数乘以复数

对于两个复数，$z_1=a+bi$和$z_2=c+di$，计算$z_1{z}_2$

1. 矩阵视角
   $$\begin{array}{c}\begin{array}{rl}{z}_1{z}_2& =(a+bi)(c+di)\\ & =(ac-bd)+(ad+bc)i\end{array}\end{array}$$
   可以看成矩阵与向量的乘积，即
   $$\left[\begin{array}{cc}a& -b\\ b& a\end{array}\right]\left[\begin{array}{c}c\\ d\end{array}\right]$$
2. 极坐标系视角
   $z_1=r_1(\cos {\theta }_1+i\sin {\theta }_1)$
   $z_2=r_2(\cos {\theta }_2+i\sin {\theta }_2)$
   $z_1{z}_2=r_1{r}_2(\cos ({\theta }_1+{\theta }_2)+i\sin ({\theta }_1+{\theta }_2))$
   可以看成将复数$z_1$旋转${\theta }_2$，并且将模长缩放$r_2$。或者是将$z_2$旋转${\theta }_1$，模长缩放$r_1$

旋转位置编码

目标：对q和k分别添加绝对位置信息，在做完点乘之后，具有相对位置信息。即q和k的位置分别是m,n，点乘之后，位置信息只与m-n有关

其中，$e^{ix}=\cos x+i\sin x$
$W_q,W_k$是quey和key对应的可学习矩阵，假设输入的query和key是$d_m$维的，那么$x_m,x_n\in {\mathcal{R}}^{d_m\times 1}$,使用$W_q,W_k$进行变换，变换之后的结果是${\mathcal{R}}^{d_m\times 1}$，这里我们只考虑2维的情况。
记$x_m,x_n$变换之后的向量是$q_m,k_n\in {\mathcal{R}}^{2\times 1}$。$q_m{e}^{im\theta }$是一个向量乘以一个复数，二维向量可以看成一个复数，那么上式可以看成两个复数相乘，那么根据第二部分的知识，两个复数的乘积等于：幅角相加，模长相乘。而$e^{im\theta }$的模长是1，即相乘之后的模长不变，只进行了旋转，即把$q_m$向量旋转了$m\theta$角度。根据第一部分的知识，旋转就是乘上一个旋转矩阵，即
f q ( x m , m ) = [ cos ⁡ m θ − sin ⁡ m θ sin ⁡ m θ cos ⁡ m θ ] [ q m 1 q m 2 ] f_q(x_m,m)=

$$\begin{bmatrix} \cos m\theta & -\sin m\theta\\ \sin m\theta & \cos m\theta \end{bmatrix}$$ $$\begin{bmatrix} q_m^1 \\ q_m^2 \end{bmatrix}$$ fq​(xm​,m)=[cosmθsinmθ​−sinmθcosmθ​][qm1​qm2​​]

f k ( x n , n ) = [ cos ⁡ n θ − sin ⁡ n θ sin ⁡ n θ cos ⁡ n θ ] [ k n 1 k n 2 ] f_k(x_n,n)=

$$\begin{bmatrix} \cos n\theta & -\sin n\theta\\ \sin n\theta & \cos n\theta \end{bmatrix}$$ $$\begin{bmatrix} k_n^1 \\ k_n^2 \end{bmatrix}$$ fk​(xn​,n)=[cosnθsinnθ​−sinnθcosnθ​][kn1​kn2​​]
对于$q_m{k}_n^{\ast }{e}^{i(m-n)\theta }$，前面的$q_m{k}_n^{\ast }$可以看成两个复数相乘，可以转化为复数的矩阵形式，即
$$\begin{array}{c}\begin{array}{rl}{q}_m{k}_n^{\ast }& =\left[\begin{array}{cc}{q}_m^1& -q_m^2\\ q_m^2& q_m^1\end{array}\right]\left[\begin{array}{c}{k}_n^1\\ -k_n^2\end{array}\right]\\ & =\left[\begin{array}{c}{q}_m^1{k}_n^1+q_m^2{k}_n^2\\ q_m^2{k}_{n}1-q_m^1{k}_n^2\end{array}\right]\end{array}\end{array}$$
再乘上$e^{i(m-n)\theta }$，相当于把上述的2维向量旋转了$(m-n)\theta$，乘上一个旋转矩阵即可，即
$$\begin{array}{c}\begin{array}{rl}{q}_m{k}_n^{\ast }{e}^{i(m-n)\theta }& =\left[\begin{array}{cc}\cos (m-n)\theta & -\sin (m-n)\theta \\ \sin (m-n)\theta & \cos (m-n)\theta \end{array}\right]\left[\begin{array}{c}{q}_m^1{k}_n^1+q_m^2{k}_n^2\\ q_m^2{k}_{n}1-q_m^1{k}_n^2\end{array}\right]\end{array}\end{array}$$
只取实数部分，即
$$\begin{array}{c}\begin{array}{rl}g(x_m,x_n,m-n)& =(q_m^1{k}_n^1+q_m^2{k}_n^2)(\cos (m-n)\theta )-(q_m^2{k}_n^1-q_m^1{k}_n^2)(\sin (m-n)\theta )\end{array}\end{array}$$
对于$<f_q(x_m,m),f_k(x_n,n)>$，可以计算：
$$\begin{array}{c}\begin{array}{rl}<f_q(x_m,m),f_k(x_n,n)>& =\left(\left[\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right]\left[\begin{array}{c}{q}_m^1\\ q_m^2\end{array}\right]{)}^T\right(\left[\begin{array}{cc}\cos n\theta & -\sin n\theta \\ \sin n\theta & \cos n\theta \end{array}\right]\left[\begin{array}{c}{k}_n^1\\ k_n^2\end{array}\right])\\ & =g(x_m,x_n,m-n)\end{array}\end{array}$$

实现

根据上述的内容，我们可以发现，实现很容易，对query和key分别乘以一个旋转矩阵就可以了。对于多维的情况，两两分组即可。

上述的矩阵过于稀疏，计算效率不高。因此有如下的高效实现

另外，transformers库中的实现与论文有点区别，高维的情况，不是相邻两个一组。

参考资料

1. 十分钟读懂旋转编码（RoPE） - 绝密伏击的文章 - 知乎
2. 旋转之一 - 复数与2D旋转 - 二圈妹的文章 - 知乎
   
3. LLM时代Transformer中的Positional Encoding - MrYXJ的文章 - 知乎