RoPE 旋转位置编码

RoPE 旋转位置编码RoPE (旋转位置编码)流程step0 计算转速step1 算每个 token 的旋转角step2 求 cos / sinstep3 用 cos / sin 旋转 token作用到 Q 和 K总结RoPE (旋转位置编码)流程step0 计算转速对于每个tokenx i xixi,转速θ为[θ0,······θ(dim/2 -1)]inv_freq 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))数学公式为θⱼ 1 / base^(2j/dim) j 0, 1, ..., d/2−1dim为每个token的维度j0时θ0 1jdim/2 - 1时θⱼ很小这是一种极好的设计相当于一个token拥有了dim/2个指针每个token旋转时有从快到慢多个指针这样不同token快指针相位重叠时慢指针会有差异这种设计还避免了慢指针过慢不同token之间的位置信息差异不够明显step1 算每个 token 的旋转角例如有两个token1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16每个 token 的旋转角只由它的位置 m 决定和 token 本身的数值 1…8 / 9…16 无关位置为 m 的 token其第 j 对的旋转角为φ(m, j) m · θⱼ j 0, 1, ..., d/2−1代码上就是把位置表 positions 和 step0 的转速表 inv_freq 做外积freqs torch.outer(positions.flatten().float(), self.inv_freq) # (batch*seq, dim/2)外积的每一行是一个 token第 j 列就是该 token 第 j 对的旋转角所以每个 token 得到 d/2 个角每个角被一对元素共用即 (x₀, x_{d/2}) 用同一个角。以上面两个 token 为例dim 8base 10000先由 step0 得到转速θ₀ 1 / 10000^(0/8) 1 θ₁ 1 / 10000^(2/8) 0.1 θ₂ 1 / 10000^(4/8) 0.01 θ₃ 1 / 10000^(6/8) 0.001两个 token 的位置 m 分别为 0 和 1代入 φ(m, j) m · θⱼtoken0 (m0)的角度: [0·1, 0·0.1, 0·0.01, 0·0.001] [0, 0, 0, 0 ] token1 (m1)的角度: [1·1, 1·0.1, 1·0.01, 1·0.001] [1, 0.1, 0.01, 0.001]可以看到位置 0 的 token 所有角为 0不旋转位置越靠后角度越大高频指针θ 大转得快、低频指针θ 小转得慢这正是 step0 说的快慢多个指针。这些角下一步会取 cos/sin 组成旋转矩阵作用到 token 上。step2 求 cos / sinstep1 得到的每个 token 只有 d/2 个角而旋转要作用在 d 个分量上所以先把角复制一份拼成 d 长emb torch.cat([freqs, freqs], dim-1) # (batch*seq, dim)这样前 d/2 维和后 d/2 维用的是同一组角——正好对应 step1 说的 (x₀, x_{d/2}) 共用一个角。再对每个角取 cos / sincos emb.cos().to(q.dtype).view(batch_size, seq_len, 1, self.dim) sin emb.sin().to(q.dtype).view(batch_size, seq_len, 1, self.dim)数学公式为cos(m) cos([φ(m,0), ..., φ(m,d/2−1), φ(m,0), ..., φ(m,d/2−1)]) sin(m) sin([φ(m,0), ..., φ(m,d/2−1), φ(m,0), ..., φ(m,d/2−1)])view 里倒数第二维是 1是为了让所有 head 共享同一份 cos/sin广播到 num_heads 维。还是上面两个 tokentoken0 (m0, 角全为 0): cos [1, 1, 1, 1, 1, 1, 1, 1] sin [0, 0, 0, 0, 0, 0, 0, 0] token1 (m1, 角 [1, 0.1, 0.01, 0.001]): cos ≈ [0.54030, 0.99500, 0.99995, 0.9999995, 0.54030, 0.99500, 0.99995, 0.9999995] sin ≈ [0.84147, 0.09983, 0.01000, 0.0010000, 0.84147, 0.09983, 0.01000, 0.0010000]可以看出角越小 cos 越接近 1、sin 越接近 0几乎不转角越大转得越明显前后两半完全一样。step3 用 cos / sin 旋转 token作用到 Q 和 K把 token 向量按 step2 的配对方式劈成两半x₁ 前 d/2 维x₂ 后 d/2 维对每一对施加一个二维旋转矩阵x₁ x₁·cos − x₂·sin x₂ x₁·sin x₂·cos即[x₁] [cos −sin] [x₁] [x₂] [sin cos] [x₂]staticmethod def _apply_rotary_pos_emb(x, cos, sin): x1, x2 x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] cos_half cos[..., : cos.shape[-1] // 2] sin_half sin[..., : sin.shape[-1] // 2] rotated torch.cat([x1 * cos_half - x2 * sin_half, x1 * sin_half x2 * cos_half], dim-1) return rotated注意只对 Q 和 K 旋转V 不旋转旋转是为了让 Q、K 的内积带上相对位置信息。接着上面的例子token0 角全为 0cos 1, sin 0旋转后原样不变token1 的 (x₁, x₂) ([9,10,11,12], [13,14,15,16])代入对应的 (cos, sin)x₁ [9·cos1 − 13·sin1, 10·cos0.1 − 14·sin0.1, 11·cos0.01 − 15·sin0.01, 12·cos0.001 − 16·sin0.001] ≈ [-6.076, 8.552, 10.849, 11.984] x₂ [9·sin1 13·cos1, 10·sin0.1 14·cos0.1, 11·sin0.01 15·cos0.01, 12·sin0.001 16·cos0.001] ≈ [14.597, 14.928, 15.109, 16.012]拼起来token1 旋转后 ≈ [-6.076, 8.552, 10.849, 11.984, 14.597, 14.928, 15.109, 16.012]。总结整个前向就是四步step0 由 base 算出 d/2 个快慢不同的转速 θstep1 用位置 m 乘 θ 得到每个 token 的 d/2 个旋转角step2 把角复制一份并取 cos / sinstep3 把 Q、K 按 (前一半, 后一半) 配对用旋转矩阵旋转。位置不同的 token 转过的角不同旋转后 Q·K 的内积只依赖两者的相对位置 (m−n)这就是 RoPE 把位置信息注入注意力分数的方式。