3、RoPE在Attention中的形式化推导：从复数乘法到内积的旋转不变性

好，咱们今天来啃一块硬骨头——RoPE在Attention中的数学形式。说实话，我第一次看RoPE论文时，也被那堆复数、旋转矩阵搞得有点晕。但后来我在项目中亲手推了一遍，发现其实没那么玄乎。说白了，它就是在做一件事：让内积带上位置信息，而且这个信息只跟相对位置有关。

3.1 从绝对位置到相对位置：我们到底想要什么？

先想一个问题：为什么绝对位置编码不够好？

你想想看，如果我用一个绝对位置向量 pi 加到词向量上，那么两个词 xi 和 xj 的注意力分数就是：

score(i,j) = (x_i + p_i) · (x_j + p_j)

展开后你会发现，这里既有 xi·xj（内容交互），也有 pi·pj（位置交互），还有交叉项。但问题是，pi·pj 只依赖于绝对位置 i 和 j，而不是它们的差值 i−j。

我当年做长文本任务时就踩过这个坑。模型在训练集上见过的位置对 (i,j) 有限，一旦推理时出现新的位置组合，注意力分数就乱套了。这就是绝对编码的致命伤——它学不到相对位置关系。

RoPE的解决思路很巧妙：让内积只依赖于 (i−j)。怎么做？用旋转！

3.2 复数视角：旋转就是乘法

我们先从二维情况说起。假设每个词向量是二维的，记作 (a, b)。在复数域里，它可以写成 z = a + bi。

现在，我想给位置 i 的词向量加一个位置信息。RoPE的做法是：旋转这个复数。

z_i' = z · e^{i·θ_i}

其中 θi = i · ω，ω 是一个预设的角频率。你看，旋转角度跟位置 i 成正比。

那么两个旋转后的向量做内积，会发生什么？

在复数域，内积对应的是 共轭乘法后取实部：

⟨z_i', z_j'⟩ = Re( z_i' · conj(z_j') )
            = Re( z·e^{iθ_i} · conj(z)·e^{-iθ_j} )
            = Re( |z|^2 · e^{i(θ_i - θ_j)} )
            = |z|^2 · cos(θ_i - θ_j)

看到了吗？结果只依赖于 θi − θj = (i−j)·ω。完美！

核心结论：在复数域中，对向量做旋转操作后，内积只保留相对位置信息。这就是旋转不变性的本质。

3.3 扩展到高维：分块旋转

实际中词向量维度 d 远大于2，比如 512、1024。怎么办？

RoPE的做法是：把 d 维向量分成 d/2 个二维子空间，每个子空间独立旋转，但旋转频率不同。

具体来说，对于第 k 个二维子空间（k = 0, 1, ..., d/2−1），旋转角度为：

θ_{i,k} = i · ω_k
ω_k = 10000^{-2k/d}

这个频率设计跟Sinusoidal位置编码如出一辙。低频子空间旋转慢，编码长距离依赖；高频子空间旋转快，编码短距离依赖。

我习惯把这个过程想象成一个「多尺度旋转盘」——每个子空间以不同速度旋转，组合起来就能编码丰富的相对位置信息。

3.4 矩阵形式：Attention中的实现

好，理论说完了，咱们看看代码里怎么实现。假设 Q 和 K 的维度是 [batch, seq_len, num_heads, head_dim]，RoPE的实现步骤是：

  1. 构造旋转矩阵：对每个位置 i，生成一个 d×d 的块对角旋转矩阵 Ri
  2. 旋转 Q 和 K：Q' = Q · Ri，K' = K · Rj
  3. 计算注意力：score = Q' · K'T

但直接构造旋转矩阵太浪费了。实际实现时，我们用一种更高效的方式：

def apply_rope(x, sin, cos):
    # x: [batch, seq_len, num_heads, head_dim]
    # sin, cos: [seq_len, head_dim]
    
    # 将x分成两半
    x1 = x[..., :head_dim//2]
    x2 = x[..., head_dim//2:]
    
    # 旋转操作
    x_rotated = torch.cat([
        x1 * cos - x2 * sin,
        x1 * sin + x2 * cos
    ], dim=-1)
    
    return x_rotated

这段代码其实就是复数乘法的矩阵版本。你仔细看：x1 * cos - x2 * sin 对应实部，x1 * sin + x2 * cos 对应虚部。嗯，这里要注意，实际实现时 sin 和 cos 需要预先计算好，避免重复计算。

避坑指南：我曾经在实现RoPE时，把 sin 和 cos 的维度搞反了。记住，sin 和 cos 的 shape 应该是 [seq_len, head_dim]，而不是 [head_dim, seq_len]。这个错误会导致位置信息完全错乱，模型训练不收敛。

3.5 旋转不变性的几何直觉

为什么旋转后内积只跟相对位置有关？我画个图你就明白了。

旋转不变性示意图 v v·R(i) v·R(j) θ_i θ_j θ_i - θ_j 内积 ⟨v·R(i), v·R(j)⟩ = |v|²·cos(θ_i - θ_j)，只依赖于相对角度

你看，向量 v 被旋转到两个不同位置。它们之间的夹角就是 (θi − θj)。内积等于 |v|²·cos(相对角度)，跟绝对位置 i、j 无关。这就是旋转不变性的几何意义。

3.6 为什么这比绝对编码好？

我总结一下RoPE的几个关键优势：

特性 绝对位置编码 RoPE
内积依赖 绝对位置 i, j 相对位置 i−j
外推能力 差（没见过的位置对） 强（旋转频率连续）
计算开销 低（加法） 中等（乘法+三角函数）
与Attention兼容性 需修改Attention结构 天然兼容，直接旋转Q/K

说白了，RoPE让模型在计算注意力时，天然就能感知到「谁离谁近」。这个特性在长文本、长序列任务中特别重要。我记得有一次做128K长度的文档理解，用绝对编码的模型在长距离依赖上完全崩了，换成RoPE后效果立竿见影。

注意：RoPE虽然好，但不是万能的。它假设旋转频率是固定的，对于某些需要动态调整位置信息的任务（比如相对位置范围变化很大的情况），可能需要结合其他技巧。我建议在实际项目中，先用RoPE作为基线，再根据任务特点做微调。

3.7 小结

这一章我们从复数乘法出发，推导了RoPE在Attention中的形式。核心就一句话：旋转操作让内积只保留相对位置信息。这个性质使得RoPE在长序列任务中表现出色，而且实现起来也不复杂。

下一章我们会深入RoPE的代码实现细节，包括如何高效计算 sin/cos、如何处理 batch 维度、以及一些常见的优化技巧。到时候我会分享一些我在实际项目中踩过的坑，保证让你少走弯路。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →