RoPE(Rotary Position Embedding)这几年已经是大语言模型位置编码的事实标准。如果你翻过主流开源模型的权重文件或推理代码,大概率会看到一堆预计算好的 cos 和 sin 缓存表;如果你手写过 Transformer 的注意力层,可能也纠结过——位置向量为什么要用“旋转”而不是直接加?如果你此刻看到“RoPE 的数学表达式”这个标题,说明正在追根溯源。这篇文章就走一遍完整推导,从复平面上的二维旋转开始,逐步推到高维分块旋转矩阵,再给出可直接复用的 PyTorch 实现和工程调优经验,把每个“为什么”都拆开说明白。
适合的人群主要有三类:一是准备自己实现位置编码的算法工程师,二是想深入理解 LLaMA 系模型内部结构的爱好者,三是在做长文本外推或位置编码改造的研究型同学。读完后你不仅能默写出 RoPE 的核心公式,还能避开实现里那些很容易踩但不查源码根本看不出来的坑。
1. 从绝对位置到相对位置:为什么 RoPE 是刚需
1.1 Transformer 的序列记忆困境
Transformer 原版架构中注意力计算本身是置换等变的,简单说就是你把序列顺序打乱,两个 token 之间的注意力分数不会变,因为 Q 和 K 的每个元素都来自 token 的 embedding,并没有位置维度。这对理解句子含义是致命的,“小明打了小狗”和“小狗打了小明”如果词向量相同,模型会一视同仁。
最早的解决办法是给每个 token 加上一个位置相关向量,也就是绝对位置编码。原版 Transformer 使用不同频率的正弦和余弦函数,每个位置的编码向量形状固定,与词向量直接相加。后来很多模型改成可学习的绝对位置编码,让模型在训练中自己调整位置向量。这类方案实现简单,在训练长度内效果也没问题,但一个核心痛点是外推性不足:训练时最长见过 1024 个位置,到了 2048 就不知道怎么给新位置赋值,强行用更高频率的三角函数会不稳定,用固定向量又会跨越训练分布。
相对位置编码的思路是:不关心 token 在第几个位置,只关心两个 token 之间的距离。这更符合语言直觉,“左侧第二个词”这种信息比“位置 17”丰富得多。但相对位置编码实现起来复杂,有的模型在 Q、K 上分别构造位置偏差,有的在 softmax 之后加一个距离相关偏置,工程上不算干净,而且容易被 attention 结构绑死,难以通用。
1.2 RoPE 的切入点:旋转即位置
RoPE 的做法是在 Q 和 K 向量上施加一个与位置相关的旋转矩阵,让向量的模长不变,方向随位置改变。这样 Q 和 K 在计算点积时,绝对值角度会自动编码两个位置的相对关系。你不需要学习任何位置参数,只需要定义好旋转角度即可。
为什么是旋转而不是加法?旋转有一个非常好的代数性质:两个经过旋转后的向量点积,可以化为原始向量点积再乘一个只与位置差相关的旋转矩阵。这意味着绝对位置信息在点积中表现为相对位置差,这非常优雅。另外,旋转矩阵是正交矩阵,保持向量的范数,所以理论上不会对词的语义破坏太大,词向量依然保留原有信息,只是方向被调整。
这个设计直接解决了两个问题:第一,不增加额外可学习参数;第二,模型能看到相对位置,因为点积结果只依赖位置差,而非绝对位置。这为后续的长文本外推和模型结构简化都打下了基础。
2. 旋转矩阵的数学原理:复数视角下的位置编码
2.1 从二维平面上的旋转开始
先看最简单的二维情况。假设一个二维向量 (x, y),我们希望把它逆时针旋转角度 θ。用复数来表示非常简洁:把 (x, y) 看作复数 z = x + iy,乘以 e^(iθ) = cosθ + i sinθ 后,新的复数 z' = z * e^(iθ) = (x + iy)(cosθ + i sinθ) = (x cosθ - y sinθ) + i(x sinθ + y cosθ)。对应的实部和虚部分开写就是:
- x' = x cosθ - y sinθ
- y' = x sinθ + y cosθ
写成矩阵形式就是:
[ x' ] [ cosθ -sinθ ] [ x ] [ y' ] = [ sinθ cosθ ] [ y ]这个旋转矩阵是正交的,逆时针旋转 θ 可以用角度叠加表示:旋转了位置 m 和位置 n,两个旋转矩阵的差就是旋转 (n - m)。
注意这里的 θ 是和位置相关的。如果令 θ_m = m * θ_base,其中 m 是绝对位置下标,θ_base 是基础角度,那么对于位置 m 的向量,旋转角度就是 m 倍的基础角。这就是“旋转位置”名字的由来。
2.2 从二维到高维:分块旋转矩阵
实际模型中 Q 和 K 向量是高维的,比如 4096 或者 8192。我们不能直接把整个向量旋转一个角度,那样所有维度都同一方式移动,表达能力不够。RoPE 的做法是把维度两两一组,d 维向量分成 d/2 个二维子向量,每个子空间使用不同的基础频率。
具体来说,假设维度下标从 0 到 d-1,我们把第 2i 和第 2i+1 视为一个二维子空间(这里 i 从 0 开始),每个子空间分配一个频率 ω_i = 10000^(-2i/d)。于是位置 m 对第 i 个子空间的旋转角度是 m * ω_i。
整个位置 m 的旋转矩阵是一个分块对角矩阵,每个对角子块是二维旋转矩阵:
R_m = diag( R(mω_0), R(mω_1), ..., R(mω_(d/2-1)) )其中每个子块 R(mω_i) 是 2x2 旋转矩阵:
[ cos(mω_i) -sin(mω_i) ] [ sin(mω_i) cos(mω_i) ]因为分块对角矩阵可以作用在向量上,等价于对每个二维子空间单独旋转,所以 RoPE 的完整公式可以写成:
q_m = R_m q k_n = R_n k这里 q 和 k 是原始的词向量经过线性投影后的向量,q_m 和 k_n 是融入位置信息后的向量。后续注意力计算时,使用 q_m 和 k_n 做点积。
2.3 核心表达式整理
将上面的推导整合,给出 RoPE 的精确定义:
- 基础频率集合:ω_i = 10000^(-2i/d),i 从 0 到 d/2 - 1
- 位置 m 对应的旋转角度:θ_{m,i} = m * ω_i
- 对每个特征维度,把向量拆成两半:第一半记为奇数维度,第二半记为偶数维度(或者反过来,取决于实现)
- 旋转后新向量:x'2i = x_2i cos θ{m,i} - x_{2i+1} sin θ_{m,i},x'{2i+1} = x{2i+1} cos θ_{m,i} + x_2i sin θ_{m,i}
常见的实现方式是把 Q 向量 reshape 成 [batch, seq, num_heads, d/2, 2],然后把这个最后一维的两个值视为复数坐标,左乘旋转矩阵。旋转操作完成后,整个注意力分数为:
score = q_m^T k_n = (R_m q)^T (R_n k) = q^T R_m^T R_n k这里 R_m^T R_n = R_{n-m},因为旋转矩阵是正交且角度可加的,所以分数只和相对位置差 (n-m) 有关。这个结果是 RoPE 的核心优点。
3. 为什么 RoPE 有效:性质推导与直觉解释
3.1 点积中的相对位置:如何把绝对位置变成相对位置
上一节已经给出了结果,这里再单独推导一遍,因为这是最容易忽略的关键转折点。先写出注意力分数的展开:
q_m^T k_n = q^T R_m^T R_n k由于 R_m 和 R_n 都是分块对角旋转矩阵,根据旋转矩阵的性质,R_m^T = R_{-m},且 R_m^T R_n = R_{n-m}。所以:
q_m^T k_n = q^T R_{n-m} k注意这里 R_{n-m} 只依赖于两个位置的差,而不是位置 m 和 n 各自的值。这意味着模型容量不需要为每个绝对位置保存独立向量,只要相对距离一致,位置关系就一致。这极大提升了模型的泛化能力,尤其对训练中没有见过的较长位置组合,模型依然能利用学习到的频率插值来估计注意力强度。
你可以把旋转矩阵想象成表盘。向量是表针,位置是表针指向的角度。两个表针之间的夹角由位置差决定,绝对时间不直接影响夹角。如果你只观察两个表针之间的夹角,就能判断两个事件的相对先后顺序。这正是 RoPE 的直觉。
3.2 远程衰减:高频低频与旋转速度
RoPE 的另一个有意思的性质是注意力分数随距离增大而衰减。为什么?因为维度分成不同频率,高频维度旋转很快,低频维度旋转很慢。对于一个特征向量来说,在距离差 (n-m) 较大的情况下,部分高频分量的 cos 和 sin 值已经来回变化,平均效应导致整体内积的绝对值变小。而低频分量对长距离变化不敏感,所以模型能保持一定的长距离注意力能力。总体来看,注意力分数随距离增加呈现一个振荡衰减的包络。
这个衰减不是人为强加的,而是旋转频率结构自然带来的。你可以类比收音机调频:不同频率的载波携带不同距离尺度的信息。高频载波适合编码精细的邻近关系,低频载波适合编码长距离依赖。RoPE 通过多频率混合实现“多尺度位置敏感”。
数学上,这个衰减函数可以写成 sinc 相关形式。具体来说,如果没有学到特殊的 q 和 k 结构,两个距离较远的位置计算出来的平均点积会趋向于零。这让模型自动倾向于关注近邻信息,而长距离信息则依赖低频分量保留。很多后续研究也验证了 RoPE 模型在中距离依赖上表现好,同时不会完全忽略长距离。
3.3 与其他位置编码的对比
为了更直观,我用一个表格把常见编码方式放在一起对比:
| 编码方式 | 绝对/相对 | 特点 | 典型问题 |
|---|---|---|---|
| 正弦绝对位置编码 | 绝对 | 无参数,结构固定 | 外推表现差,无法微调相对距离 |
| 可学习绝对位置编码 | 绝对 | 参数随位置增多 | 外推零能力,需要截断或插值 |
| 相对位置偏置 | 相对 | 直接学习距离偏置 | 增加了模型复杂度,不同注意力头不共享位置信息 |
| RoPE | 相对(隐式) | 不引入参数,天然相对位置 | 表达式略复杂,高频部分在长上下文会有插值问题 |
RoPE 最吸引人的是它把相对位置信息隐含在旋转矩阵中,不需要在 attention 计算图外额外加偏置,和当前主流的高效注意力(比如 Flash Attention)天然兼容。这也是为什么很多大模型选择 RoPE 而不是传统相对位置编码。
4. 工程实现:从公式到 Tensor 代码
4.1 预处理旋转角度:生成 cos 和 sin 缓存
实际在 PyTorch 里实现 RoPE 一般分两步:生成角度的 cos 和 sin 缓存,然后在 forward 时对 Q 和 K 应用旋转。先看角度生成。
import torch def precompute_freqs_cis(dim: int, max_seq_len: int, base: float = 10000.0): # dim 是单个 attention 头的维度,需要是 2 的倍数 i = torch.arange(0, dim, 2).float() # 论文中是 0, 2, 4, ..., dim-2 freqs = 1.0 / (base ** (i / dim)) # 创建位置索引 0 到 max_seq_len-1 m = torch.arange(max_seq_len).float() # 外积:每个位置 m,每个频率一个角度 m * freq angles = m[:, None] * freqs[None, :] # [max_seq_len, dim/2] # 对每个角度,cos 和 sin 需要扩展到完整的 dim 维度(奇偶分开) # 标准实现中会把 cos 和 sin 交替排列,也可以写成两个半维 cos_cached = torch.cos(angles) sin_cached = torch.sin(angles) # 扩展为 [max_seq_len, 1, dim] 以便broadcast # 具体处理和后续应用方式有关,这里先返回原始角度信息 return cos_cached, sin_cached这里需要说明的是,论文里的频率指数使用的是 2i/d,其中 i 从 0 到 dim/2-1。由于 i 是按 2 的步长取的,所以 i 的实际值相当于“维度对编号”。这个细节很多人第一次看会搞混。如果你的 dim 不是 2 的倍数,RoPE 无法使用,一般要求 dim 是偶数。
生成 cos 和 sin 的缓存后,注意角度 shape 是 [max_seq_len, dim/2],我们需要把它扩展成 [max_seq_len, dim]。扩展方式有两种:一种是按维度顺序预成倍,比如把 cos_cached 用 repeat_interleave 得到每个维度对应的值,另一种是后续在 apply 时把向量 reshape 成两半再广播。我推荐后者,下面会讲到。
4.2 apply_rotary_pos_emb 的核心代码
最常见的实现方式是把 Q 和 K 的最后一维分成两半,前半作为旋转的“实数坐标”,后半作为“虚数坐标”,再按复数形式乘以旋转因子。这里有多种等价写法,我给出一种清晰且不容易出错的版本:
def rotate_half(x: torch.Tensor): # x shape: [..., dim] # 把最后一维分成前后两半 x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] # 交叉组合,对应公式:x'_2i = x_2i cos - x_2i+1 sin # 这里把前半后半看成交替维度 return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, position_ids): # q, k shape: [batch, seq, heads, dim] 或 [batch, heads, seq, dim] # cos, sin shape: [seq, dim] 或 [batch, seq, dim] # position_ids: [batch, seq] # 从位置 id 中获取 cos 和 sin cos_k = cos[position_ids] # [batch, seq, dim] sin_k = sin[position_ids] # [batch, seq, dim] # 旋转后的 q 和 k q_embed = (q * cos_k) + (rotate_half(q) * sin_k) k_embed = (k * cos_k) + (rotate_half(k) * sin_k) return q_embed, k_embed这里的核心是rotate_half的实现。为什么是(-x2, x1)?回到二维旋转公式:
x'_1 = x_1 cosθ - x_2 sinθ x'_2 = x_2 cosθ + x_1 sinθ如果我把向量分成前后两半,前半表示所有偶数维度,后半表示所有奇数维度,那么旋转后的前半相当于:原前半乘以 cos 减去原后半乘以 sin;旋转后的后半相当于:原后半乘以 cos 加上原前半乘以 sin。上面代码中(-x2, x1)恰好就是旋转后的“补充项”:
q * cos_k得到前半方向上的值rotate_half(q) = (-x2, x1),再乘以sin_k,就完成了后半和前半的交叉加法
这个写法比torch.view(dim // 2, 2)再重新拼接要快一些,也更容易理解。
4.3 维度排列陷阱与常见错误
实现 RoPE 最容易踩的坑是维度排列。如果你使用[batch, seq, heads, dim]布局,cos 和 sin 的 shape 是[seq, dim],你需要先按 position_ids 做索引,再在头维度上广播。如果你使用[batch, heads, seq, dim]布局,position_ids 索引后的 shape 也要保持匹配。
第二个常见错误是旋转 Q 和 K 时,位置索引没有对齐。Q 的第 m 个位置要使用角度 m,K 的第 n 个位置要使用角度 n,这本来无比自然,但在实现 attention 时,有可能因为缓存或转置导致位置错位,尤其是在加入 KV Cache 做增量推理时。增量阶段需要把当前新 token 的位置索引传给 apply 函数,而不能沿用训练时完整序列的位置索引,否则旋转角度会叠加错误。
第三个常见错误是 split 方向错误。有的实现是按[..., 0::2]和[..., 1::2]切分,这和[..., :dim//2]切分在数学上表示相同,但代码写法不同,容易少用一个维度或反转。建议固定一种风格,并且用非常简单的正弦信号测试验证旋转的正确性。
5. 实战中的扩展与调优
5.1 上下文外推:NTK 缩放与频率调整
RoPE 的一个著名难点是长文本外推。训练时上下文长度是 L,测试时如果直接给超过 L 的位置 m,由于角度超出缓存范围,模型没有见过这些角度组合,注意力通常会退化。原版位置编码在这方面表现更差,而 RoPE 因为有个天然的频率结构,可以通过插值手法缓解。
位置插值方法很简单:把位置索引压缩,比如原来 4096 位置的上下文,测试想扩展到 8192,可以把位置 m 缩小为 m / 2,使角度范围保持在训练范围内。这相当于把所有 token 位置折中,模型被迫适应更缓慢的角度变化。这种方法效果尚可,但会导致高频信息被压缩,模型可能失去对相邻 token 的敏感度。
NTK-aware 缩放是另一种思路。它不改变角度整体范围,而是调整基础频率 base。通过把 base 从 10000 增大,各个维度的频率下降,整体角度增长速度减慢。这等价于在频域做了一种“非线性插值”,高频维度保持一定锐度,低频维度被压缩。实现很简单,把计算 freqs 的 base 测试后替换为新的 base。常见的转换公式是:
base_new = base * (scale) ^ (dim/(dim-2))其中 scale 是长度扩展比例。这个经验公式在很多开源项目里被验证能稳定地让模型外推到长文本外。我个人推荐在微调阶段使用 NTK 方式,而不是直接位置插值,因为前者对模型原始分布改动更小。
5.2 高维 RoPE:3D 旋转位置编码
RoPE 虽然最初用于一维文本序列,但后来在多模态模型里也被扩展到了图像和视频。图像是二维位置,视频是三维位置(时间加空间)。一个直接的扩展方式是把二维坐标编码拆成行和列两个方向,分别用一维 RoPE 计算,然后组合起来。这种做法已经能看到很多视觉 Transformer 模型在采用,但它们往往直接把位置 id 拼接后通过一维 RoPE 处理,其实丢失了二维的旋转直感。
更严格的高维 RoPE 是使用 3D 旋转矩阵。三维中的位置由 (x, y) 或 (x, y, t) 表示,每个子空间可能包含三个维度而不是两个。这时需要采用三维旋转方程,比如向量的各分量可以按三个轴分别旋转,或者使用四元数。工程上通常不直接实现三维旋转矩阵,而是把不同轴上的位置信息分别作用到不同子空间,比如某些维度编码 x,另一些维度编码 y,最后合并。这样实现简单,且能保留旋转的相对位置性质。
5.3 优化技巧:预计算与缓存
推理性能是 RoPE 的关键挑战。每个 token 都要计算旋转,如果每次前传都去生成 cos 和 sin,开销会非常明显。实际部署时应该在模型加载阶段预计算最长上下文的 cos/sin 缓存,并对位置索引做索引查询。如果启用了 KV Cache,只需要缓存past_length之后的新 token 位置对应的 cos/sin 切片。
低精度推理(比如 FP16 或 INT8)时,cos 和 sin 的精度损失通常可以忽略,但要注意角度 m * ω_i 随着 m 增大可能出现浮点溢出。比如位置 m 达到 200 万,基础频率 1.0 时,角度可能超过 1e6,FP16 下的精度就崩了。所以很多长文本实现会先把位置索引缩放到一个安全范围,再查表。
此外,FlashAttention 通常支持自定义位置编码融合算子。如果你做的是显存受限的推理,建议把 RoPE 的旋转操作写到 attention kernel 里,避免产生额外的中间张量。这一步优化能显著减少显存占用,对长序列尤其重要。
6. 常见问题与排查技巧实录
6.1 典型错误速查表
| 现象 | 可能原因 | 解决方案 |
|---|---|---|
| 训练 loss 正常,但测试时换长度效果崩塌 | 位置索引超出预计算范围,或缓存表长度不够 | 动态扩展 cos/sin 缓存,或用 NTK 缩放 |
| 注意力分数完全不随距离变化 | 旋转矩阵未应用到 Q 或 K,或者位置索引恒为 0 | 检查 forward 中是否同时传入了 position_ids |
| 相同 token 在不同位置的输出几乎相同 | 频率 base 设置过大,导致各维度旋转角度过小 | 降低 base 或在二维子空间中验证角度差异 |
| 多卡训练时 loss 出现微妙波动 | 位置缓存在不同 rank 上生成方式不一致 | 确保 all ranks 用相同随机种子或相同基础频率生成 |
| 推理阶段 kv cache 在第一个 token 时计算正确,后续出错 | 增量阶段 position_ids 未基于累计长度更新 | 用past_length + new_token_index计算新位置 |
排查位置编码问题最简单的方法是用一个双正弦或线性序列做 sanity check。比如构造一个简单模型,仅保留 attention 层,输入一对相同词向量、不同距离的 token,观察注意力输出是否在距离上单调衰减。如果输出不满足预期,直接打印中间旋转后的 Q 和 K,检查角度是否按公式变化。
6.2 我的实操心得
我自己做早期实验时,花了半天时间对着一模一样的 loss curve 发愣。后来发现错误出在position_ids是形状[batch, 1]而不是[batch, seq],导致每个序列里所有 token 都用同一个位置 0 来取 cos 和 sin。这种错误不会让训练崩溃,因为模型可以强行忽略位置信息,但最后测试长上下文时完全没法看。
还有一个容易被忽略的细节是在 attention mask 里用了-inf,但 RoPE 的旋转角度在长距离时可能产生比较大的正负波动,导致 softmax 之前某些 token 的 logits 即使在 masked 位置仍然有非零值,如果没有配合正确的 mask 一个 token 一个 token 地填,输出会莫名其妙。建议在调试时先用无 mask 的简单序列做单测,再逐步加复杂性。
如果你在改写别人的实现,一定要确认他们的“维度对编号”是不是从 0 开始。某些开源实现里i是直接取torch.arange(0, dim, 2),然后指数用i/dim,也就是0/dim, 2/dim, 4/dim,这样频率集合是单调递减。有些实现会用i从 0 到dim/2-1,再乘 2,两种写法数学等价,但如果你混着抄就容易出 bug。
另外,我建议把 RoPE 的公式写成注释放在代码里,因为几个月后回来维护,连原作者都可能忘记为什么rotate_half是(-x2, x1)而不是(x2, -x1)。旋转方向的正负号决定的是相对位置编码的正负,一旦反了,模型会认为距离左边和距离右边是反的,训练时虽然损失能下降一些,但泛化能力会很差。
我个人在实际操作中的体会是,RoPE 看起来只是位置编码,但它和模型深度、注意力头数、训练长度都有交互。调整频率 base 去适配长上下文时,应该先在小的验证集上观察困惑度变化,而不是只看 loss。因为 loss 在 2% 以内的波动可能不代表泛化能力改善。你需要对比不同长度的绝对位置差上的注意力分布,看它是否平滑。
如果你手头有大模型的权重,可以试着用代码提取出某几个头的注意力权重,绘制成“距离 vs 注意力强度”的曲线,能直观看到 RoPE 的衰减形态。不同头高频和低频占比不同,有的头几乎只关注相邻 50 个 token,有的头直到 512 距离依然能保持一定强度。这正是多尺度建模的魅力。
最后再分享一个小技巧:如果你想测试自己实现的 RoPE 是否正确,不需要跑完整模型。构造一个全零向量,位置设置为 m,旋转后模长应该还为零;再构造一个单位向量,位置 m 旋转后再与位置 m 的原向量点积,理论上应当为 1(同向)。异向测试用位置差 π/2,点积应该为 0。这三个测试通过,你的数学实现基本没有大问题。根据我的经验,这是最快找出维度错乱的方式。