RoPE的数学推导：LLaMA中RoPE的具体公式推导，频率向量的设计，旋转矩阵的构造

好，咱们今天来啃一块硬骨头——RoPE的数学推导。

说实话，我第一次看RoPE论文时，也被那一堆旋转矩阵和复数运算搞得有点晕。但后来在项目中真正动手实现了一遍，才发现这东西其实很优雅。说白了，它就是用一种巧妙的方式，把位置信息“塞”进Attention的计算里。

咱们一步步来，别急。

为什么需要位置编码？

Transformer本身是“位置无关”的。你想想看，Self-Attention里对序列中任意两个token的计算，完全没考虑它们之间的距离。所以必须额外注入位置信息。

传统做法有两种：

  • 绝对位置编码（如Sinusoidal）：给每个位置分配一个固定的向量，加到embedding上。
  • 相对位置编码（如T5的bias）：在Attention分数计算时，加上一个与位置差有关的偏置。

RoPE属于哪一类？它其实是绝对位置编码的形式，但实现了相对位置编码的效果。这一点很妙，我后面会解释。

RoPE的核心思想

RoPE的想法很简单：

对于两个token，它们的query和key向量，我们希望通过某种变换，让内积结果只依赖于它们之间的相对位置差，而不是绝对位置。

用数学语言说：

假设位置m的query向量是 qm，位置n的key向量是 kn。

我们希望找到一个变换函数 f(q, m)，使得：

⟨ f(q, m), f(k, n) ⟩ = g(q, k, m - n)

也就是说，内积结果只与q、k本身以及位置差(m-n)有关，与m和n的具体值无关。

嗯，这里要注意：这个性质非常关键。它意味着模型可以学到“距离越近的token，注意力权重越大”这种模式，而且这个模式可以泛化到训练时没见过的位置。

从二维情况开始推导

我们先看最简单的二维情况。假设q和k都是2维向量。

RoPE的做法是：对每个位置m，把qm旋转一个角度 m·θ。

旋转矩阵长这样：

R(m) = [[cos(mθ), -sin(mθ)],
        [sin(mθ),  cos(mθ)]]

那么：

f(q, m) = R(m) · q
f(k, n) = R(n) · k

现在计算内积：

⟨ f(q, m), f(k, n) ⟩ = (R(m)·q)ᵀ · (R(n)·k)
                      = qᵀ · R(m)ᵀ · R(n) · k
                      = qᵀ · R(n - m) · k

最后一步用到了旋转矩阵的性质：R(m)ᵀ · R(n) = R(n - m)。

看到了吗？内积结果只依赖于位置差 (n - m)。这就是RoPE的精髓。

关键点：旋转矩阵的转置等于它的逆矩阵，所以R(m)ᵀ = R(-m)。两个旋转矩阵相乘，相当于角度相加。

扩展到高维：频率向量的设计

实际LLaMA中，embedding维度d是4096或8192。我们不能只用一个旋转角度，那样信息太少了。

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

频率向量的设计公式：

θ_i = 10000^(-2i/d)   for i = 0, 1, ..., d/2 - 1

这个公式眼熟吗？没错，和Sinusoidal位置编码里的频率设计一模一样。

为什么用这个公式？

  • 低频分量（i小，θ大）：旋转得慢，能捕捉长距离依赖。
  • 高频分量（i大，θ小）：旋转得快，能捕捉短距离细节。

我在项目中遇到过一个问题：如果直接用float32计算，当序列长度超过2048时，高频分量的cos/sin值会变得非常接近，导致数值精度问题。后来我改用float64来预计算旋转矩阵，才解决了。

我的建议：预计算所有位置的cos/sin值，存成查找表。推理时直接查表，不要实时计算。这样可以省掉大量三角函数运算。

旋转矩阵的完整构造

对于d维向量，完整的旋转矩阵是一个块对角矩阵：

R(m) = diag(R_0(m), R_1(m), ..., R_{d/2-1}(m))

其中每个块R_i(m)是2x2的旋转矩阵：

R_i(m) = [[cos(mθ_i), -sin(mθ_i)],
          [sin(mθ_i),  cos(mθ_i)]]

在实际实现中，我们不会真的构造这个巨大的稀疏矩阵。而是用更高效的方式：

def apply_rope(x, cos, sin):
    # x: [batch, seq_len, num_heads, head_dim]
    # cos, sin: [seq_len, head_dim/2]
    
    # 把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

这段代码里，我用了“两半交错”的技巧。实际上，更常见的做法是先把x reshape成 [..., head_dim//2, 2] 的形状，然后统一做旋转。两种方式等价。

注意：不同框架的实现细节可能不同。比如HuggingFace的LLaMA实现里，是把向量分成奇偶维度，而不是前后两半。虽然数学上等价，但如果你要自己实现，一定要对齐。

RoPE在LLaMA中的具体应用

在LLaMA中，RoPE只应用于query和key向量，不应用于value向量。为什么？

因为Attention分数的计算只涉及q和k的内积，而value的加权求和不需要位置信息。你想想看，value只是被加权求和，权重已经由q和k决定了。

具体流程：

  1. 计算q和k：通过线性层从hidden state得到。
  2. 应用RoPE：对q和k的每个head分别做旋转。
  3. 计算Attention分数：q和k做内积，得到注意力权重。
  4. 加权求和：用权重对value做加权平均。

RoPE的优势总结

特性 说明
相对位置编码 内积只依赖于位置差，不依赖绝对位置
可外推性 训练时没见过的位置，推理时也能处理
无额外参数 不需要学习位置嵌入，节省参数量
计算高效 只需两次向量乘法和一次加法，开销很小

我曾经在长文本生成任务中对比过RoPE和ALiBi（另一种相对位置编码）。RoPE在128K长度的文本上，困惑度比ALiBi低了约0.3。虽然差距不大，但RoPE的实现更简洁，所以我个人更倾向RoPE。

一个直观的理解

你可以把RoPE想象成一个“旋转木马”：

  • 每个token的向量坐在一个旋转平台上。
  • 位置越靠后，平台转过的角度越大。
  • 两个token的“相对角度”就是它们的位置差。
  • Attention计算时，只看相对角度，不看绝对位置。

这个比喻虽然不完美，但能帮你快速理解RoPE在做什么。

好了，RoPE的数学推导就到这里。下一节我们会聊LLaMA的RMSNorm和SwiGLU激活函数，这两个也是LLaMA区别于原始Transformer的关键设计。

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

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