RoPE
基础知识
定义一个长度为 $N$ 的输入序列为： $$ \mathbb{S}N = {w_i}{i=1}^N $$
其中 $w_i$ 表示输入序列中第 $i$ 个 token，而输入序列 $\mathbb{S}_N$ 对应的 embedding 表示为： $$ \mathbb{E}_N = {\mathbf{x}i}{i=1}^N $$ 其中 $\mathbf{x}_i$ 表示第 $i$ 个 token $w_i$ 对应的 $d$ 维词嵌入向量。
Transformer 作为一个序列模型，通过 Self-Attention 机制，避免了类似于 RNN 的逐步输入导致计算效率低的问题，可以实现全局并行。Transformer 每个位置都能直接与所有其他位置交互，第 $i$ 个 token 可以直接「看见」整个序列，通过注意力机制可以捕捉全局上下文关系。 $$ Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_{k}}}) V $$ 但是，Transformer 没有像 RNN 那样的递归结构，也没有像 CNN 那样的局部感知卷积核，因此它需要人为地加入位置信息。
在输入到 Transformer 之前，我们通常将每个 token 的 embedding $E(x_i)$ 与一个 位置向量 $P_i$ 相加： $$z_i = E(x_i) + P_i$$ 这样，模型可以同时感知到 内容（词向量） 和 位置信息（位置编码）。
绝对位置编码
|
|
构造位置编码的方法
用整型值标记位置
一种自然而然的想法是，给第一个 token 标记1，给第二个 token 标记2…，以此类推。
这种方法产生了以下几个主要问题：
- 模型可能遇见比训练时所用的序列更长的序列。不利于模型的泛化。
- 模型的位置表示是无界的。随着序列长度的增加，位置值会越来越大。
用 [0,1] 范围标记位置
为了解决整型值带来的问题，可以考虑将位置值的范围限制在 [0, 1] 之内，其中，0表示第一个 token，1表示最后一个 token。比如有3个 token，那么位置信息就表示成 [0, 0.5, 1]，若有四个 token，位置信息就表示成 [0, 0.33, 0.69, 1]。
但这样产生的问题是，当序列长度不同时，token间的相对距离是不一样的。例如在序列长度为3时，token间的相对距离为0.5；在序列长度为4时，token间的相对距离就变为0.33。
因此，我们需要这样一种位置表示方式，满足于：
- 它能用来表示一个 token 在序列中的绝对位置
- 在序列长度不同的情况下，不同序列中 token 的相对位置/距离也要保持一致
- 可以用来表示模型在训练过程中从来没有看到过的句子长度
用二进制向量标记位置
考虑到位置信息作用在 input embedding 上，因此比起用单一的值，更好的方案是用一个和 input embedding 维度一样的向量来表示位置。这时我们就很容易想到二进制编码。如下图，假设 d_model = 3，那么我们的位置向量可以表示成：
这下所有的值都是有界的（位于0，1之间），且 transformer 中的 d_model 本来就足够大，基本可以把我们要的每一个位置都编码出来了。
但是这种编码方式也存在问题：这样编码出来的位置向量，处在一个离散的空间中，不同位置间的变化是不连续的。假设 d_model = 2，我们有4个位置需要编码，这四个位置向量可以表示成 [0,0],[0,1],[1,0],[1,1]。我们把它的位置向量空间做出来：
用周期函数（sin）来表示位置
回想一下，现在我们需要一个有界又连续的函数，最简单的，正弦函数 sin 就可以满足这一点。我们可以考虑把位置向量当中的每一个元素都用一个 sin 函数来表示，则第 t 个 token 的位置向量可以表示为以下这个 $d_{model}$ 的向量
$$
PE_t = \left[ \sin\left( \frac{1}{2^0} t \right), \sin\left( \frac{1}{2^1} t \right), \dots, \sin\left( \frac{1}{2^{i-1}} t \right), \dots, \sin\left( \frac{1}{2^{d_{\text{model}} - 1}} t \right) \right]
$$
结合下图，来理解一下这样设计的含义。图中每一行表示一个 $PE_t$，每一列表示 $PE_t$ 中的第 i 个元素。旋钮用于调整精度，越往右边的旋钮，需要调整的精度越大，因此指针移动的步伐越小。每一排的旋钮都在上一排的基础上进行调整（函数中 t 的作用）。
通过频率 $\frac{1}{2^{i-1}}$ 来控制 sin 函数的波长，频率不断减小，则波长不断变大，此时 sin 函数对 t 的变动越不敏感，以此来达到越向右的旋钮，指针移动步伐越小的目的。这也类似于二进制编码，每一位上都是0和1的交互，越往低位走（越往左边走），交互的频率越慢。
为了避免这种情况，我们尽量将函数的波长拉长。一种简单的解决办法是同一把所有的频率都设成一个非常小的值。因此在 transformer 的论文中，采用了这个频率 $\frac{1}{10000^{i / (d_{\text{model}}-1)}}$，这里的 i 对应为每个 token 的位置向量中 $d_{model}$ 中的第 i 个值。
总结一下，到这里我们把位置向量表示为： $$ PE_t = \left[ \sin(w_0 t), \sin(w_1 t), \dots, \sin(w_{i-1} t), \dots, \sin(w_{d_{\text{model}} - 1} t) \right] $$
其中 $w_{i} =\frac{1}{10000^{i / (d_{\text{model}}-1)}}$
用 sin 和 cos 交替来表示位置
目前为止，我们的位置向量实现了如下功能：
- 每个 token 的向量唯一（每个 sin 函数的频率足够小）
- 位置向量的值是有界的，且位于连续空间中。模型在处理位置向量时更容易泛化，即更好处理长度和训练数据分布不一致的序列（sin函数本身的性质）
那现在我们对位置向量再提出一个要求，不同的位置向量是可以通过线性转换得到的。这样，我们不仅能表示一个 token 的绝对位置，还可以表示一个 token 的相对位置，即我们想要： $$ PE_{t+\Delta t} = T_{\Delta{t}} * PE_t $$ 这里，T 表示一个线性变换矩阵。观察这个目标式子，联想到在向量空间中一种常用的线形变换——旋转。在这里，我们将 t 想象为一个角度，那么 $\Delta{t}$ 就是其旋转的角度，则上面的式子可以进一步写成： $$ \begin{pmatrix} \sin(t + \triangle t) \ \cos(t + \triangle t) \end{pmatrix} = \begin{pmatrix} \cos \triangle t & \sin \triangle t \ -\sin \triangle t & \cos \triangle t \end{pmatrix} \begin{pmatrix} \sin t \ \cos t \end{pmatrix} $$ 有了这个构想，我们就可以把原来元素全都是 sin 函数的 $PE_t$ 做一个替换，我们让位置两两一组，分别用 sin 和 cos 的函数对来表示它们，则现在我们有： $$ PE_t = \left[ \sin(w_0 t), \cos(w_0 t), \sin(w_1 t), \cos(w_1 t), \dots, \sin\left(w_{\frac{d_{\text{model}}}{2} - 1} t\right), \cos\left(w_{\frac{d_{\text{model}}}{2} - 1} t\right) \right] $$
这里的
- $w_{i} =\frac{1}{10000^{2i / (d_{\text{model}})}}$
- $i = 0, 1, 2, …, \frac{d_{model}}{2}-1$
在这样的表示下，我们可以很容易用一个线性变换，把 $PE_t$ 转换成 $PE_{t+\Delta{t}}$
$$ PE_{t+\triangle t} = \begin{pmatrix} \begin{bmatrix} \cos(w_0 \triangle t) & \sin(w_0 \triangle t) \ -\sin(w_0 \triangle t) & \cos(w_0 \triangle t) \end{bmatrix} & & & \dots & 0 \ & \ddots & & & \vdots \ \vdots & & \begin{bmatrix} \cos(w_i \triangle t) & \sin(w_i \triangle t) \ -\sin(w_i \triangle t) & \cos(w_i \triangle t) \end{bmatrix} & & \vdots \ \vdots & & & \ddots & \vdots \ 0 & \dots & \dots & & \begin{bmatrix} \cos\left(w_{\frac{d_{\text{model}}}{2}-1} \triangle t\right) & \sin\left(w_{\frac{d_{\text{model}}}{2}-1} \triangle t\right) \ -\sin\left(w_{\frac{d_{\text{model}}}{2}-1} \triangle t\right) & \cos\left(w_{\frac{d_{\text{model}}}{2}-1} \triangle t\right) \end{bmatrix} \end{pmatrix} \begin{pmatrix} \sin(w_0 t) \ \cos(w_0 t) \ \vdots \ \sin(w_i t) \ \cos(w_i t) \ \vdots \ \sin\left(w_{\frac{d_{\text{model}}}{2}-1} t\right) \ \cos\left(w_{\frac{d_{\text{model}}}{2}-1} t\right) \end{pmatrix} = \begin{pmatrix} \sin\left(w_0 (t+\triangle t)\right) \ \cos\left(w_0 (t+\triangle t)\right) \ \vdots \ \sin\left(w_i (t+\triangle t)\right) \ \cos\left(w_i (t+\triangle t)\right) \ \vdots \ \sin\left(w_{\frac{d_{\text{model}}}{2}-1} (t+\triangle t)\right) \ \cos\left(w_{\frac{d_{\text{model}}}{2}-1} (t+\triangle t)\right) \end{pmatrix} $$
Sinusoidal 编码
刚才的 sin 和 cos 交替来表示位置，其实就是 Sinusoidal 的位置编码。这是一种经典的绝对位置编码，最初是由谷歌在论文 Attention is All You Need 中提出的方案，用于 Transformer 的位置编码。具体计算方式如下所示： $$\begin{split} P_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right) \ P_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right) \end{split} $$ 其中 pos 是位置，i 表示维度。
$i = 0, 1, 2, …, \frac{d_{model}}{2}-1$
看起来是通过 sin 和 cos 函数将位置编码的取值固定在了 [-1, 1] 之前，但是为什么不用线性函数？而且这里面的10000是怎么想到的？
谷歌在论文里面给出了解释：
- 具有相对位置表达能力：Sinusoidal 可以学习到相对位置，对于固定位置距离的 $k$，$PE(i+k)$ 可以表示成 $PE(i)$ 的线性函数。
- 两个位置向量的内积只和相对位置 $k$ 有关。
- Sinusoidal编码具有对称性。
- 随着 $k$ 的增加，内积的结果会直接减少，即会存在远程衰减。
下图是一个 sequence length 为 50 ，$d_{model}$ 为 128 的位置编码
- 每个 token 对应 1 行，每一行在 0～127 的每一个值表示对应的 sin/cos 的值
- 由于 sin/cos 函数的性质，位置向量的每一个值都位于
[-1, 1]之间 - 纵向来看，图的右半边几乎都是蓝色（值十分接近 1）和白色（值十分接近 0）交替，这是因为在 $d_{model}$ 越往后的位置，频率越小，波长越长。这个时候无论 pos 多大，sin 值都为 0，cos 值都为 1
- 越往左边，颜色交替越频繁
- 以 i = 0 为例，随着 pos 增加，就是 sin 函数从 0 到 1 再到 0 再到 -1 的循环
- 以 i = 1 为例，随着 pos 增加，就是 cos 函数从 1 到 0 再到 -1 再到 0 的循环
- 对于每一列，对应到位置向量的 sin/cos 值，都是周期性的
另一个具体的例子：
我们已经知道了三角函数编码的一些优秀特性。
- 在设置合适的𝜃值的前提下，每个位置都能取到唯一的位置编码（绝对性）。
- 一个位置编码可以由另一个位置编码旋转而来（相对性）。
我们接下来看看三角函数编码的一些其它特性。
Sinusoidal 编码的点积只取决于偏移量
两个位置编码的点积(dot product) 仅取决于偏移量 ，也即两个位置编码的点积可以反应出两个位置编码间的距离。
$$ \begin{aligned}
PE_t^T * PE_{t+\triangle t} &= \sum_{i=0}^{\frac{d_{\text{model}}}{2}-1} \left[ \sin(w_i t)\sin(w_i (t+ \triangle t)) + \cos(w_i t)\cos(w_i (t+ \triangle t)) \right] \ &= \sum_{i=0}^{\frac{d_{\text{model}}}{2}-1} \cos\left(w_i \left(t - (t+ \triangle t)\right)\right) \ &= \sum_{i=0}^{\frac{d_{\text{model}}}{2}-1} \cos\left(w_i \triangle t\right) \end{aligned} $$
Sinusoidal 编码的无向性
位置编码的点积是无向的，即 $PE_t^T * PE_{t+\triangle t} = PE_{t+\triangle t} * PE_t^T$
具体如下图所示，图上一些基础信息如下：
- 横轴表示 Δ。
- 纵轴表示固定某个 $PE_t$ 的情况下，改变 Δ 后得到的 $PE_t$ 和 $PE_{t+\Delta t}$ 的内积。
- d 表示不同的 hidden_size
我们可以看到 $PE_t^T∗PE_{t+\Delta t}$ 的变动趋势为：
- 在固定某个 $PE_t$ 的情况下，两个位置编码的内积具有对称性。
- 在固定某个 $PE_t$ 的情况下，两个位置编码的内积具有远程衰减性，即两个位置编码相距越远，点积值越小，但是不具备单调性。
这说明位置向量的点积虽然可以反映相对距离，但因为点积的结果是对称的，它并没有学习到位置的方向性。
Sinusoidal 编码的远程衰减性
远程衰减性（Distant Decay）是指位置编码应能捕获到序列中位置信息和长距离依赖关系（相隔较远的单词之间的关系）。具体来说就是：
- 当两个位置相隔较远时，它们的编码相差较大，相似度较低，这有助于模型区分这些位置，更好地理解和区分序列中的词序关系。
- 而当两个位置相隔较近时，它们的相似度较高，这有助于模型理解它们之间的关系。
远程衰减的先验是：文本是离散时序数据，我们通常会假设文字之间距离越近相关性越强。即位置相近的 Token 平均来说获得更多的注意力，而距离比较远的 Token 平均获得更少的注意力。
在此先验条件下，良好的震荡曲线应该具备如下特点：
- 可以在无限长度下保持连续单调衰减；
- 衰减曲线是非线性的，近距离衰减变化迅速，远距离衰减平缓；
- 在衰减过程中，尽可能少震荡。
还是看这张图，可以看到，对于三角函数位置编码：
- 在固定某个 $PE_t$ 的情况下，两个位置编码的内积具有远程衰减性，即两个位置编码相距越远，点积值越小
- 在长距离下，衰减性并非预期单调下降，而是包括两种震荡：一种是局部窗口内的震荡，一种是随距离收敛到一定波动范围。
Sinusoidal 编码的外推性
长度外推能力（extrapolation，也称 length extrapolation）：如果模型在不经微调的情况下，在超过训练长度的文本上测试，依然能较好的维持其训练效果，我们就称该模型具有长度外推能力。相反，如果大模型由于训练和预测时输入的长度不一致，导致模型泛化能力下降，我们就说模型的外推性存在问题。
如何让位置编码在保证分布内表现的前提下，提升其外推表现，是设计位置编码的一个重要权衡。因为 Sinusoidal 位置编码中的正弦余弦函数具备如下特点，所以理论上也具备一定长度外推的能力。
但是，三角函数编码在外推性上的表现不甚理想。
论文 TENER: Adapting Transformer Encoder for Named Entity Recognition 对这个问题做了精彩的实验和分析。为了更详细探究这个问题，论文假设：
- $x_t,x_{t+\Delta{t}}$ 分别为两个不同位置的原始 token 向量，其尺寸为
(hidden_size, 1) - $PE_t,PE_{t+\Delta{t}}$ 分别为两个不同位置的原始 PE 向量，其尺寸为
(hidden_size, 1) - $W_Q,W_K$分别为尺寸为
(hidden_size, hidden_size)的 Q、K 矩阵
应用了sinusoidal位置编码的q和k点积如下： $$ \begin{align} \boldsymbol{q}t^* \boldsymbol{k}{t+\Delta t} &= \left[ W_Q * \left( \boldsymbol{x}t + \boldsymbol{PE}t \right) \right]^T \left[ W_K * \left( \boldsymbol{x}{t+\Delta t} + \boldsymbol{PE}{t+\Delta t} \right) \right] \ &= \left[ \boldsymbol{x}t^T W_Q^T + \boldsymbol{PE}t^T W_Q^T \right] \left[ W_K \boldsymbol{x}{t+\Delta t} + W_K \boldsymbol{PE}{t+\Delta t} \right] \end{align} $$
从中我们不难发现，经过 attention 层后，位置编码真正起作用的不再是 $PE_t^T∗PE_{t+\Delta t}$ 这两个位置变量都相关的部，而是引入了线性变化后的 $PE_t^T∗W_Q^T * W_k * PE_{t+\Delta t}$ 。
那么在引入这种线性变化后，位置编码还能保持上述所说的绝对性、相对性和远距离衰减性这种优良性值吗？论文用实验的方式来细看这一点。由于 $W^T_Q W_K$ 本质上可以合成一种线性变化，所以我们可以随机初始化一个线性矩阵 W 来代替它。从下图我们可以看出看固定住某个 t 之后，变动 Δ 的点积结果。
- 浅蓝色代表：两个不同位置 sinusoidal 位置编码的点积 $PE_t^T∗PE_{t+\Delta t}$，确实有很好的远程衰减，可以反映出两个位置编码间的距离。
- 黄色和绿色代表：引入两个随机初始化的线性变化的矩阵的点积。PE 矩阵乘以 W 权重矩阵之后，没有明显规律了。表明引入线性变化后，$PE_t^T∗PE_{t+\Delta t}$ 变成了 $PE_t^T∗W*PE_{t+\Delta t}$，原始位置编码的优良性质（远程衰减性等）都受到了极大程度的破坏，其内积所反映的距离因素效果就被破坏。
旋转位置编码 RoPE
理想情况下，一个好的位置编码应该满足以下条件：
- 每个位置输出一个唯一的编码
- 具备良好的外推性
- 任何位置之间的相对距离在不同长度的句子中应该是一致的
这两条比较好理解，最后一条是指如果两个 token 在句子 1 中的相对距离为 k，在句子 2 中的相对距离也是 k，那么这两个句子中，两个 token 之间的相关性应该是一致的，也就是 attention_sample1 (token 1, token 2) = attention_sample2 (token 1, token 2)。
论文中提出为了能利用上 token 之间的相对位置信息，假定 query 向量 $q_m$ 和 key 向量 $k_n$ 之间的内积操作可以被一个函数 $g$ 表示，该函数 $g$ 的输入是词嵌入向量 $x_m$，$x_n$ 和它们之间的相对位置 $m - n$： $$ \left< f_q(x_m, m), f_k(x_n, n) \right> = g(x_m, x_n, m - n) $$ 接下来的目标就是找到一个等价的位置编码方式，从而使得上述关系成立。
假定现在词嵌入向量的维度是两维 $d{=}2$，这样就可以利用上2维度平面上的向量的几何性质，然后论文中提出了一个满足上述关系的 $f$ 和 $g$ 的形式如下：
$$ \begin{aligned} f_q(x_m, m) &= (W_q x_m) e^{i m \theta} \ f_k(x_n, n) &= (W_k x_n) e^{i n \theta} \ g(x_m, x_n, m - n) &= \mathrm{Re}\left[ (W_q x_m) (W_k x_n)^* e^{i(m-n)\theta} \right] \end{aligned} $$
二维 RoPE 理论推导的证明
根据欧拉公式有： $$ e^{ix} = \cos x + i\sin x $$ 则是上述指数函数可以表示为实部为 $\cos x$，虚部为 $\sin x$ 的一个复数，欧拉公式建立了指数函数、三角函数和复数之间的桥梁。
则上述 $f$ 和 $g$ 公式中的
$$ \begin{aligned} e^{im\theta} &= \cos(m\theta) + i\sin(m\theta) \ e^{in\theta} &= \cos(n\theta) + i\sin(n\theta) \ e^{i(m-n)\theta} &= \cos((m - n)\theta) + i\sin((m - n)\theta) \end{aligned} $$
然后我们看回公式： $$ f_q(x_m, m) = (W_q x_m) e^{im\theta} $$
其中 $W_q$ 是个二维矩阵，$x_m$ 是个二维向量，相乘的结果也是一个二维向量，这里用 $q_m$ 表示： $$ q_m = \begin{pmatrix} q_m^{(1)} \ q_m^{(2)} \end{pmatrix} = W_q x_m = \begin{pmatrix} W_q^{(11)} & W_q^{(12)} \ W_q^{(21)} & W_q^{(22)} \end{pmatrix} \begin{pmatrix} x_m^{(1)} \ x_m^{(2)} \end{pmatrix} $$ 然后首先将 $q_m$ 表示成复数形式：
$$ q_m = \left[ q_m^{(1)}, q_m^{(2)} \right] = \left[ q_m^{(1)} + i q_m^{(2)} \right] $$
接着
$$ f_q(x_m, m) = (W_q x_m) e^{im\theta} = q_m e^{im\theta} $$
其实就是两个复数相乘：
$$ q_m e^{im\theta} = \left( q_m^{(1)} + i q_m^{(2)} \right) * \left( \cos(m\theta) + i \sin(m\theta) \right) $$
我们首先来复习一下复数乘法*的性质：
$$ (a + ib) \cdot (c + id) = ac + ibc + iad + i^2 bd = (ac - bd) + i(bc + ad) $$
可以看到，复数乘法也是用的分配律，还有用到了复数的一个性质：$i^2 = -1$
然后就有：
$$ \begin{aligned} q_m e^{im\theta} &= \left( q_m^{(1)} + i q_m^{(2)} \right) * \left( \cos(m\theta) + i \sin(m\theta) \right) \ &= \left( q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta) \right) + i \left( q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta) \right) \end{aligned} $$
将结果重新表达成实数向量形式就是： $$ q_m e^{im\theta} = \left[ q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta),\ q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta) \right] $$ 相信读者看到这里会发现这不就是 query 向量乘以了一个旋转矩阵吗？
$$ \begin{aligned} f_q(x_m, m) &= (W_q x_m) e^{im\theta} = q_m e^{im\theta} \ &= \left[ q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta),\ q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta) \right] \ &= \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \ \sin(m\theta) & \cos(m\theta) \end{pmatrix} \begin{pmatrix} q_m^{(1)} \ q_m^{(2)} \end{pmatrix} \end{aligned} $$
这就是为什么叫做旋转式位置编码的原因。
同理可得 key 向量 $k_n$： $$ \begin{aligned} f_k(x_n, n) &= (W_k x_n) e^{in\theta} = k_n e^{in\theta} \ &= \left[ k_n^{(1)} \cos(n\theta) - k_n^{(2)} \sin(n\theta),\ k_n^{(2)} \cos(n\theta) + k_n^{(1)} \sin(n\theta) \right] \ &= \begin{pmatrix} \cos(n\theta) & -\sin(n\theta) \ \sin(n\theta) & \cos(n\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \ k_n^{(2)} \end{pmatrix} \end{aligned} $$
最后还有个函数 $g$：
$$ g(x_m, x_n, m - n) = \mathrm{Re}\left[ (W_q x_m) (W_k x_n)^* e^{i(m-n)\theta} \right] $$
其中 $\mathrm{Re}[x]$ 表示一个复数 $x$ 的实部部分，而 $(W_k x_n)^*$ 则表示复数 $W_k x_n$ 的共轭。
复习一下共轭复数的定义：
$$ \begin{aligned} z &= a + ib \ z^* &= a - ib \end{aligned} $$
所以可得：
$$ \begin{aligned} W_q x_m &= q_m = q_m^{(1)} + i q_m^{(2)} \ W_k x_n &= k_n = k_n^{(1)} + i k_n^{(2)} \ (W_k x_n)^* &= k_n^* = k_n^{(1)} - i k_n^{(2)} \ e^{i(m-n)\theta} &= \cos((m - n)\theta) + i \sin((m - n)\theta) \end{aligned} $$
继续可得： $$ \begin{aligned} g(x_m, x_n, m - n) &= \mathrm{Re}\left[ (W_q x_m) (W_k x_n)^* e^{i(m-n)\theta} \right] \ &= \mathrm{Re}\left[ \left( q_m^{(1)} + i q_m^{(2)} \right) \left( k_n^{(1)} - i k_n^{(2)} \right) \left( \cos((m - n)\theta) + i \sin((m - n)\theta) \right) \right] \ &= \mathrm{Re}\left[ \left( (q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)}) + i(q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)}) \right) \left( \cos((m - n)\theta) + i \sin((m - n)\theta) \right) \right] \ &= \left( q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)} \right) \cos((m - n)\theta) - \left( q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)} \right) \sin((m - n)\theta) \end{aligned} $$
ok, 接下来我们就要证明函数 $g$ 的计算公式是成立的。
首先回顾一下 attention 操作，位置 m 的 query 和位置 n 的 key 做一个内积操作：
$$ \begin{aligned} f_q(x_m, m) &= \left[ q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta),\ q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta) \right] \ f_k(x_n, n) &= \left[ k_n^{(1)} \cos(n\theta) - k_n^{(2)} \sin(n\theta),\ k_n^{(2)} \cos(n\theta) + k_n^{(1)} \sin(n\theta) \right] \end{aligned} $$ 则内积计算为： $$ \begin{aligned} \left< f_q(x_m, m), f_k(x_n, n) \right> &= \ &\quad \left( q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta) \right) \left( k_n^{(1)} \cos(n\theta) - k_n^{(2)} \sin(n\theta) \right) \ &\quad + \left( q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta) \right) \left( k_n^{(2)} \cos(n\theta) + k_n^{(1)} \sin(n\theta) \right) \ &= q_m^{(1)} \cos(m\theta) k_n^{(1)} \cos(n\theta) - q_m^{(1)} \cos(m\theta) k_n^{(2)} \sin(n\theta) \ &\quad - q_m^{(2)} \sin(m\theta) k_n^{(1)} \cos(n\theta) + q_m^{(2)} \sin(m\theta) k_n^{(2)} \sin(n\theta) \ &\quad + q_m^{(2)} \cos(m\theta) k_n^{(2)} \cos(n\theta) + q_m^{(2)} \cos(m\theta) k_n^{(1)} \sin(n\theta) \ &\quad + q_m^{(1)} \sin(m\theta) k_n^{(2)} \cos(n\theta) + q_m^{(1)} \sin(m\theta) k_n^{(1)} \sin(n\theta) \end{aligned} $$
根据三角函数的性质： $$ \begin{aligned} \sin(a + b) &= \sin a \cos b + \cos a \sin b \ \sin(a - b) &= \sin a \cos b - \cos a \sin b \ \cos(a + b) &= \cos a \cos b - \sin a \sin b \ \cos(a - b) &= \cos a \cos b + \sin a \sin b \end{aligned} $$
上式整理为：
$$ \begin{aligned} \left< f_q(x_m, m), f_k(x_n, n) \right> &= \ &\quad q_m^{(1)} k_n^{(1)} \left( \cos(m\theta) \cos(n\theta) + \sin(m\theta) \sin(n\theta) \right) \ &\quad + q_m^{(1)} k_n^{(2)} \left( -\cos(m\theta) \sin(n\theta) + \sin(m\theta) \cos(n\theta) \right) \ &\quad + q_m^{(2)} k_n^{(1)} \left( -\sin(m\theta) \cos(n\theta) + \cos(m\theta) \sin(n\theta) \right) \ &\quad + q_m^{(2)} k_n^{(2)} \left( \sin(m\theta) \sin(n\theta) + \cos(m\theta) \cos(n\theta) \right) \ &= q_m^{(1)} k_n^{(1)} \cos\left( (m - n)\theta \right) \ &\quad + q_m^{(1)} k_n^{(2)} \sin\left( (m - n)\theta \right) \ &\quad - q_m^{(2)} k_n^{(1)} \sin\left( (m - n)\theta \right) \ &\quad + q_m^{(2)} k_n^{(2)} \cos\left( (m - n)\theta \right) \ &= \left( q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)} \right) \cos\left( (m - n)\theta \right) + \left( q_m^{(1)} k_n^{(2)} - q_m^{(2)} k_n^{(1)} \right) \sin\left( (m - n)\theta \right) \ &= \left( q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)} \right) \cos\left( (m - n)\theta \right) - \left( q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)} \right) \sin\left( (m - n)\theta \right) \ &= g(x_m, x_n, m - n)
\end{aligned}
$$
这就证明上述关系是成立的，位置 m 的 query 和位置 n 的 key 的内积就是函数 g。
写成矩阵的形式： $$ \begin{aligned}
\left< f_q(x_m, m), f_k(x_n, n) \right>
&= \left( \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \ \sin(m\theta) & \cos(m\theta) \end{pmatrix} \begin{pmatrix} q_m^{(1)} \ q_m^{(2)} \end{pmatrix} \right)^T \left( \begin{pmatrix} \cos(n\theta) & -\sin(n\theta) \ \sin(n\theta) & \cos(n\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \ k_n^{(2)} \end{pmatrix} \right) \
&= \begin{pmatrix} q_m^{(1)} & q_m^{(2)} \end{pmatrix} \begin{pmatrix} \cos(m\theta) & \sin(m\theta) \ -\sin(m\theta) & \cos(m\theta) \end{pmatrix} \begin{pmatrix} \cos(n\theta) & -\sin(n\theta) \ \sin(n\theta) & \cos(n\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \ k_n^{(2)} \end{pmatrix} \
&= \begin{pmatrix} q_m^{(1)} & q_m^{(2)} \end{pmatrix} \begin{pmatrix} \cos(m\theta)\cos(n\theta) + \sin(m\theta)\sin(n\theta) & -\cos(m\theta)\sin(n\theta) + \sin(m\theta)\cos(n\theta) \ -\sin(m\theta)\cos(n\theta) + \cos(m\theta)\sin(n\theta) & \sin(m\theta)\sin(n\theta) + \cos(m\theta)\cos(n\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \ k_n^{(2)} \end{pmatrix} \
&= \begin{pmatrix} q_m^{(1)} & q_m^{(2)} \end{pmatrix} \begin{pmatrix} \cos((m-n)\theta) & -\sin((m-n)\theta) \ \sin((m-n)\theta) & \cos((m-n)\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \ k_n^{(2)} \end{pmatrix}
\end{aligned}
$$
推广到 n 维
将2维推广到任意维度，可以表示如下：
$$ f_{{q,k}}(x_m, m) = \boldsymbol{R}{\Theta,m}^d \boldsymbol{W}{{q,k}} x_m \tag{12} $$
内积满足线性叠加性，因此任意偶数维的RoPE，我们都可以表示为二维情形的拼接，即
$$ \boldsymbol{R}_{\Theta,m}^d = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 & \dots & 0 & 0 \
\sin m\theta_1 & \cos m\theta_1 & 0 & 0 & \dots & 0 & 0 \
0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \dots & 0 & 0 \
0 & 0 & \sin m\theta_2 & \cos m\theta_2 & \dots & 0 & 0 \
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \
0 & 0 & 0 & 0 & \dots & \cos m\theta_{d/2} & -\sin m\theta_{d/2} \
0 & 0 & 0 & 0 & \dots & \sin m\theta_{d/2} & \cos m\theta_{d/2}
\end{pmatrix} $$
$$ \Theta = \left{ \theta_i = 10000^{-2(i-1)/d},\ i \in [1, 2, \dots, d/2] \right} $$
将 RoPE 应用到前面的 Self-Attention 计算，可以得到包含相对位置信息的 Self-Attention：
$$ \boldsymbol{q}m^\mathrm{T} \boldsymbol{k}n = \left( \boldsymbol{R}{\Theta,m}^d \boldsymbol{W}q x_m \right)^\mathrm{T} \left( \boldsymbol{R}{\Theta,n}^d \boldsymbol{W}k x_n \right) = x_m^\mathrm{T} \boldsymbol{W}q^\mathrm{T} \boldsymbol{R}{\Theta,n-m}^d \boldsymbol{W}k x_n $$ 其中，$\boldsymbol{R}{\Theta,n-m}^d = \left( \boldsymbol{R}{\Theta,m}^d \right)^\mathrm{T} \boldsymbol{R}{\Theta,n}^d$。
值得指出的是，由于 $\boldsymbol{R}_\Theta^d$ 是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性。
RoPE 的长程衰减
从图中我们可以看到随着相对距离的变大，内积结果有衰减趋势的出现。因此，选择 $\theta_i = 10000^{−2(i-1)/d}$，确实能带来一定的远程衰减性。论文中还试过以 $\theta_i = 10000^{−2(i-1)/d}$ 为初始化，将 $\theta_i$ 视为可训练参数，然后训练一段时间后发现并没有显著更新，因此干脆就直接固定 $\theta_i = 10000^{−2(i-1)/d}$ 了。
LLaMA 中 RoPE 的实现
LLaMA 中对 RoPE 的实现采用复数的公式来计算 $f_q(x_m, m)=(W_qx_m)e^{im\theta}$。该方式速度较快，但不方便后续修改。
具体而言，是把每个向量（Key 或者 Query）两维度一组切分，分成元素对 $(q^1,q^2),(q^3,q^4),…$，每对都解释为二维向量。然后 RoPE 以角度 $\theta_i$ 对每个二维向量 维度对 $(q^i,q^{i+1})$ 分别进行旋转，旋转角的取值与 Sinusoidal 位置编码相同，即采样频率 θ 乘上 token 下标（$m\theta_i=m×base^{−2(i-1)/d}$），旋转完将所有切分拼接，就得到了含有位置信息的特征向量。
$$ \begin{pmatrix}cos(m\theta_1) & -sin(m\theta_1) & 0 & 0 & … & 0 & 0\sin(m\theta_1) & cos(m\theta_1)& 0 & 0 &… & 0 & 0 \0 & 0 & cos(m\theta_2) & -sin(m\theta_2) & … & 0 & 0 \0 & 0 & sin(m\theta_2) & cos(m\theta_2) & … & 0 & 0 \0 & 0 &0 & 0 & … & 0 & 0 \. & . &. & . & .\ \ & . & . \. & . &. & . & \ .\ & . & . \. & . &. & . & \ \ . & . & . \0 & 0 &0 & 0 & … & cos(m\theta_{d/2}) & -sin(m\theta_{d/2}) \0 & 0 &0 & 0 & … & sin(m\theta_{d/2}) & cos(m\theta_{d/2}) \\end{pmatrix} \begin{pmatrix}q_m^{(1)} \q_m^{(2)} \q_m^{(3)} \q_m^{(4)} \.\.\.\q_m^{(d-1)} \q_m^{(d)}\end{pmatrix} $$
|
|
precompute_freqs_cis() 函数会生成旋转矩阵，即给定维度预计算频率θ。θ 完全由 Q、K、V 的向量长度 d 决定。位置 m 对应我们的 query 长度，实际代码中由 max_position_embeddings 参数决定，可以理解为模型支持的最长 query 的长度，因此 max 有了，m 的范围也就有了。结合上面的信息，针对一个固定了最长 query 长度 m 和向量维度 d 的 LLM，我们可以提前将其对应的旋转变换矩阵构造完成。
freqs = torch.outer(t, freqs) 的矩阵如下。
$$
freqs =\begin{bmatrix} 1\theta 1 & 1\theta_2 & 1\theta_3 & … & 1\theta{d/2-1} & 1\theta_{d/2}\ 2\theta_1 & 2\theta_{2} & 2\theta_{3} & … & 2\theta_{d/2-1} & 2\theta_{d/2}\ \vdots & \vdots & \vdots &\ddots & \vdots & \vdots \ \vdots& \vdots& \vdots & \vdots & \vdots & \vdots\ m\theta_{1} & m\theta_{2} & m\theta_{3} & … & m\theta_{d/2-1} &m\theta_{d/2} \ \end{bmatrix}
$$
结合这个 Rd 的变换矩阵，分别执行 cos 和 sin，便可以得到我们计算所需的全位置全维度的变换矩阵。
torch.polar 之后的 freqs 如下：
$$
freqs =\begin{bmatrix} cos(\theta_1)+i \cdot sin(\theta_1) & cos(\theta_2)+i \cdot sin(\theta_2) & … & cos(\theta_{d/2})+i \cdot sin(\theta_{d/2})\ cos(2\theta_1)+i \cdot sin(2\theta_1) & cos(2\theta_2)+i \cdot sin(2\theta_2) & … & cos(2\theta_{d/2})+i \cdot sin(2\theta_{d/2})\ \vdots & \vdots &\ddots & \vdots & \ cos(m\theta_1)+i \cdot sin(m\theta_1) & cos(m\theta_2)+i \cdot sin(m\theta_2) & … & cos(m\theta_{d/2})+i \cdot sin(m\theta_{d/2}) \end{bmatrix}
$$
对于 q 向量和 k 向量应用 rotary embedding
|
|
Attention 的计算：
|
|
GPT-J Style RoPE 实现
另一种实现方法中，不通过引入复数计算，而是考虑到 $\boldsymbol{R}_{\Theta,m}^d$ 的稀疏性，直接拆成以下两部分计算，其中 $\otimes$ 是逐位对应相乘。
$$ \begin{bmatrix}q_m^{(1)} \q_m^{(2)} \q_m^{(3)} \q_m^{(4)} \{\vdots}\q_m^{(d-1)} \q_m^{(d)} \\end{bmatrix} = \begin{bmatrix}cos(m\theta_1) \cos(m\theta_1) \cos(m\theta_2) \cos(m\theta_2) \{\vdots}\cos(m\theta_{d/2}) \cos(m\theta_{d/2}) \\end{bmatrix} \otimes\begin{bmatrix}q_m^{(1)} \q_m^{(2)} \q_m^{(3)} \q_m^{(4)} \{\vdots}\q_m^{(d-1)} \q_m^{(d)} \\end{bmatrix}+ \begin{bmatrix}sin(m\theta_1) \sin(m\theta_1) \sin(m\theta_2) \sin(m\theta_2) \{\vdots}\sin(m\theta_{d/2}) \sin(m\theta_{d/2}) \\end{bmatrix} \otimes\begin{bmatrix}-q_m^{(2)} \q_m^{(1)} \-q_m^{(4)} \q_m^{(3)} \{\vdots}\-q_m^{(d)} \q_m^{(d-1)} \\end{bmatrix} $$
|
|
|
|
Huggingface RoPE 的实现
用的更多的是 Huggingface 的这种 RoPE 实现，和 GPT-J style 的区别不是相邻两个元素为一组，而是 $q_0$ 和 $q_{d/2−1}$ 为一组。对应到区别在 rotate_half ，是否是 interleaved。
这样实现也是合理的，原因参考这个解释： 把 Transformer升级之路：2、博采众长的旋转式位置编码 - 科学空间|Scientific Spaces 的式11右边的向量
[q0,q1,q2,q3,q4,q5]改成[q0,q3,q1,q4,q2,q5]，推倒结果重拍后实际上就和 chatglm 等等方式一样了，就是选择了不同的维度组成复数，可以选前一半和后一半组成复数，而同一个 token 维度内 shuffle 无所谓的。苏神给我的解释是”神经元是无序的（dot attention 做内积，不依赖于维度顺序）
$$ \begin{bmatrix} q_m^{(1)} \ q_m^{(2)} \ {\vdots}\ q_m^{(d/2)} \ q_m^{(d/2+1)} \ {\vdots}\ q_m^{(d)}\ \end{bmatrix} =
\begin{bmatrix} cos(m\theta_1) \ cos(m\theta_2) \ {\vdots}\ cos(m\theta_{d/2}) \
cos(m\theta_1) \ cos(m\theta_2) \ {\vdots}\ cos(m\theta_{d/2}) \ \end{bmatrix}
\otimes
\begin{bmatrix} q_m^{(1)} \ q_m^{(2)} \ {\vdots}\ q_m^{(d/2)} \ q_m^{(d/2+1)} \ {\vdots}\ q_m^{(d)}\ \end{bmatrix}
\begin{bmatrix} sin(m\theta_1) \ sin(m\theta_2) \ {\vdots}\ sin(m\theta_{d/2}) \
sin(m\theta_1) \ sin(m\theta_2) \ {\vdots}\ sin(m\theta_{d/2}) \ \end{bmatrix}
\otimes \begin{bmatrix} -q_m^{(d/2+1)} \ {\vdots}\ -q_m^{(d)} \ q_m^{(1)} \ q_m^{(2)} \ {\vdots}\ q_m^{(d/2)} \ \end{bmatrix} $$
其 rotate_half 实现如下：
|
|
计算出 query_states 和 key_states 之后，可以直接 apply_rotary_pos_emb
|
|
具体计算 sin 和 cos 函数如下，本质上就是拼出
m * [theta_0, theta_1, ..., theta_(dim/2-1), theta_0, theta_1, ..., theta_(dim/2-1)]- 然后直接每一位直接计算 sin 和 cos 就行
|
|
Meta 版本 RoPE 与 Huggingface 版 RoPE 权重转换
https://github.com/huggingface/transformers/issues/30872#issuecomment-2118683720
|
|
RoPE 长度外推
虽然 RoPE 理论上可以编码任意长度的绝对位置信息，并且 sin/cos 计算就能将任意长度的相对位置信息表达出来，但是实验发现 RoPE 仍然存在外推问题，即测试长度超过训练长度之后，模型的效果会有显著的崩坏，具体表现为困惑度（Perplexity，PPL） 等指标显著上升。对此，就有很多人提出了如何扩展 LLM 的 Context 长度，比如有人实验了“位置线性内插“的方案，显示通过非常少的长文本微调，就可以让已有的 LLM 处理长文本。几乎同时，Meta 也提出了同样的思路，在论文 Extending Context Window of Large Language Models via Positional Interpolation 中提出了位置线性内插 Position Interpolation，PI）。随后，又有人提出了 NTK-aware Scaled RoPE，实现了不用微调就可以扩展 Context 长度的效果。
从 Sinusoidal 编码到旋转位置编码 RoPE，$\theta_i$ 的取值一直都是 $\theta_i = 10000^{−2(i-1)/d}$，也就是 base = 10000。而最新的位置线性内插和 NTK 都打破了这一传统。
位置线性内插 Position Interpolation
传统的 RoPE 算法在超出最大距离后，PPL 就会爆炸，因此直接推广的效果一定是很差的，因此 Position Interpolation 绕过了推广的限制，通过内插的方法，如下图所示。
上图是关于“位置线性内插（Position Interpolation）”方法的图示说明。对于预训练 2048 长度，外推 4096 长度：
左上角图示：
- 这部分是 LLM 预训练的长度。
- 蓝色点代表输入的位置索引，它们都在 0 到 2048 的范围内，即预训练范围内。
- 这意味着给模型提供的输入序列长度没有超过它预训练时的最大长度。
右上角图示：
- 这部分展示了所谓的“长度外推”（Extrapolation）情况。
- 在这种情况下，模型被要求处理位置索引超出 2048 的情况，直到 4096（由红色点表示）。
- 这些位置对于模型来说是“未见过的”，因为在预训练时它们并没有涉及。
- 当模型遇到这种更长的输入序列时，其性能可能会受到影响，因为它没有被优化来处理这种情况。
左下角图示：
- 这部分展示了“位置线性内插”方法的应用。
- 为了解决模型处理更长序列时的问题，研究者们提出了一个策略：将位置索引本身进行下采样或缩放。
- 具体来说，他们将原本在
[0, 4096]范围内索引（由蓝色和绿色点表示）缩放到[0, 2048]的范围内。 - 通过这种方式，所有的位置索引都被映射回了模型预训练时的范围内，从而帮助模型更好地处理这些原本超出其处理能力的输入序列。
NTK-Aware Scaled RoPE
参考资料
-
No backlinks found.