RoPE 推导详解：从向量旋转到复数表示

Rotary Positional Embedding（RoPE）是目前大语言模型中最常用的位置编码方式之一，被 GPT、LLaMA、ChatGLM 等模型广泛使用。

相比传统 绝对位置编码（Absolute Position Encoding），RoPE 的核心特点是：

在 Attention 中自然引入相对位置（Relative Position）。

本文将从 两个角度推导 RoPE：

二维向量旋转（几何视角）
复数乘法（数学视角）

同时明确写出推导过程中使用的关键公式。

1 Transformer 中的位置问题

Transformer 的 Attention 计算为：

$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$

问题是：

Attention 本身不包含位置信息。

如果输入序列顺序改变，模型无法感知。

因此需要 Positional Encoding。

2 RoPE 的核心思想

RoPE 的思想非常简单：

将 embedding 的每两个维度视为一个二维向量，并根据 token 位置对其进行旋转。

假设 embedding 向量：

$x = (x_1, x_2, x_3, x_4, \dots)$

RoPE 会把它拆成二维对：

$(x_1,x_2), (x_3,x_4), (x_5,x_6), \dots$

每一对进行 二维旋转变换。

3 二维向量旋转推导

3.1 二维旋转矩阵

在二维空间中，向量旋转角度 $\theta$ 的公式为：

$\begin{pmatrix} x'_1 \\ x'_2 \end{pmatrix}= \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix}$

展开可得：

$x'_1 = x_1\cos\theta - x_2\sin\theta$ $x'_2 = x_1\sin\theta + x_2\cos\theta$

这就是 RoPE 的基本变换公式。

3.2 位置依赖的旋转角度

RoPE 的旋转角度依赖于 token 位置 $p$ 和 维度频率 $\theta_k$：

$\theta = p \cdot \theta_k$

其中：

$\theta_k = 10000^{-2k/d}$

参数说明：

符号	含义
$p$	token 位置
$k$	维度 index
$d$	embedding 维度

因此旋转角为：

$p \cdot 10000^{-2k/d}$

4 RoPE 在 Attention 中的效果

Attention 计算核心是：

$QK^T$

如果对 Query 和 Key 同时进行旋转：

$Q' = R_{\theta_i} Q$ $K' = R_{\theta_j} K$

其中：

$R_\theta = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix}$

则内积为：

$(Q')^T K'$

利用旋转矩阵性质：

$R_\theta^T = R_{-\theta}$

可得：

$(Q')^T K' = Q^T R_{\theta_i-\theta_j} K$

关键结论：

Attention 只依赖于位置差 $(i-j)$

因此 RoPE 天然编码相对位置。

5 复数视角推导

上面的推导基于 二维向量。
但 RoPE 论文使用 复数形式，因为表达更简洁。

5.1 向量与复数的对应

定义：

$z = x_1 + i x_2$

则二维向量：

$(x_1,x_2)$

可以表示为复数：

$z$

5.2 欧拉公式

复数旋转使用 欧拉公式：

$e^{i\theta} = \cos\theta + i\sin\theta$

5.3 复数旋转

二维旋转等价于：

$z' = z e^{i\theta}$

展开：

$(x_1 + i x_2)(\cos\theta + i\sin\theta)$

计算得到：

$= x_1\cos\theta - x_2\sin\theta + i(x_1\sin\theta + x_2\cos\theta)$

因此：

实部：

$x'_1 = x_1\cos\theta - x_2\sin\theta$

虚部：

$x'_2 = x_1\sin\theta + x_2\cos\theta$

这 完全等价于二维旋转矩阵。

6 RoPE 的复数形式（完整推导）

上一节说明了：

二维旋转可以表示为复数乘法：

$z' = z e^{i\theta}$

RoPE 论文就是利用这个性质，将二维向量旋转写成复数形式。

6.1 向量与复数表示

对于 embedding 的一对维度：

$(x_1,x_2)$

定义复数表示：

$z = x_1 + i x_2$

因此 Query 与 Key 可以写成：

$q = q_1 + i q_2$ $k = k_1 + i k_2$

6.2 RoPE 的旋转

对于位置 (i) 的 Query：

$q_i' = q_i e^{i\theta_i}$

对于位置 (j) 的 Key：

$k_j' = k_j e^{i\theta_j}$

其中：

$\theta_i = i \cdot \theta$ $\theta_j = j \cdot \theta$

6.3 Attention 内积

Attention 的核心是：

$QK^T$

在复数形式中，对应 复数内积。

为了得到实数结果，需要使用 共轭复数：

$\langle q,k\rangle = \text{Re}(q \overline{k})$

其中：

(\overline{k}) 为复数共轭
(\text{Re}(\cdot)) 表示取实部

6.4 代入 RoPE

首先计算：

$q_i' \overline{k_j'}$

代入旋转：

$q_i' = q_i e^{i\theta_i}$ $k_j' = k_j e^{i\theta_j}$

则：

$\overline{k_j'} = \overline{k_j} e^{-i\theta_j}$

因此：

$q_i' \overline{k_j'} = q_i e^{i\theta_i} \cdot \overline{k_j} e^{-i\theta_j}$

整理得到：

$= q_i \overline{k_j} e^{i(\theta_i-\theta_j)}$

6.5 取实部

Attention score 为：

$\text{Re}(q_i' \overline{k_j'})$

