★ FEATURED ARTICLE 当年调试两层 Transformer 时我踩过一个非常典型的坑把输入句子里几个词的顺序打乱之后模型输出竟然几乎没变。查了半天才意识到问题不在实现而在位置信息。没有位置编码的注意力机制本质上是对称的它不关心词与词之间隔了多远只关心内容像不像。短句上这种毛病不明显序列一长就彻底露馅。后来我开始逐个翻主流开源语言模型的训练和推理代码发现这几年位置编码几乎都从“加法偏置”换成了旋转位置嵌入Rotary Position EmbeddingRoPE。RePE 不再把位置当成一个加在词向量上的噪声而是当成一个作用在 Query 和 Key 上的旋转变换。这个设计有个非常优雅的结果注意力分数天然只依赖两个 token 之间的相对距离而不依赖它们各自的绝对位置。这篇文章不谈空泛结论只讲几件实际用得上的事为什么 Transformer 必须有位置编码、RoPE 的旋转矩阵到底做了什么、怎么从零手写一遍能跑的 PyTorch 实现以及在实际训练和推理里 RoPE 最常见的那几个坑该怎么绕。1. 先讲动机没有位置编码的 Transformer 是什么样子的1.1 自注意力的排列不变性Transformer 的核心是自注意力。给定一组输入向量每一层的注意力输出是“对整条序列做加权求和”权重由两两之间的点积决定Attention(Q, K, V) softmax(QK^T / sqrt(d)) · V这里面有一个容易被忽略的性质如果把输入序列的打乱顺序再喂进去输出也会跟着打乱但每个位置的输出值本身不变。用术语说就是 permutation equivariant排列等变。也就是说模型本身根本不知道“谁在谁后面”。把这个问题类比成语义理解就很好懂了。同样一句话“张三打了李四”和“李四打了张三”用的词一模一样唯一的差别在词序。如果模型没有位置信息它看到的就是一个词袋只能知道出现了张三、李四、打却分不清谁打了谁。对短句来说预训练数据里的大量语序先验还能补救但一旦句子变长、结构变得复杂位置信息的缺失就是致命的。1.2 为什么早期方案不彻底初代 Transformer 用的是正弦加法式位置编码直接把一个随位置变化的正弦波向量加到词向量上。这种方式在当时的机器翻译任务上够用但它有几个明显的先天问题第一位置信息是被“加”进去的和词向量混在一起。模型必须自己从混合向量里把“内容”和“位置”拆开这增加了学习负担。第二注意力分数计算时两个 token 的点积里虽然包含了它们的绝对位置信号却很难直接化简成只依赖相对位置的形式。换句话说模型要学出“距离越近关系越强”这种归纳偏置需要靠大量数据硬学而不是由结构天然保证。后来也有人用可学习位置编码、相对位置偏置、T5 式相对 bias、ALiBi 线性距离惩罚等方案。这些方案各有进展但要么仍然是在加法层面打转要么虽然显式建模了相对位置却没有和 Q、K 的旋转几何结合起来无法同时满足“形式简单”“可外推”“不增加推理状态”这三个要求。这才轮到 RoPE 出场。1.3 我们到底想要怎样一种位置编码站在工程角度理想位置编码应该满足四个条件。一是相对位置敏感两个 token 的注意力分数最好只由它们的相对位移决定这和语言直觉一致。二是形式轻量不能给 KV Cache 或推理流程增加额外的复杂度。三是可外推训练时最长 2048推理时最好能体面地跑到 4096 甚至更长而不是直接断崖式崩溃。四是向量范数稳定对向量做变换时不要把向量的模长也跟着改掉否则注意力在长序列上的尺度会失控。RoPE 恰恰同时满足这四点。它的核心思路是给每个 Query 和 Key 单独施加一个只和它自身位置相关的旋转变换再计算两者的点积。三角函数和旋转矩阵恰好满足“先旋转再点积 旋转一点的差后点积”于是相对位置就自动浮现了。2. 旋转矩阵如何把“位置”藏进夹角2.1 从二维旋转说起回想二维平面上的旋转矩阵。把一个向量逆时针旋转角度m·θ可以得到R(mθ) [[cos(mθ), -sin(mθ)], [sin(mθ), cos(mθ)]]这里m是 token 在序列里的位置θ是一个基础频率。比如位置 1 旋转θ位置 5 旋转5θ位置差为 4 的两个向量它们之间的相对角度就是4θ。如果把注意力计算里的二维 Query 和二维 Key 都各自旋转再来做点积就会得到⟨R(mθ) q, R(nθ) k⟩ q^T R(mθ)^T R(nθ) k q^T R((n-m)θ) k最终等式里只剩n-m跟 m 和 n 的绝对值无关。这就是 RoPE 的全部魔法只要 Q 和 K 用同一套旋转体系点积天然只看相对位置。2.2 从二维推广到高维语言模型的 head_dim 通常是偶数比如 64 或 128。RoPE 的做法是把高维向量拆成d/2个二维子向量每个子向量用各自独立的频率θ_i旋转。这些频率设计成从大到小的指数分布θ_i 10000^(-2i / d), i 0, 1, ..., d/2 - 1当i0时频率最高旋转角度变化最快当i接近d/2-1时频率最低旋转角度变化非常缓慢。高频率通道负责捕捉相邻词的局部关系低频率通道负责传递跨很长的位置信息。这和小波分析里的多尺度分解思路类似不同通道各司其职。完整的高维变换写成公式大概是f(q, m) [ [cos(mθ_0), -sin(mθ_0)], [sin(mθ_0), cos(mθ_0)], ... [cos(mθ_{d/2-1}), -sin(mθ_{d/2-1})], [sin(mθ_{d/2-1}), cos(mθ_{d/2-1})] ] · q这块对角旋转矩阵虽然庞大但实际运算时完全不需要显式构建矩阵。我们可以借用复数观点把相邻的两个维度看成一个复数旋转就是乘上e^(i·m·θ_i)。这正是原始论文和很多开源实现采用复数形式的原因。2.3 三个容易被忽略的几何性质第一个性质是模长不变。旋转不改变向量的长度所以|R(m)q| |q|。这一点很重要它保证了注意力分数不会因为位置编码的引入而出现尺度漂移。第二个性质是相对位置分解。这一点前面已经推导过是 RoPE 名称里“旋转”的底气。它让模型在结构上就具备相对位置建模能力而不是靠数据硬学。第三个性质是远程衰减。由于用了多层不同频率的旋转当位置差变大时高频率分量之间的点积会快速震荡抵消整体上会出现一个随距离增加而衰减的包络。这个衰减不是某个规则写死的而是由频率分布自然形成的类似于原始正弦编码里的衰减行为但尺度更平滑。实际调试时我经常看到有人问“为什么 RoPE 不直接把位置加在 embedding 上”。答案是加法会改向量长度旋转不会加法无法化简成n-m旋转可以。RoPE 是乘法式位置编码里目前最优雅的一种。3. 从公式到 PyTorch一份可直接运行的实现3.1 先写旋转函数高维旋转的落地写法有两种。论文原汁原味的写法是用复数每两个相邻维度组成一个复数和预计算的复数频率相乘。主流框架里更常见的是rotate_half写法它利用向量前后两半的错位拼接省去复数操作。先看rotate_halfimport torch def rotate_half(x): 把向量的后半段取负接到前半段后面。 x1 x[..., : x.shape[-1] // 2] x2 x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim-1)再看应用函数def apply_rotary_pos_emb(q, k, cos, sin): # 假设 q/k 形状都是 [batch, num_heads, seq_len, head_dim] # cos/sin 形状为 [seq_len, head_dim] cos cos[: q.size(-2), :].unsqueeze(0).unsqueeze(0) sin sin[: q.size(-2), :].unsqueeze(0).unsqueeze(0) q_embed (q * cos) (rotate_half(q) * sin) k_embed (k * cos) (rotate_half(k) * sin) return q_embed, k_embed这里的关键是把向量x拆成前后两半然后用(q * cos) rotate_half(q) * sin组合。它实际上把不同维度对按一定分组做了旋转和复数乘法是同一族旋转嵌入只是分组顺序不同。在工程上这种写法更直观也完全不用处理复数类型。3.2 预计算 cos 和 sin 表位置编码是确定的不需要学习所以可以预先算出一张表运行时按位置取用。频率公式里的10000是 base后面会讲它怎么影响外推。def precompute_rope_cache(dim, max_seq_len, base10000.0): # 生成 [θ_0, θ_1, ..., θ_(dim/2-1)] inv_freq 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) # 生成每个位置的角度 m * θ_i t torch.arange(max_seq_len).float().unsqueeze(1) angles t * inv_freq.unsqueeze(0) # [max_seq_len, dim // 2] angles torch.cat((angles, angles), dim-1) # [max_seq_len, dim] cos_cached angles.cos() sin_cached angles.sin() return cos_cached, sin_cached这里torch.cat((angles, angles), dim-1)是工程里很常见的加速技巧先算出dim // 2个角再复制一份拼到后面让每个旋转“配对”有两个维度共享同一个频率省去逐维度查表。3.3 复数写法更贴近论文原理解读如果你从论文或某些复现代码里见到复数实现也不奇怪。它的核心是把相邻两个维度读成一个复数def apply_rotary_emb(xq, xk, freqs_cis): # xq: [batch, num_heads, seq_len, head_dim] xq_ xq.float().view(*xq.shape[:-1], -1, 2) xk_ xk.float().view(*xk.shape[:-1], -1, 2) xq_ torch.view_as_complex(xq_) # [batch, heads, seq, dim/2] xk_ torch.view_as_complex(xk_) freqs_cis freqs_cis[: xq_.size(-2)] # [seq, dim/2] freqs_cis freqs_cis.unsqueeze(0).unsqueeze(0) xq_out torch.view_as_real(xq_ * freqs_cis).flatten(-2) xk_out torch.view_as_real(xk_ * freqs_cis).flatten(-2) return xq_out.type_as(xq), xk_out.type_as(xk)复数写法对数学推导更友好xq_ * freqs_cis就是复平面上的一次旋转。但复数在有些后端比如某些半精度推理脚本上会带来额外的类型兼容问题所以在工程上rotate_half版本更常用。3.4 接入多头注意力把上面的片段串进一个常规多头注意力模块核心就三行cos, sin precompute_rope_cache(head_dim, max_seq_len) q_embed, k_embed apply_rotary_pos_emb(q, k, cos, sin) attn_weights torch.matmul(q_embed, k_embed.transpose(-2, -1)) / math.sqrt(head_dim)这里有一个很容易搞错的细节只旋转 Q 和 KV 完全不旋转。V 只是被注意力分数加权求和的对象位置关系已经由权重表达了再给 V 加位置反而会造成信息重复。4. 和绝对位置、相对位置、ALiBi 的横向对比4.1 一张表看清主流方案位置编码方案作用位置相对位置依赖训练后外推能力KV Cache 负担典型场景正弦绝对编码加到词嵌入上隐式需要模型自己学弱超出训练长度迅速退化无早期 Transformer可学习绝对编码加到词嵌入上隐式很弱训练长度是最硬上限无经典双向编码器T5 相对 bias加到注意力 logits 上显式中等超出 bias 表范围即失效需要缓存 bias 表T5 类模型ALiBi对注意力 logits 做距离惩罚显式较好但没有旋转几何的深层性质无长文本推理类模型RoPE作用于 Q、K 向量本身显式由旋转推导中等但配合 base 调整能显著增强无额外状态现代大模型默认方案4.2 为什么加法式方案输在了结构上加法式位置编码的问题在于位置信息是“叠”在内容上的。模型对同一个位置的向量施加的偏置会被投到所有下游计算里最后注意力分数确实包含了位置信号但这个信号和内容信号纠缠在一起无法分离。相对位置 bias 和 ALiBi 把位置直接叠加到注意力得分上已经比加法式强很多但它们仍然是一种“修正项”没有改变 Q、K 本身的几何结构。RoPE 则是在计算 Q、K 的点积之前先把它们做了旋转。旋转之后Q、K 的内容信息一点没丢但点积的表达式里自动出现了相对位置。这就是“乘法式”和“加法式”的本质区别RoPE 把位置信息变成了向量之间的夹角关系而不是一个附加数值。4.3 长上下文场景下谁更省心从工程角度看RoPE 在 KV Cache 上的优势也很明显。可学习位置编码在推理时需要为每个位置维护独立的嵌入向量ALiBi 和 T5 bias 也要在运行时查表或构造偏置。RoPE 则只需要在加载 K、V 时把当前增量位置的旋转角算好缓存里无需额外存位置状态。推理效率上RoPE 几乎没有额外开销这也是大模型普遍选择它的一个重要原因。当然 RoPE 不是万能药。它放弃了对模型“某绝对位置有特殊含义”的显式建模如果你做的是必须依赖绝对位置的合成任务比如“第 5 个位置一定是某种 token”RoPE 可能需要额外输入特征来补。但对绝大多数自然语言序列建模来说相对位置远比绝对位置重要。5. 工程落地避坑外推、base 和 KV Cache5.1 训练 2048推理 8192为什么直接崩很多人在自己的模型上把训练序列长度定为 2048用了 RoPE 之后天真地以为“相对位置编码应该能外推”结果把输入改成 8192 一跑困惑度直接飙升。原因是 RoPE 的频率是按训练时见过的位置范围设计的。位置 4096 对应的最低频通道角度可能是训练时从未出现过的值模型虽然在结构上是相对位置的但训练到的统计特性并没有覆盖这么大的旋转角度范围。绝对位置虽然被旋转“隐藏”了却仍然会通过频率范围影响模型。5.2 调 base 才是外推的真正入口长度外推最常见的方案是修改频率公式里的 base。NTK-aware 放缩是目前最通用的经验法则假设你打算把上下文从训练长度 L 扩展到目标长度L令放缩比s L / L那么新的 base 大致取new_base base × s^(d / (d - 2))这里的d是 head_dim。为什么分母是d-2因为频率最低的那个通道i d/2-1对应的角频率是base^(-(d-2)/d)。要让它在目标长度内也至少能转满必要的圈数就必须按这个幂次抬高 base。这个公式本质上是保证最高频率不变、最低频率按比例压低让整个频谱能均匀覆盖到更长的序列。实际调参时我不会一次性把 base 从 10000 调到百万级。比如从 2048 扩展到 8192s4head_dim64那么新的 base 大约是10000 × 4^(64/62) ≈ 10000 × 4.09 ≈ 40900这个量级是相对温和的短序列上的表现也不会明显退化。5.3 动态插值和 YaRN 的思路NTK-aware 的 base 改法只针对“目标长度”单一值。如果推理时输入长度是 4096过一会儿变成 8192再变成 32768固定 base 就不太够用了。动态 NTK 的思路是随当前序列长度实时计算 base长度越长 base 越大避免最坏情况出现。更精致的方案是 YaRN。它的思路是对不同频率通道区别对待高频通道波长很短训练中已经见过大量完整周期尽量保持不动低频通道波长很长训练中可能只见过不足一圈按放缩比线性拉伸中间频段做平滑过渡。这样做的物理直觉是高频通道编码的是局部精细关系微调它会破坏既有能力低频通道编码的是全局粗粒度关系拉伸它代价小、收益大。我实测下来YaRN 在 2 倍到 4 倍长度扩展时效果不错但每次换配置都必须跑一遍验证集对比困惑度没有银弹。5.4 KV Cache 上最容易犯的错RoPE 在自回归推理时KV Cache 里缓存的 Key 必须是已经做过旋转的 Key。当你生成第t1个 token 时新 Query 要用位置t1的 cos/sin 旋转新 Key 也要用位置t1的 cos/sin 旋转但缓存里旧的 Key 绝不能再用当前位置的 cos/sin 重新乘一遍。如果你忘记区分“位置索引”和“缓存序号”就会出现一个经典 bug缓存里的 Key 被反复旋转注意力分数完全错乱。我在调试解耦注意力时遇到过类似问题最后是打印每个位置 Key 的 norm 和夹角才发现不对。正确做法是给 RoPE 模块传入position_offset让新 token 的旋转角从offset开始取而缓存区只负责追加不做重复旋转。5.5 推荐的外推流程我自己现在做长上下文微调时一般先跑一轮短脚本把目标长度和训练长度的困惑度差量测出来。如果差量不大直接微调少量 step 就能恢复如果差量大先按 NTK 公式改 base再评估。直白的经验是base 改小了短序列受损base 改大了长序列不平滑需要找一个在两者之间平衡好的点并带着目标长度一起做几次短训练来校准。6. 我实际调试 RoPE 时的一些经验清单最后分享几条零散的实操体会都是常规文档里不太会写的东西。第一RoPE 的实现通常不需要学习参数但它极度依赖 dtype。在 float16 下累计几百个位置的 cos/sin 表精度损失会在长序列上累积。如果发现 4096 长度下中间层输出有奇怪抖动优先把 RoPE 的缓存表用 float32 存储计算时再转换。第二head_dim 必须是偶数。如果你的注意力头维度是奇数要么补一个维度要么用 padding 截断。实际模型里 head_dim 几乎都是 2 的幂但自己搭小实验时容易忽略。第三位置频率的顺序会影响训练效果。有的实现把高频维度排在前面有的做交错排列。高频在前、低频在后的排列是最常见的但如果你从开源代码里抄实现一定要保持预计算表和rotate_half的分组方式一致。混用不同的分组会导致模型权重不匹配加载即失效。第四微调时如果只增加了上下文长度但没有增加数据多样性外推效果仍然有限。RoPE 只负责“结构上能支持更长的相对距离”模型要真正学会利用这些长距离信息数据里必须要有足够的跨长距离依赖样本。否则你可以把位置编码调得再完美模型也只是把注意力均匀摊开并不会产出更智能的长文能力。第五如果只是为了快速验证 RoPE 是否导致 bug最直接的方法是构造一个固定权重的小模型分别输入“位置 0 到 256”和“位置 1 到 257”对比两组输出的相对差异是否一致。因为 RoPE 是相对位置驱动的好的实现应当保证这两组结果几乎只在绝对位置分量上有微小差别。这个小实验能帮你把实现问题和模型问题快速分开。RoPE 并不是什么复杂的高深理论它本质上就是一个聪明的旋转矩阵应用只是把“位置”做进了 Q、K 的夹角里。理解清楚旋转的几何含义之后不管是在自定义注意力里手写位置编码还是在长上下文模型里调外推你都会比只看文档照抄的人多一层底气和判断力。
阅读完成 · 觉得有帮助？
