旋转位置编码

论文名字：ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING

这篇论文主要介绍了一种名为RoFormer的增强型Transformer模型，其核心是提出了旋转位置嵌入（RoPE）方法来改进Transformer中的位置编码机制，具体内容如下：

\[\begin{equation} \begin{split} q_m &= f_q(x_m, m) \\ k_n &= f_k(x_n, n) \\ v_n &= f_v(x_n, n) \end{split} \end{equation}\]

$ q_m $ 和 $ k_n $ 是m位置查询向量和n位置键向量，$ v_n $ 是n位置对应的值向量。

下面计算注意力权重：

\[\begin{equation} \begin{split} a_{m,n} &=\frac{\exp\left(\frac{q_m^\top k_n}{\sqrt{d}}\right)}{\sum_{j=1}^N \exp\left(\frac{q_m^\top k_j}{\sqrt{d}}\right)} \\ o_m &=\sum_{n=1}^N a_{m,n} v_n \end{split} \end{equation}\]

一、绝对位置编码

绝对位置编码，是在特定位置上面，添加一个固定的偏置项，这个偏置项是确定的不可训练的，绝对位置编码对能处理的序列长度有限制，且不能很好的处理位置间的相对关系。

\[\begin{equation} f_{t: t \in \{q,k,v\}}(x_i, i) := W_{t: t \in \{q,k,v\}}(x_i + p_i) \end{equation}\] \[\begin{equation} \begin{split} p_{i,2t} &= \sin(k / 10000^{2t/d}) \\ p_{i,2t+1} &= \cos(k / 10000^{2t/d}) \end{split} \end{equation}\]

二、相对位置编码

相对位置编码，每个相对位置关系，对应的emb是可以训练的

\[\begin{equation} \begin{split} f_q(x_m) &= W_q x_m \\ f_k(x_n, n) &= W_k(x_n + \tilde{p}_r^k) \\ f_v(x_n, n) &= W_v(x_n + \tilde{p}_r^v) \end{split} \end{equation}\]

公式中 \(\tilde{p}_r^k\)，\(\tilde{p}_r^v\)是可以训练的相对位置emb，其中\(r=clip(m-n, r_{min}, r_{max})\)代表了m和n的相对距离。 参考公式3的形式，$x_i$ 和 $p_i$相加后映射成q，k，v向量，公式2计算权重重新进行展开后，得到下面的公式：

\[\begin{equation} q_m^\top k_n = x_m^\top W_q^\top W_k x_n + x_m^\top W_q^\top W_k p_n + p_m^\top W_q^\top W_k x_n + p_m^\top W_q^\top W_k p_n \end{equation}\]

这个公式有4项，第一项为不加位置emb情况下的原始展开项，后面的三项，分别为 $x_m$ 和 $p_n$ 交叉项，$p_m$ 和 $x_n$ 交叉项，$p_m$ 和 $p_n$交叉项。

相对位置编码的核心思路是，对于$p_n$项，替换成正弦位置编码的相对位置编码$\tilde{p}_{m-n}$，而向量$p_m$分别替换成两个向量u,v。 进一步的，$W_k$也进行了拆解，变为基于内容的$W_k$和基于位置的$\tilde{W}_k$，进行替换后，公式变为如下形式：

\[\begin{equation} q_m^\top k_n = x_m^\top W_q^\top W_k x_n + x_m^\top W_q^\top \tilde{W}_k \tilde{p}_{m-n} + u^\top W_q^\top W_k x_n + v^\top W_q^\top \tilde{W}_k \tilde{p}_{m-n} \end{equation}\]

在一些研究中，发现位置信息，更适合添加在q，k向量上面，而不适合添加在v向量上面，这是因为对于非结构话的文本，将位置信息重点放在 q 和 k 的计算中，而减少 v 对位置信息的依赖，有助于模型更清晰地分离 “位置相关性” 与 “内容语义”，提升注意力机制的效率。 但是，这个结论也并非绝对的，在某些特定的任务上，v向量也可以添加位置信息。 我们对公式6进行改写，得到如下形式，其中$b_{i,j}$是可以训练的向量矩阵。

\[\begin{align} q_m^\top k_n &= x_m^\top W_q^\top W_k x_n + b_{i,j} \\ q_m^\top k_n &= x_m^\top W_q^\top W_k x_n + p_m^\top U_q^\top U_k p_n + b_{i,j} \end{align}\]

经理论分析，He et al.[2022]提出，公式6的中间两项，足以表达相对位置关系，只需要将$p_m$和$p_n$替换成相对位置编码$p_{m-n}$

\[\begin{equation} q_m^\top k_n = x_m^\top W_q^\top W_k x_n + x_m^\top W_q^\top W_k \tilde{p}_{m-n} + \tilde{p}_{m-n}^\top W_q^\top W_k x_n \end{equation}\]

三、旋转位置编码

在attention注意力机制中，为了将相对位置信息融入到$q_m$和$k_n$的计算中，我们需要一个仅仅依赖$q_m$，$k_n$和m-n的函数：

