简介

因为RoPE的优异性能，其已成为各种大模型中位置编码的首选，包括多模态模型；在一些多模态模型或视频理解模型中，甚至会用到多维度RoPE。虽然RoPE已广泛应用，之前也看了不少针对其原理解析的文章，特别是苏神的推导帖，可能是我个人能力问题，也可能是太过于追求原理的问题，对RoPE一直没有理解透彻，总感觉差那么一点。本文会适当弱化RoPE为什么会有效的原理推导，将重点放在RoPE应该如何实现以及具体的实现步骤和一些不同实现方式的解析。

RoPE的核心思想是通过旋转操作将位置信息融入到词向量。在复数域中，一个复数由实部和虚部组成，故一个复数可以看作复数域中的向量。将一个复数与复数 $e^{i\theta }=\cos \theta +i\sin \theta$相乘时，相当于在复平面上绕原点旋转了 $\theta$度角。假设有一个复数 $v=x+iy$，按上述说明，将 $v$与 $e^{i\theta }$相乘就是将 $v$绕原点旋转了 $\theta$度角，可以通过公式推导验证：
$$\begin{array}{rl}{v}^{\prime }& =v\ast e^{i\theta }\\ & =(x+iy)\ast (\cos \theta +i\sin \theta )\\ & =\cos \theta x+i\sin \theta x+i\cos \theta y-\sin \theta y\\ & =(\cos \theta x-\sin \theta y)+i(\sin \theta x+\cos \theta y)\end{array}\text{\hspace{1em}(1)}$$

通过公式(1)有 $x^{\prime }=\cos \theta x-\sin \theta y,y^{\prime }=\sin \theta x+\cos \theta y$。从二维向量的角度看，变换前后的复数 $v,v^{\prime }$其实均可以看作列向量 $\left[\begin{array}{c}x\\ y\end{array}\right]，\left[\begin{array}{c}{x}^{\prime }\\ y^{\prime }\end{array}\right]$，公式(1)表示 $v^{\prime }$是由 $v$旋转而来，且恰好对应了二维向量空间的旋转矩阵，如下所示：
$$\begin{array}{c}\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]\end{array}\text{\hspace{1em}(2)}$$

如公式(2)所示，在二维向量空间中，将一个向量绕原点旋转 $\theta$度就是乘以一个旋转矩阵，即 $R(\theta )=\left[\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]$。通过上述推导可知，二维向量空间中的旋转和复数乘法在进行旋转时是等价的。

RoPE是将隐向量序列中不同位置的token的特征向量与不同的旋转角对应的旋转矩阵相乘，即将不同位置token的特征向量旋转不同的角度。注意力计算的主要部分是token向量之间的点积运算，可以通过推导展示RoPE是如何在自注意力计算中引入相对位置信息的。

假设目前token的维度为2，有序列中m处的查询向量 $q_m=\left[\begin{array}{c}{q}_{m,1}\\ q_{m,2}\end{array}\right]$和n处的键向量 $k_n=\left[\begin{array}{c}{k}_{n,1}\\ k_{n,2}\end{array}\right]$，设旋转基础角度为 $\theta$，序列中不同位置token的旋转角由位置索引与基础角度相乘得到，即上述两个向量的旋转角度为别是 $m\theta ,n\theta$，RoPE就是给不同位置的token向量与对应位置的旋转角的旋转矩阵相乘，基于前文的推理，有，
$$\begin{array}{c}{q}_m^{\prime }=\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]=R(m\theta )q_m\end{array}\text{\hspace{1em}(3)}$$
$$\begin{array}{c}{k}_n^{\prime }=\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]=R(n\theta )k_m\end{array}\text{\hspace{1em}(4)}$$

旋转后的两个向量进行以下点积计算：
$$\begin{array}{rl}{q}_m^{{}^{\prime }T}{k}_n^{\prime }& =(R(m\theta )q_m{)}^T\cdot R(n\theta )k_n\\ & =q_m^T\cdot R(m\theta {)}^T\cdot R(n\theta )\cdot k_n\\ & =q_m^T\cdot \left[\begin{array}{cc}\cos (m\theta )& \sin (m\theta )\\ -\sin (m\theta )& \cos (m\theta )\end{array}\right]\left[\begin{array}{cc}\cos (n\theta )& -\sin (n\theta )\\ \sin (n\theta )& \cos (n\theta )\end{array}\right]\cdot k_n\\ & =q_m^T\cdot \left[\begin{array}{cc}\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{array}\right]\cdot k_n\\ & =q_m^T\cdot \left[\begin{array}{cc}\cos ((n-m)\theta )& -\sin ((n-m)\theta )\\ \sin ((n-m)\theta )& \cos ((n-m)\theta )\end{array}\right]\cdot k_n\\ & =q_m^T\cdot R((n-m)\theta )\cdot k_n\end{array}\text{\hspace{1em}(5)}$$

上述公式(5)推导过程要使用高中数学知识–三角函数和差定理。可以看到，两个旋转不同角度后向量的点积与它们之间旋转的角度之差有关系，是一种相对位置的体现，即体现了RoPE给序列中不同位置token引入了相对位置信息。

