旋转位置编码（Rotary Position Embedding, RoPE）
核心思想
RoPE 的核心思想是通过旋转矩阵对词向量进行变换，从而将位置信息编码到词向量中，与传统的位置编码方式不同，RoPE 不是简单地给词向量加上位置向量，而是通过旋转操作来编码位置信息。
数学原理
1 基本公式
对于位置 $m$ 处的词向量 $\mathbf{x} = [x_1, x_2, ..., x_d]$，RoPE 将其变换为：
$$f(\mathbf{x}, m) = \mathbf{R}_m \cdot \mathbf{x}$$
$\mathbf{R}_m$ 是旋转矩阵：
$$\mathbf{R}_m = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 & \cdots & 0 \ \sin m\theta_1 & \cos m\theta_1 & 0 & 0 & \cdots & 0 \ 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \cdots & 0 \ 0 & 0 & \sin m\theta_2 & \cos m\theta2 & \cdots & 0 \ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots \ 0 & 0 & 0 & 0 & \cdots & \cos m\theta{d/2} \end{pmatrix}$$
2 分块处理
通常将词向量分为 $d/2$ 个二维子空间，每个子空间应用一个旋转矩阵：
$$\begin{pmatrix} x{2k} \ x{2k+1} \end{pmatrix}' = \begin{pmatrix} \cos m\theta_k & -\sin m\theta_k \ \sin m\theta_k & \cos m\thetak \end{pmatrix} \begin{pmatrix} x{2k} \ x_{2k+1} \end{pmatrix}$$
$\theta_k = 10000^{-2k/d}$ （与Transformer中的正弦波编码类似）
关键特性
1 相对位置编码
RoPE 最显著的特性是能够自然地表示相对位置关系：
$$f(\mathbf{x}, m)^T f(\mathbf{y}, n) = f(\mathbf{x}, 0)^T f(\mathbf{y}, n-m)$$
这意味着两个位置 $m$ 和 $n$ 的词向量点积只依赖于它们的相对位置差 $(n-m)$。
2 旋转不变性
$$f(\mathbf{Rx}, m) = \mathbf{R}_m \cdot \mathbf{R} \cdot \mathbf{x} = \mathbf{R} \cdot \mathbf{R}_m \cdot \mathbf{x} = \mathbf{R} \cdot f(\mathbf{x}, m)$$
旋转操作与位置编码可以交换顺序。
实现方式
1 前向传播中应用
def apply_rotary_pos_emb(x, cos, sin):
# x: [batch_size, seq_len, num_heads, head_dim]
# 将x分成两半并交错旋转
x_rotated = x * cos + rotate_half(x) * sin
return x_rotated
def rotate_half(x):
x1 = x[..., :x.shape[-1]//2]
x2 = x[..., x.shape[-1]//2:]
return torch.cat([-x2, x1], dim=-1)
2 预计算cos和sin
def precompute_freqs_cis(dim, max_seq_len, theta=10000.0):
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
t = torch.arange(max_seq_len, device=freqs.device)
freqs = torch.outer(t, freqs) # [max_seq_len, dim//2]
freqs_cos = torch.cos(freqs)
freqs_sin = torch.sin(freqs)
return freqs_cos, freqs_sin
在Transformer中的应用
1 自注意力机制中的使用
def attention_with_rope(query, key, value, cos, sin):
# 在Q和K上应用RoPE
query = apply_rotary_pos_emb(query, cos, sin)
key = apply_rotary_pos_emb(key, cos, sin)
# 计算注意力分数
attn_weights = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
attn_weights = F.softmax(attn_weights, dim=-1)
# 注意力输出
output = torch.matmul(attn_weights, value)
return output
2 位置编码的计算时机
- 在输入层不需要添加位置编码
- 仅在每一层的自注意力计算中应用RoPE
- 不同层可以共享相同的旋转矩阵预计算结果
优势与局限性
1 优势
- 相对位置编码：自然地捕获相对位置关系
- 外推能力：比绝对位置编码更好的长度外推性能
- 计算效率：只需额外的几个逐元素操作，计算开销小
- 兼容性：可以方便地集成到现有的Transformer架构中
2 局限性
- 非因果位置关系：没有明确区分过去和未来的位置
- 维度限制：需要在偶数维度上工作
- 长距离衰减：虽然理论上有，但实践中可能需要特殊处理
应用实例
1 LLama系列
class LlamaAttention(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.head_dim = self.hidden_size // self.num_heads
self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
def forward(self, hidden_states, attention_mask, position_ids):
# 计算Q、K、V
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
# 应用RoPE
cos, sin = self.rotary_emb(query_states, position_ids)
query_states = apply_rotary_pos_emb(query_states, cos, sin)
key_states = apply_rotary_pos_emb(key_states, cos, sin)
# 标准注意力计算
attn_output = self._attention(query_states, key_states, value_states, attention_mask)
return self.o_proj(attn_output)
与其他位置编码的比较
| 特性 | RoPE | 绝对位置编码 | 相对位置编码 |
|---|---|---|---|
| 外推能力 | 好 | 差 | 中 |
| 计算效率 | 高 | 高 | 低 |
| 内存占用 | 低 | 低 | 高 |
| 实现复杂度 | 中 | 低 | 高 |
参数选择建议
- 基础频率 $\theta$：通常使用10000.0
- 缩放因子：对于长序列，可以使用更大的缩放因子
- 维度分配：确保维度为2的倍数
RoPE 通过旋转矩阵优雅地解决了位置编码问题，在保持相对位置关系的同时，实现了高效的并行计算和良好的外推能力，已成为现代Transformer架构中的标准方案之一。
