当前位置：首页 > news >正文

旋转位置编码RoPE

news 2025/4/29 5:42:08

Rotary Position Embedding (RoPE)

相关论文地址

https://arxiv.org/abs/2104.09864

RoPE介绍

扭转位置编码（Rotary Position Embedding, RoPE）是一种改进的位置编码方法，用于增强模型对序列顺序的敏感性。它通过对输入序列的特征向量进行旋转变换，将位置信息隐式编码进来，使得模型能够更好地捕获相对位置信息。RoPE 通常用于对比 Transformer 中传统的位置编码，能够提供更强的位置感知能力。

RoPE的数学公式

RoPE 使用了旋转矩阵来对输入向量进行变换。具体公式如下：
$$
\text{RoPE}(x) = \begin{bmatrix}
\cos(\theta_{pos}) - \sin(\theta_{pos}) \
\sin(\theta_{pos}) + \cos(\theta_{pos})
\end{bmatrix}

\begin{bmatrix}
\mathbf{x}{2i}\
\mathbf{x}{2i+1}\
\end{bmatrix}
$$

其中， $\theta_{pos}$ 是与位置 $p os$ 相关的旋转角度，定义为：
$\theta_{pos} = pos \cdot \frac{1}{10000^{2i/d_{model}}}$
在此公式中， $p os$ 表示位置索引， $d_{model}$ 是模型的维度。

代码实现

基于上述公式，实现 RoPE。

import torch
from torch import nnclass RotaryPositionEmbedding(nn.Module):"""Rotary Position Embedding (RoPE)"""def __init__(self, d_model):super().__init__()self.d_model = d_model## shape: [d_model //2 ]self.theta = 1 / 10000 ** (torch.arange(0, d_model, 2) / d_model)def forward(self, x):"""shape of x: [batch_size, seq_len, d_model]"""batch_size, seq_len, d_model = x.shape# 创建位置编码矩阵 [seq_len, 1]pos = torch.arange(seq_len).unsqueeze(1)# 计算旋转角度 [seq_len, d_model//2]angle = pos * self.thetacos_angle = torch.cos(angle)sin_angle = torch.sin(angle)# 分割奇偶维度x_even = x[..., 0::2]  # 偶数维度 [batch_size, seq_len, d_model//2]x_odd = x[..., 1::2]  # 奇数维度 [batch_size, seq_len, d_model//2]# 应用旋转矩阵x_rotated = torch.zeros_like(x)x_rotated[..., 0::2] = x_even * cos_angle - x_odd * sin_anglex_rotated[..., 1::2] = x_even * sin_angle + x_odd * cos_anglereturn x_rotated

相关解释
1. 初始化参数解释：
  - d_model: 模型维度。
2. forward方法说明：
  - 计算每个位置 $p os$ 的旋转角度 $\theta_{pos}$ 。
  - 对输入 $x$ 的偶数和奇数索引分别进行旋转变换。
  - 返回经过RoPE处理后的向量。

操作步骤	张量形状变化示例
输入数据	[batch, seq_len, d_model]
创建位置编码矩阵	[seq_len, 1]
计算旋转角度	[seq_len, d_model//2]
分割奇偶维度	2x [batch, seq_len, d_model//2]
应用旋转矩阵	2x [batch, seq_len, d_model//2]
合并旋转结果	[batch, seq_len, d_model]

使用示例

测试代码

if __name__ == "__main__":batch_size = 4seq_len = 100d_model = 512rope_layer = RotaryPositionEmbedding(d_model=d_model)input_data = torch.randn(batch_size, seq_len, d_model)output_data = rope_layer(input_data)print(output_data.shape)  # 输出应为 (4, 100, 512)