当年调试两层 Transformer 时,我踩过一个非常典型的坑:把输入句子里几个词的顺序打乱之后,模型输出竟然几乎没变。查了半天才意识到,问题不在实现,而在位置信息。没有位置编码的注意力机制本质上是对称的,它不关心词与词之间隔了多远,只关心内容像不像。短句上这种毛病不明显,序列一长就彻底露馅。后来我开始逐个翻主流开源语言模型的训练和推理代码,发现这几年位置编码几乎都从“加法偏置”换成了旋转位置嵌入(Rotary Position Embedding,RoPE)。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当i=0时频率最高,旋转角度变化最快;当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_half:import 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, base=10000.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 和 K,V 完全不旋转。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 扩展到 8192,s=4,head_dim=64,那么新的 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。当你生成第t+1个 token 时,新 Query 要用位置t+1的 cos/sin 旋转,新 Key 也要用位置t+1的 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 的夹角里。理解清楚旋转的几何含义之后,不管是在自定义注意力里手写位置编码,还是在长上下文模型里调外推,你都会比只看文档照抄的人多一层底气和判断力。
当年调试两层 Transformer 时,我踩过一个非常典型的坑:把输入句子里几个词的顺序打乱之后,模型输出竟然几乎没变。查了半天才意识到,问题不在实现,而在位置信息。没有位置编码的注意力机制本质上是对称的,它不关心词与词之间隔了多远,只关心内容像不像。短句上这种毛病不明显,序列一长就彻底露馅。后来我开始逐个翻主流开源语言模型的训练和推理代码,发现这几年位置编码几乎都从“加法偏置”换成了旋转位置嵌入(Rotary Position Embedding,RoPE)。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当i=0时频率最高,旋转角度变化最快;当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_half:import 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, base=10000.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 和 K,V 完全不旋转。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 扩展到 8192,s=4,head_dim=64,那么新的 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。当你生成第t+1个 token 时,新 Query 要用位置t+1的 cos/sin 旋转,新 Key 也要用位置t+1的 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 的夹角里。理解清楚旋转的几何含义之后,不管是在自定义注意力里手写位置编码,还是在长上下文模型里调外推,你都会比只看文档照抄的人多一层底气和判断力。