因此：

$\text{Re}(q_i \overline{k_j} e^{i(\theta_i-\theta_j)})$

关键点在于：

旋转只依赖于 ( \theta_i-\theta_j )

6.6 相对位置出现

由于：

$\theta_i = i\theta$ $\theta_j = j\theta$

因此：

$\theta_i - \theta_j = (i-j)\theta$

所以最终 Attention 为：

$\text{Re}(q_i \overline{k_j} e^{i(i-j)\theta})$

6.7 关键结论

Attention score 只依赖：

$(i-j)$

即：

相对位置

而不是绝对位置。

因此 RoPE 天然实现 relative position encoding。

6.8 总结

复数形式下 RoPE 的推导可以总结为：

1️⃣ 二维旋转

$z' = z e^{i\theta}$

2️⃣ Query / Key 旋转

$q_i' = q_i e^{i\theta_i}$ $k_j' = k_j e^{i\theta_j}$

3️⃣ Attention 内积

$q_i' \overline{k_j'} = q_i \overline{k_j} e^{i(\theta_i-\theta_j)}$

4️⃣ 位置关系

$\theta_i-\theta_j=(i-j)\theta$

最终得到：

Attention 只依赖 token 的相对位置。

7 为什么不同维度使用不同频率

RoPE 中：

$\theta_k = 10000^{-2k/d}$

这与 Transformer 原始位置编码一致。

原因是：

不同维度代表不同频率。

类似 傅里叶基底（Fourier Basis）：

低维 → 低频
高维 → 高频

这样模型能够同时表示：

长距离关系
短距离关系

8 RoPE 总结

RoPE 的核心可以总结为一句话：

将 embedding 每两个维度看作二维向量，根据 token 位置进行旋转。

关键数学工具包括：

使用到的公式

1️⃣ 二维旋转矩阵

$\begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix}$

2️⃣ 欧拉公式

$e^{i\theta} = \cos\theta + i\sin\theta$

3️⃣ 复数旋转

$z' = z e^{i\theta}$

4️⃣ RoPE 频率公式

$\theta_k = 10000^{-2k/d}$

9 直观理解

RoPE 可以用一句非常直观的话理解：

Token 的表示在 embedding 空间中随着位置不断旋转。

当两个 token 计算 attention 时：

它们的相对旋转角度决定注意力强度。

因此模型可以自然建模：

相对距离
顺序关系

pytorch实现

def apply_rotary_pos_emb(x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
    # x: (batch, num_heads, seq_len, head_dim)
    # pos: (batch, seq_len) position indices
    
    batch, num_heads, seq_len, head_dim = x.shape
    
    # Compute frequencies: theta_i = 1 / (10000^(2i/d))
    freqs = 1.0 / (10000 ** (torch.arange(0, head_dim, 2, device=x.device).float() / head_dim))
    
    # pos: (batch, seq_len) -> (batch, seq_len, 1)
    pos = pos.unsqueeze(-1).float()
    
    # angles: (batch, seq_len, head_dim/2)
    angles = pos * freqs.unsqueeze(0).unsqueeze(0)
    
    # cos and sin
    cos = torch.cos(angles)  # (batch, seq_len, head_dim/2)
    sin = torch.sin(angles)
    
    # Reshape x for rotation: split into pairs
    x1 = x[..., 0::2]  # (batch, num_heads, seq_len, head_dim/2)
    x2 = x[..., 1::2]
    
    # Apply rotation
    cos = cos.unsqueeze(1)  # (batch, 1, seq_len, head_dim/2)
    sin = sin.unsqueeze(1)
    
    x_rotated_1 = x1 * cos - x2 * sin
    x_rotated_2 = x1 * sin + x2 * cos
    
    # Interleave back
    x_out = torch.stack([x_rotated_1, x_rotated_2], dim=-1).flatten(-2)
    
    return x_out