注意：上述旋转均是逆时针旋转！

拆解

上述只是对RoPE的原理进行了简短解释，尽可能直白地说明RoPE原理，但上述讨论局限在二维空间内，而LLMs模型中token向量的维度都很大，接下来对常规的RoPE实现进行拆解，将整个实现过程解剖出来。

1. 为了沿用二维旋转矩阵的性质，RoPE会将一个token的高维向量看作多个复数，假设向量特征维度为 $d$(基本所有LLMs中的特征维度数是偶数)，那么会将一个$d$维的token特征向量视为 $d/2$个复数

• 假设一个token向量为 $[1,2,3,4]$，那么 $[1,2]$是第一个复数，$[3,4]$是第二个复数
• 同一个token中的不同复数对的旋转角度不同
• 基础旋转角 ${\theta }_i=10000^{-2i/d}=\frac{1}{10000^{\frac{i}{\frac{d}{2}}}}$；注意，同一token向量中不同位置复数的旋转角是由 $i$来决定的。上文已经说到，一个token向量会分成 $d/2$个复数， $i\in [0,d/2)$，所以 ${\theta }_i$就对应第 $i$对复数的旋转角
  • 10000是一个超参数

2. 同一序列中不同位置的token向量的旋转角也不同，就是基础旋转角与token对应的序列索引乘积，即一个token在序列中的索引为m，那么其对应的最终的旋转角为 ${\theta }_i\cdot m$。

• 实现时一般会先将 ${\theta }_i$计算出来，再基于序列长度获取不同位置token索引的一维向量，然后两者相乘，得到不同位置token中不同复数维度的准确旋转角度

3. 将输入沿着特征维度，按顺序每两个为一组作为一对复数进行拆分，向量维度从[batch_size, seq_len, dim]–>[batch_size, seq_len, dim//2, 2]

• 最后一个维度中，第一组向量为所有复数的实部，第二组向量为所有复数的虚部

4. 按照二维矩阵中的计算方式，进行旋转矩阵乘法，然后将最终的实部和虚部重新合并为原来的维度[batch_size, seq_len, dim]，即完成了RoPE的应用

上述对RoPE的具体操作过程进行了叙述，以下是一种较原始的朴素实现：

import torchdef get_rotary_matrix(seq_len, dim, base=10000):"""生成RoPE的旋转矩阵"""# 生成不同频率的正弦和余弦值theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))  # shape为[dim//2]# 生成位置索引position = torch.arange(seq_len).float()  # shape为[seq_len]# 计算每个位置和维度对应的角度theta = torch.outer(position, theta)  # 计算外积，其中第(i, j)个元素是position[i] * theta[j]；shape为[seq_len, dim//2]# 计算正弦和余弦值cos = torch.cos(theta)  # shape为[seq_len, dim//2]sin = torch.sin(theta)  # shape为[seq_len, dim//2]return cos, sindef apply_rotary_embedding(x, cos, sin):"""应用旋转位置编码"""# 假设x的形状为[batch_size, seq_len, dim]# 将向量视为复数，每两个维度一组x_reshape = x.view(*x.shape[:-1], -1, 2)  # shape为[batch_size, seq_len, dim//2, 2]，即沿着特征维度拆分# 构建正弦和余弦矩阵，使其与x_reshape形状匹配cos_expanded = cos.view(1, cos.shape[0], cos.shape[1], 1)  # shape为[1, seq_len, dim//2, 1]sin_expanded = sin.view(1, sin.shape[0], sin.shape[1], 1)  # shape为[1, seq_len, dim//2, 1]# 旋转操作（复数乘法）# [x_real, x_imag] * (cos + i*sin) = [x_real*cos - x_imag*sin, x_real*sin + x_imag*cos]x_out_1 = x_reshape[:, :, :, 0:1] * cos_expanded - x_reshape[:, :, :, 1:2] * sin_expandedx_out_2 = x_reshape[:, :, :, 0:1] * sin_expanded + x_reshape[:, :, :, 1:2] * cos_expanded# 合并结果x_out = torch.cat([x_out_1, x_out_2], dim=-1)  # shape为[batch_size, seq_len, dim//2, 2]return x_out.view(*x.shape)# 示例用法
def apply_rope(x):"""对输入向量应用RoPE位置编码"""batch_size, seq_len, dim = x.shapecos, sin = get_rotary_matrix(seq_len, dim)return apply_rotary_embedding(x, cos, sin)# 测试代码
if __name__ == "__main__":# 创建一个随机输入张量batch_size, seq_len, dim = 2, 10, 512x = torch.randn(batch_size, seq_len, dim)# 应用RoPEx_with_rope = apply_rope(x)print(f"输入形状: {x.shape}")print(f"输出形状: {x_with_rope.shape}")

一些tricks

• 在对输入的嵌入维度上应用RoPE时，当特征维度很长，或者希望节约资源资源时，可以不对所有维度全部应用，而是可以设置一个比例，即只在前 $d_{rope}$维度上应用RoPE，后面保持原始数值不变
• 在计算 ${\theta }_i$时原始论文使用的base为10000，但可以对其进行缩放，实现长度外推；如增加base，可以处理训练过程中未见过的长序列