RoPE 旋转位置编码 (Rotary Position Embedding)
旋转位置编码(Rotary Position Embedding,RoPE)
1. 发展背景与动机
1.1 位置编码的必要性
Transformer模型中的自注意力机制本身是"位置无关"的，这意味着它将输入的token序列视为一个"词袋"，不关心每个token的位置顺序。例如，序列[A, B, C]和[B, A, C]在自注意力机制看来可能没有区别，这显然不符合语言的实际特性。为了让模型理解序列的顺序信息，必须显式地注入位置信息。
1.2 传统位置编码方法及其局限
1.2.1 绝对位置编码
-
正弦/余弦编码：原始Transformer使用固定三角函数生成位置向量（如 $ PE(pos, 2i) = \sin(pos/10000^{2i/d}) $），直接加到词嵌入上。优点是简单且无需训练参数，但缺点是无法直接建模相对位置关系，且长度外推能力差（处理长于训练序列的输入时性能下降）。
-
可学习位置编码：如BERT中为每个位置分配可训练向量。优点是可适应数据，但缺点是需要额外参数，且无法泛化到未见过的序列长度。
1.2.2 相对位置编码
- 如Transformer-XL和T5中，通过偏置项或"距离桶"直接建模token间的相对位置（例如，在注意力分数中添加与相对距离相关的偏差）。优点是能显式捕捉相对位置，但实现复杂（计算复杂度 $ O(L^2) $），且可能破坏注意力机制的线性特性。
1.3 RoPE的提出背景
Rotary Position Embedding (RoPE) 由苏剑林等人在2021年提出，旨在解决传统位置编码方法的局限性。RoPE的核心思想是：通过绝对位置编码的方式实现相对位置编码的效果，即对Query和Key向量进行旋转变换，使注意力分数仅依赖于相对位置差。
RoPE已被LLaMA、ChatGLM、Baichuan等主流大模型广泛采用，成为现代Transformer架构的重要组件。
2. 数学原理
2.1 核心思想
RoPE的核心思想是：通过旋转变换将位置信息编码到查询（Query）和键（Key）向量中，使得变换后的向量内积仅依赖于token间的相对位置差（如 $ m-n $），而与它们的绝对位置（如 $ m $ 和 $ n $）无关。
2.2 数学模型
对于位置 $ m $ 的token，其Query或Key向量 $ \mathbf{x}m $（维度为 $ d $）被划分为 $ d/2 $ 个二维向量对。对第 $ i $ 个向量对 $ [x, x_{2i+1}] $，应用旋转变换：
其中 $ \theta_i = 10000^{-2i/d} $ 是预设的频率参数（与原始Transformer的正弦编码类似），$ i $ 是维度组索引（$ i = 0, 1, ..., d/2-1 $）。
该变换等价于在复数空间中乘以旋转因子 $ e^{im\theta_i} $：
- 将向量对 $ [x_{2i}, x_{2i+1}] $ 视为复数 $ z = x_{2i} + ix_{2i+1} $
- 旋转后的复数为 $ z' = z \cdot e^{im\theta_i} = (x_{2i} + ix_{2i+1})(\cos(m\theta_i) + i\sin(m\theta_i)) $
- 展开后实部和虚部分别对应变换后的向量分量 $ [x'{2i}, x'] $
2.3 相对位置依赖的证明
旋转变换后，位置 $ m $ 的Query向量 $ \mathbf{q}_m $ 和位置 $ n $ 的Key向量 $ \mathbf{k}_n $ 的内积为：
其中 $ \mathbf{R}_{\theta, n-m} $ 是相对位置差 $ n-m $ 对应的旋转矩阵。这表明注意力分数仅依赖于相对位置差，而非绝对位置。
3. 几何解释
3.1 二维旋转的直观理解
在二维平面上，RoPE的旋转变换可视为将向量逆时针旋转角度 $ m\theta_i $：
- 模长不变：旋转是正交变换，不改变向量的长度（保持语义信息）
- 方向改变：旋转角度由绝对位置 $ m $ 决定，位置越靠后，旋转角度越大
示例：假设位置 $ m=3 $ 的某个二维向量分量为 $ [2.0, 1.0] $，频率 $ \theta_i = 0.8 $ 弧度：
- 旋转角度：$ m\theta_i = 3 \times 0.8 = 2.4 $ 弧度（约137.5°）
- 变换计算：\[\begin{aligned} x'_{2i} &= 2.0 \times \cos(2.4) - 1.0 \times \sin(2.4) \approx 2.0 \times (-0.737) - 1.0 \times 0.675 \approx -2.149 \\ x'_{2i+1} &= 2.0 \times \sin(2.4) + 1.0 \times \cos(2.4) \approx 2.0 \times 0.675 + 1.0 \times (-0.737) \approx 0.613 \end{aligned} \]
- 模长验证：原始模长 $ \sqrt{2.0^2 + 1.0^2} = \sqrt{5} \approx 2.236 $，旋转后模长 $ \sqrt{(-2.149)^2 + 0.613^2} = \sqrt{5} \approx 2.236 $（保持不变）
3.2 高维扩展与多频率旋转
对于 $ d $ 维向量，RoPE将其视为 $ d/2 $ 个二维向量对，每组应用相同的旋转变换，但使用不同的旋转频率 $ \theta_i $：
- 高频维度（低 $ i $ ）：$ \theta_i $ 较大（如 $ i=0 $ 时 $ \theta_0 = 10000^0 = 1 $），旋转速度快，捕捉局部位置关系
- 低频维度（高 $ i $ ）：$ \theta_i $ 较小（如 $ i=d/2-1 $ 时 $ \theta_{d/2-1} = 10000^{-(d-2)/d} $），旋转速度慢，捕捉长程位置关系
这种多频率设计使RoPE能同时建模短程和长程依赖。
3.3 相对位置的几何体现
两个向量旋转后的内积取决于它们的相对旋转角度差：
- 位置 $ m $ 的向量旋转 $ m\theta_i $，位置 $ n $ 的向量旋转 $ n\theta_i $
- 它们的夹角差为 $ (m-n)\theta_i $，仅与相对位置差 $ m-n $ 有关
- 内积大小由夹角余弦决定，因此注意力分数自然体现相对位置信息
可以这样想象：RoPE的旋转变换相当于给每个token的向量装上了一个"陀螺"，陀螺的转速由绝对位置决定（位置越靠后，转得越快）。但当计算两个向量的点积（内积）时，真正起作用的是两个陀螺之间的"转速差"。只要转速差相同，点积结果就相同。
4. PyTorch代码实现
以下是一个简化的RoPE实现，基于LLaMA官方代码：
import torch
import torch.nn as nn
import math
def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):
"""
预计算旋转频率向量（复数形式）。
Args:
dim: 词向量维度（d）
seq_len: 序列最大长度
theta: 频率基数（默认10000.0）
Returns:
freqs_cis: 复数形式的旋转向量，形状为 [seq_len, dim//2]
"""
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: dim//2].float() / dim))
t = torch.arange(seq_len, device=freqs.device)
freqs = torch.outer(t, freqs).float()
freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
return freqs_cis
def apply_rotary_emb(
xq: torch.Tensor,
xk: torch.Tensor,
freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""
对Query和Key应用旋转位置编码。
Args:
xq: Query向量，形状为 [batch_size, seq_len, num_heads, head_dim]
xk: Key向量，形状同xq
freqs_cis: 预计算的旋转频率，形状为 [seq_len, head_dim//2]
Returns:
旋转后的Query和Key向量
"""
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)
xq_complex = torch.view_as_complex(xq_)
xk_complex = torch.view_as_complex(xk_)
freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)
xq_rotated = torch.view_as_real(xq_complex * freqs_cis)
xk_rotated = torch.view_as_real(xk_complex * freqs_cis)
xq_out = xq_rotated.flatten(-2).type_as(xq)
xk_out = xk_rotated.flatten(-2).type_as(xk)
return xq_out, xk_out
# 在Attention模块中的使用示例
class Attention(nn.Module):
def __init__(self, args):
super().__init__()
self.wq = nn.Linear(args.hidden_dim, args.dim)
self.wk = nn.Linear(args.hidden_dim, args.dim)
self.wv = nn.Linear(args.hidden_dim, args.dim)
self.freqs_cis = precompute_freqs_cis(
dim=args.head_dim,
seq_len=args.max_seq_len,
theta=args.theta
)
def forward(self, x: torch.Tensor):
batch_size, seq_len, _ = x.shape
q = self.wq(x)
k = self.wk(x)
v = self.wv(x)
q_rotated, k_rotated = apply_rotary_emb(q, k, self.freqs_cis[:seq_len])
scores = torch.matmul(q_rotated, k_rotated.transpose(-2, -1)) / math.sqrt(args.head_dim)
scores = nn.functional.softmax(scores.float(), dim=-1)
output = torch.matmul(scores, v)
return output
代码说明
- precompute_freqs_cis：预计算所有位置和维度组的旋转角度（以复数形式存储），避免重复计算
- apply_rotary_emb：将Query和Key向量重塑为复数形式，与旋转因子相乘（等效于旋转变换），然后转回实数
- 高效性：旋转操作仅涉及逐元素乘法和重塑，计算开销低
5. 优势与总结
5.1 主要优势
- 相对位置感知：注意力分数仅依赖相对位置差，更符合语言建模直觉
- 长度外推（Length Extrapolation）：由于三角函数的周期性，RoPE能处理比训练时更长的序列（通过调整旋转基数或NTK-aware插值可进一步增强）
- 计算高效：无需可学习参数，旋转矩阵可预计算，与线性注意力兼容
- 模型兼容性：可直接集成到现有Transformer中，无需修改注意力机制结构
- 保持模型稳定性：旋转矩阵是正交矩阵，不改变向量的模长，有助于训练过程的稳定
5.2 总结
RoPE通过旋转变换将位置信息优雅地编码到Query和Key向量中，实现了绝对位置编码的简洁性和相对位置编码的表达力。其数学原理深刻（基于复数旋转），实现高效，且具备优异的外推能力，成为现代大模型位置编码的事实标准。
RoPE的几何直观理解是其最大魅力所在：通过绝对位置触发的旋转变换，自然地使注意力分数只依赖于相对位置差。这种设计不仅巧妙且有效，还具有良好的外推性和计算效率，因此被广泛应用于大型语言模型和跨模态模型中。随着对长上下文和多模态建模需求的增长，RoPE及其演进变体将继续发挥关键作用。
浙公网安备 33010602011771号