\[\begin{equation} \langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, m - n) \end{equation}\]

公式11就是本文想寻找的一个理想的形式，通过函数$f_q(x_m, m)$和$f_k(x_n, n)$，对$x_m$和$x_n$进行操作，使其后续操作，能够自然的包含了位置信息。

\[\begin{equation} \begin{split} f_q(x_m, m) &= (W_{q}x_m)e^{im\theta}\\ f_k(x_n, n) &= (W_{k}x_n)e^{in\theta} \\ g(x_m, x_n, m - n) &= \text{Re}[(W_{q}x_m)(W_{k}x_n)^*e^{i(m - n)\theta}] \end{split} \end{equation}\]

公式中的Re[.]，$(W_{k}x_n)^*是W_{k}x_n的共轭复数$。

通过共轭复数与原复数的乘积运算，将复数乘法的结果投影到实部，从而将基于复数旋转的位置编码转换为实数空间中的向量运算，实现对相对位置信息的建模。 这个公式的形式太复杂了，进一步展开。

\[\begin{equation} f_{q,k}(x_m,m) = \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} W_{\{q,k\}}^{(11)} & W_{\{q,k\}}^{(12)} \\ W_{\{q,k\}}^{(21)} & W_{\{q,k\}}^{(22)} \end{pmatrix} \begin{pmatrix} x_m^{(1)} \\ x_m^{(2)} \end{pmatrix} \end{equation}\]

文中后续的公式，不好理解，我们从实际例子出发来看看旋转位置编码如何操作，最终$f_q$和$f_k$的函数形式没有差异，所以只需要举一个例子说明如何旋转即可。 假设有一个长度为5，维度为d=4的序列，序列的位置下标，分别为[0,1,2,3,4]，旋转的基础频率f设置为1/10000（经验值）。 基于以上的设置，在位置n的第i个位置对应的旋转角度为$n*(1/10000)^{-2i/d}$

计算步骤

1、拆分向量 需要把向量拆分成两个一组，我们这里维度为4，所以会拆出来两组。 具体如下： 假设序列对应的值如下：

序列编号	序列向量值	拆分后组1	拆分后组2
0	[1,0,1,0]	[1,0]	[1,0]
1	[0,1,0,1]	[0,1]	[0,1]
2	[1,1,1,1]	[1,1]	[1,1]
3	[1,-1,1,-1]	[1,-1]	[1,-1]
4	[0.5,0.5,0.5,0.5]	[0.5, 0.5]	[0.5, 0.5]

计算旋转角度，组1,2的旋转角度分别如下：

\[\theta_{m,0}=m*base^{-0/4}=m\] \[\theta_{m,1}=m*base^{-2/4}=m/100\]

旋转矩阵如下：

\[\begin{bmatrix} cos(\theta) & sin(\theta) \\ -sin(\theta) & cos(\theta) \end{bmatrix}\]

这边我们写一段python代码，计算最终的向量：

import math 
vec_list = [[1,0,1,0],[0,1,0,1],[1,1,1,1],[1,-1,1,-1],[0.5,0.5,0.5,0.5]]
for idx, vec in enumerate(vec_list):
    vec1 = vec[:2]
    vec2 = vec[2:]
    theta0 = idx
    theta1 = idx/100
    cos0 = math.cos(theta0)
    sin0 = math.sin(theta0)
    cos1 = math.cos(theta1)
    sin1 = math.sin(theta1)
    vec1 = [cos0*vec1[0] - sin0*vec1[1], sin0*vec1[0] + cos0*vec1[1]]
    vec2 = [cos1*vec2[0] - sin1*vec2[1], sin1*vec2[0] + cos1*vec2[1]]
    print([vec1, vec2])

对向量一次应用旋转矩阵，如下：

序列编号	序列向量值	拆分后组1	拆分后组2	$\theta_0, \theta_1$	旋转后组1	旋转后组2
0	[1,0,1,0]	[1,0]	[1,0]	0, 0	[1,0]	[1,0]
1	[0,1,0,1]	[0,1]	[0,1]	1, 0.01	[-0.8415, 0.5403]	[-0.0100, 0.9999]
2	[1,1,1,1]	[1,1]	[1,1]	2, 0.02	[-1.3254, 0.4932]	[0.9798, 1.0198]
3	[1,-1,1,-1]	[1,-1]	[1,-1]	3, 0.03	[-0.8489, 1.1311]	[1.0296, -0.9696]
4	[0.5,0.5,0.5,0.5]	[0.5, 0.5]	[0.5, 0.5]	4, 0.04	[0.0516, -0.7052]	[0.4796, 0.5196]

至此我们完成的了不同位置的向量旋转，我们可以发现，旋转后的向量，包含了位置信息，并且旋转后的向量，和原始向量的维度是一致的。 对于一个高维向量（维度d为偶数），我们可以把它拆成两两一组的形式，每组形成一个二维空间中的点。对每组应用旋转操作（旋转角度可能不同）后，整体效果等效于在高维空间中进行一个特定旋转，该旋转由多个相互正交的二维旋转组成。