RoPE的旋转矩阵推导：从复数到二维旋转的数学之美

说实话，我第一次接触RoPE的时候，也被它的数学绕得有点晕。但后来我发现，这东西本质上就是一个很优雅的数学小把戏——用复数旋转来编码位置信息。今天我就带你一步步拆解这个过程，保证你看完会觉得：哦，原来这么简单。

为什么非得用旋转？

先想一个问题：在Transformer里，位置编码要满足什么条件？

我个人习惯从两个角度来理解：

• 相对位置敏感：模型要知道词A和词B之间的距离，而不是它们各自在哪个绝对位置
• 距离衰减：离得越近的词，相互影响应该越大

你想想看，如果我用一个简单的加法编码，比如位置1加向量p1，位置2加向量p2，那模型怎么知道p1和p2之间的关系？它得自己去学，这效率就低了。

RoPE的思路很直接：让位置信息通过旋转矩阵嵌入到注意力计算中。说白了，就是让两个token的query和key在做点积时，天然就能反映出它们的相对位置。

从复数乘法说起

还记得复数的乘法吗？

z1 = r1 * e^(iθ1)
z2 = r2 * e^(iθ2)
z1 * z2 = r1*r2 * e^(i(θ1+θ2))

这里有个关键性质：复数乘法等价于旋转加缩放。如果我把一个复数看作二维平面上的点，乘以e^(iθ)就相当于把这个点逆时针旋转θ角度。

我在项目中遇到过一个问题：为什么不用实数矩阵而用复数？答案其实很简单——复数天然支持旋转操作，而且数学上更简洁。你想想看，如果用实数矩阵表示旋转，你得写一个2x2的矩阵，多麻烦。

二维旋转矩阵的推导

好，现在我们进入正题。假设我们有一个二维向量 [x, y]，想把它旋转θ角度。用复数表示就是：

z = x + i*y
z' = z * e^(iθ) = (x + i*y) * (cosθ + i*sinθ)

展开这个乘法：

z' = x*cosθ + i*x*sinθ + i*y*cosθ - y*sinθ
   = (x*cosθ - y*sinθ) + i*(x*sinθ + y*cosθ)

所以旋转后的坐标是：

x' = x*cosθ - y*sinθ
y' = x*sinθ + y*cosθ

写成矩阵形式就是：

[x']   [cosθ  -sinθ] [x]
[y'] = [sinθ   cosθ] [y]

嗯，这里要注意：这个矩阵就是二维旋转矩阵R(θ)。它满足一个漂亮的性质——旋转矩阵是正交矩阵，也就是说R(θ)^T * R(θ) = I。这意味着旋转不会改变向量的长度，只是改变了方向。

扩展到高维空间

实际应用中，我们的embedding维度d通常是几百甚至上千。怎么把二维旋转扩展到高维？

做法其实很粗暴：把d维向量分成d/2个二维子空间，每个子空间独立旋转。每个子空间分配一个不同的旋转频率。

比如d=4时：

原始向量: [x0, x1, x2, x3]
分成两组: [x0, x1] 和 [x2, x3]
旋转角度: θ1 和 θ2

旋转后: 
[x0']   [cosθ1  -sinθ1   0       0   ] [x0]
[x1'] = [sinθ1   cosθ1    0       0   ] [x1]
[x2']   [0       0       cosθ2  -sinθ2] [x2]
[x3']   [0       0       sinθ2   cosθ2] [x3]

这就是一个块对角矩阵。每个2x2块负责一个子空间的旋转。

RoPE的完整数学形式

现在我们把所有东西串起来。对于一个位置pos的token，它的query向量q经过RoPE后变成：

q'_pos = R(θ·pos) * q_pos

其中R是一个块对角旋转矩阵，θ是一个包含d/2个频率的向量。

那么两个位置pos1和pos2的query和key做点积时：

q'_pos1 · k'_pos2 = (R(θ·pos1) * q) · (R(θ·pos2) * k)
                   = q^T * R(θ·pos1)^T * R(θ·pos2) * k
                   = q^T * R(θ·(pos2 - pos1)) * k

看到了吗？点积结果只依赖于相对位置(pos2 - pos1)，而不是绝对位置。这就是RoPE最牛的地方——它让注意力机制天然地关注相对位置。

可视化理解

为了让你更直观地理解，我画了一张图来展示RoPE的核心逻辑：

代码实现

最后，给你一个PyTorch的实现示例。这个代码我在实际项目中用过，可以直接拿来用：

import torch
import math

def precompute_freqs_cis(dim: int, max_seq_len: int, base: float = 10000.0):
    """
    预计算旋转矩阵的频率
    dim: embedding维度
    max_seq_len: 最大序列长度
    """
    # 计算每个子空间的频率
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    
    # 生成位置序列
    t = torch.arange(max_seq_len, dtype=torch.float)
    
    # 外积得到每个位置每个子空间的旋转角度
    freqs = torch.outer(t, freqs)
    
    # 转换为复数形式
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor):
    """
    应用RoPE旋转
    x: [batch, seq_len, num_heads, dim]
    freqs_cis: [seq_len, dim//2]
    """
    # 将x转换为复数形式
    x_complex = torch.view_as_complex(
        x.float().reshape(*x.shape[:-1], -1, 2)
    )
    
    # 应用旋转
    x_rotated = x_complex * freqs_cis.unsqueeze(0).unsqueeze(2)
    
    # 转回实数
    x_out = torch.view_as_real(x_rotated).reshape(*x.shape)
    return x_out.type_as(x)

好了，这就是RoPE旋转矩阵的完整推导。从复数乘法到二维旋转，再到高维扩展，每一步都有它存在的道理。你想想看，这么优雅的数学结构，居然能完美解决位置编码的问题，是不是挺神奇的？