RoPE这三个字母最近两年在模型圈子里出现的频率高得吓人。我自己的习惯是拿到一个新模型先翻它的注意力实现十有八九会看到一段对query和key做旋转的代码边上注释写着Rotary Position Embedding。不少朋友问过我它到底是什么原理为什么大家都用它以及长文本扩展里那些NTK、YaRN之类的花活为什么都建立在它上面。这篇文章就把RoPE一次性讲清楚看完你不仅能理解它还能自己动手写一份能跑的版本。只想动手的话直接跳到第3节代码在项目里改改就能用。1. 先说结论RoPE到底解决了什么问题Transformer里的自注意力机制本身是不认位置的。你把一句话的词序打乱每个词算出来的query、key、value其实完全一样注意力矩阵也只取决于词和词之间的内容相似度至于谁在谁前面、谁离谁多远模型一无所知。这对语言建模来说是致命的因为我爱你和你爱我用的词一模一样意思却完全不同。所以需要一种机制把位置信息注入到模型里。位置编码Position Encoding / Embedding干的就是这件事。它要解决的需求可以分成两类一类是绝对位置即这个token是第几个另一类是相对位置即这两个token隔了多远、谁在前谁在后。早期方案往往只侧重其中一个而RoPE比较特别它用一组旋转矩阵同时把两种信息都表达了出来而且是通过乘法而不是加法注入的。这里先给结论后面逐个展开。RoPE的核心做法是对每个位置的query向量和key向量按照它们各自的位置编号做一次旋转。旋转角度由位置编号和维度编号共同决定。做完这个操作之后某个位置的query和另一个位置的key做点积时结果里只会出现两个位置的差值而不是它们各自的绝对值。也就是说注意力分数天然变成相对位置相关的了。这个过程不引入额外参数不改变向量模长而且频率参数是显式可调的为后来的长文本外推留下了一个非常顺手的旋钮。2. 原理拆解怎么把位置转进去2.1 二维旋转是最小单元先回忆一下二维平面里的旋转。一个点(x, y)绕着原点逆时针旋转角度φ会变成x x * cos(φ) - y * sin(φ)y x * sin(φ) y * cos(φ)写成矩阵形式就是标准的旋转矩阵R(φ)。RoPE的想法非常简单粗暴把token的位置编号m当成旋转角度的一部分让第m个token的向量旋转m * θ这么多弧度。这里的θ不是随便给的它是每个维度对独有的频率后面会专门讲。你可以把这一对维度想象成表盘上的一根指针。位置1转一格位置2转两格位置3转三格。token之间隔得越远两根指针的夹角就越大。注意力机制做的事恰恰就是比较query指针和key指针之间的夹角。把位置信息编码成角度本质上就是在说我通过指针之间的夹角来感知距离。2.2 对Q和K做旋转点积只跟相对位置有关这一小节是RoPE最漂亮的地方。假设某个维度对的两个分量是(x, y)位置m的query向量旋转mθ位置n的key向量旋转nθ。现在计算旋转后的query和旋转后的key之间的点积你会发现一个奇妙的现象两个位置的绝对编号m、n在结果里消失了只剩下它们的差m-n。用复数来写会更干净。把(x, y)看成复数z x iy旋转就是乘上e的imθ次方。旋转后的点积取实部Re[ z_q * e^{imθ} * conj(z_k * e^{inθ}) ] Re[ z_q * conj(z_k) * e^{i(m-n)θ} ]右边只跟m-n有关。也就是说无论第5个token和第6个token还是第105个token和第106个token只要它们的距离一样注意力分数里来自位置的那部分贡献就一样。这正是相对位置编码梦寐以求的性质。别小看这个结论。它意味着模型不用凭空记住第100个位置长什么样只需要学会距离为k时应该给予多大关注。这种距离感知是平移不变的也就是不管这段话从哪个位置开始模型对局部语法结构的建模能力都不受影响。很多实验都观察到RoPE训练出来的模型在句子中后部的注意力模式非常稳定就是这个性质在起作用。2.3 不同维度用不同频率高频抓近处低频抓全局问题来了如果所有维度对都用同一个旋转频率θ那指针转到一定程度就会转回原点位置信息开始混叠。想象一根秒针转一圈是60秒但你无法从指针角度判断现在是第1分钟还是第61分钟。所以RoPE让不同的维度对使用不同的频率。具体公式是θ_i base的(-2i/d)次方其中base默认取10000i是维度对的编号d是注意力头维度。编号小的维度对θ大旋转速度快适合区分相邻token编号大的维度对θ小旋转速度慢负责感知远距离的粗略位置。这跟复平面上多个指针转动的效果差不多秒针负责精细计秒时针负责粗略计小时三个指针合起来任意时刻的表盘读数都能唯一确定时间。为什么base取10000这是从早期Transformer的正弦位置编码里继承下来的经验值。在那个方案里base同样是10000用来控制不同维度正弦波的波长范围。RoPE直接沿用了这个设定波长覆盖范围从2π到10000 * 2π刚好能匹配常见训练序列长度。到后面做长文本扩展时很多人改的第一个参数也是这个base原因就在这里——它决定了位置感知的量程。3. 手写RoPE也就二十行代码3.1 预计算cos和sin表RoPE本身不参与训练所有旋转角度只跟位置和维度有关所以可以提前算好一张cos/sin查找表用时直接查。下面这段代码就是标准做法import numpy as np def precompute_rope(dim, max_seq_len, base10000.0): # 每个维度对的旋转频率一共 dim//2 个 inv_freq 1.0 / (base ** (np.arange(0, dim, 2) / dim)) # 每个位置上每个维度对的旋转角度 t np.arange(max_seq_len, dtypenp.float32) angles np.outer(t, inv_freq) # [max_seq_len, dim//2] # 扩成 dim 维方便和向量直接相乘 cos np.cos(angles) sin np.sin(angles) # 拼接成 [max_seq_len, dim] cos np.concatenate([cos, cos], axis-1) sin np.concatenate([sin, sin], axis-1) return cos.astype(np.float32), sin.astype(np.float32)注意最后把cos和sin各复制了一份拼起来。因为待会要对整个head_dim维度的向量做旋转每一维都需要一个对应的cos/sin值而旋转时是两两成对使用的所以前一半和后一半复用同一组值。如果你用的是前后两半配对的写法这一步必须做对。3.2 rotate-half小技巧真正旋转时最直接的做法是遍历每个二维对手动套旋转矩阵。但这样既慢又啰嗦。工程上有个取巧的办法叫rotate-half。思路是把向量看成前后两半前半是实部后半是虚部然后整体乘以e的iφ次方。def rotate_half(x): d x.shape[-1] x1, x2 x[..., : d // 2], x[..., d // 2 :] return np.concatenate([-x2, x1], axis-1) def apply_rope(q, k, cos, sin, seq_len): # q/k: [batch, heads, seq_len, head_dim] # cos/sin: [max_seq_len, head_dim] cos cos[None, None, :seq_len, :] sin sin[None, None, :seq_len, :] q_out q * cos rotate_half(q) * sin k_out k * cos rotate_half(k) * sin return q_out, k_out这里的rotate_half把后半段取负后挪到前面前半段原样挪到后面。表面上看像个奇怪的切片操作实际上它就是在做二维旋转对每一对(x_i, x_{id/2})新的值等于旧的乘以cos再加上-旧的后半、旧的前半乘以sin。这就是复数乘法展开后的样子。用这个技巧整个旋转过程变成了一次逐元素乘法和一次逐元素加法在GPU上非常高效。3.3 接入注意力的完整流程实际使用时位置编码只加在query和key上value向量不做旋转。原因在于value本身不参与位置计算它只是被注意力权重加权求和的对象。位置信息已经通过权重注入了最终输出再旋转value反而是画蛇添足还会破坏内容语义。完整的接入顺序是这样的先对输入做投影得到q、k、v然后对q和k施加RoPE再做缩放点积注意力最后和v加权求和。算softmax之前q和k的匹配度由两部分组成一部分是内容相似度一部分是相对位置贡献。训练过程中模型会自动学会在两者之间取舍比如某些任务更看重局部n-gram模型就会加强相邻位置的注意力权重。注意RoPE要求head_dim必须是偶数。因为它是两两配对旋转的如果head_dim是奇数你会在rotate_half时得到形状不匹配的报错。目前主流模型的head_dim基本都是64或128不会有这个问题但自己写小模型时容易踩。4. 和其他位置编码方案对比RoPE赢在哪4.1 加法式绝对位置编码实现简单但有明显副作用最早期的Transformer用的是正弦绝对位置编码做法是把位置向量直接加到token向量上。后来的经典预训练模型改用可学习的绝对位置向量训练时直接学一个位置0、位置1、位置2……的嵌入表。这类方案的优点是实现极简几行代码就能搞定。但缺点也很明显。第一加法会把位置信息混进token语义里模型需要费劲把两者解耦第二它是绝对位置的两个位置距离远并不意味着注意力分数天然变小位置关系完全靠模型硬学第三可学习位置表只能覆盖训练时的最大长度一旦序列变长就得重新训练或插值处理不好精度会崩。4.2 相对位置偏置表达能力更强但工程成本高后来出现了相对位置编码做法是在注意力分数上直接加一个关于相对距离的偏置项。比如某些方案按距离分桶每一个距离桶学一个标量偏置能处理的最大相对距离还要人为截断。这种方案表达能力很强却引入了额外参数而且每个token对的距离都要算一遍实现起来比RoPE繁琐。4.3 RoPE凭什么成为主流把主流方案放在一起对比RoPE的优势就很清晰了对比项加法绝对相对偏置RoPE相对距离感知需要模型硬学显式建模显式建模对向量语义的干扰会混入无不改变模长额外参数有或可学习有无实现成本低中低长文本扩展难一般容易可调频三点最关键。第一旋转不改变向量模长只改变方向语义信息完整保留不会像加法那样把词义和位置搅在一起。第二相对距离是数学上自然的产物不是学出来的所以对训练中没见过的距离组合也有一定泛化能力。第三频率是显式的、可调的后面所有长文本技巧都是在它的公式上做文章这是其他方案给不了的操作空间。唯一需要注意的是RoPE在点积中丢失了绝对位置信息。模型确实不太容易区分第1个token和第1001个token这种绝对序号差异。但实践中模型可以通过其他通路比如token本身的语义、上下文的隐式线索把这个信息补回来通常不影响使用。5. 长文本扩展RoPE耐玩的另一面5.1 外推和插值的思路假设一个模型用RoPE在4096长度上训练现在想让它支持8192的输入。直接硬跑会出现什么问题位置m跑到4096以上之后对应的旋转角度是从未见过的值注意力分数分布会变得异常Perplexity困惑度直接飙升。位置上没见过和内容上没见过一样致命。于是一类做法是位置插值。思想很简单既然超过4096的角度没训练过那就把所有角度除以一个缩放因子s比如s2把8192的位置全部映射回训练过的角度范围里。代价是相邻位置的旋转角度差变小了模型对近距离的区分能力会有所下降。这个方法实现起来极容易几乎不用改模型只需要重建一张长度为新上限的cos/sin表。5.2 高频部分先饱和所以调法不同如果每种频率都统一缩放问题就来了。还记得频率分布吗高频维度每走一个token就转很多角度在长序列里转几圈就绕晕了角度完全相同位置关系直接混叠。这就是高频饱和。它对远距离位置感知无所谓但对相邻token的精细区分是致命的。所以更好的做法是不均匀调整。NTK感知缩放的核心思路是把base从10000调大而不是直接缩放位置。调大base之后低频维度负责远距离的旋转变慢能够覆盖更长的序列高频维度负责近距离几乎不动保持了原有的局部精度。一条近似公式是让新base约等于原base乘以缩放因子的d/(d-2)次方在常见配置下基本就是乘以缩放因子本身。这个思路和信号处理里的频率延伸一脉相承也是后来一系列进阶方案共同的出发点。5.3 实践中的有效配置我在实际项目里试过几轮体感是这样的如果只是小幅外推比如1.5倍以内只改base就能撑住如果是要翻几倍长度优先用线性插值效果稳定再往上追求极致效果可以把位置插值、NTK调频和注意力温度修正组合起来用。所有方案都只需要在加载模型前重算一次cos/sin表训练参数一个都不用动这在工程上太友好了。提醒一个细节调整频率后一定要确认模型加载时使用的cos/sin表长度覆盖了你实际推理的最大长度否则推理到一半出现越界或者静默截断输出的文本会莫名其妙变差。6. 实操中踩过的坑6.1 fp16下的精度问题cos/sin表如果用fp16保存位置号一大角度差值太小fp16的表示精度会把这些差值抹平。我在一个长文本任务里就遇到过这种情况序列长度超过2048之后注意力集中在几个位置附近输出质量明显下降排查到最后发现是位置表精度不够。解决方案是保持cos/sin表为fp32只在真正和q、k相乘的那一刻转成fp16或bf16。这一步改动成本极低效果差别却很大。6.2 维度配对方式必须一致RoPE在实现上有个容易踩的坑维度配对顺序。论文描述的是相邻维度两两配对而很多工程实现用的是前后两半配对也就是rotate_half那个写法。两种写法数学上等价但必须保证预计算inv_freq、生成cos/sin表时用的配对方式和apply_rope里的rotate方式一致。如果混用了旋转就会作用在错误的维度对上模型性能会明显下降且很难定位。遇到模型加载后精度异常先检查head_dim是否偶数再检查配对顺序是否和原始checkpoint的实现一致。6.3 推理时的缓存与偏移问题在自回归推理里会把历史token的K和V缓存起来避免重复计算。这里有个常见的陷阱K缓存的应该是什么必须是旋转后的K。因为位置信息已经在旋转时注入进去了缓存旋转后的K新来的query直接和它做点积就行。如果缓存的是旋转前的K就得记录每个K的位置号再在推理时逐个旋转得不偿失。还有个隐蔽的偏移问题。如果缓存不是从位置0开始的比如多轮对话里每轮重新取会话起始位置索引cos/sin表时必须用原始绝对位置而不是当前批次内的相对位置。我见过有代码在第二轮对话时从0开始取cos/sin导致位置全部错位生成的回复语义混乱。排查方法是逐token打印query使用的实际角度和理论值比对错位问题一目了然。RoPE最打动我的地方是它把一个看似需要大量参数和工程技巧的问题简化成了一个干净的几何事实旋转。你不需要一个庞大的可学习位置表不需要复杂的相对距离分桶只需要记住按位置转动query和key这一句话。每次看到新模型还在用RoPE我都会想好的设计大概就是这种样子——原理讲起来五分钟代码写起来二十行效果却能撑起一代又一代大模型的底座。如果你在自研小模型或者改别人的推理脚本建议亲手跑一遍上面的代码把cos/sin表打印出来看看高频维度和低频维度的差异那种直观感受比读十篇博客都有用。