旋转位置编码（RoPE）详解：从一维序列到三维空间

旋转位置编码（Rotary Position Embedding, RoPE）作为一种新颖的相对位置编码技术，因其能将位置信息以旋转变换的方式融入嵌入向量，并自然地在自注意力机制中体现相对位置关系而备受关注。本文将详细介绍 RoPE 的基本原理、实现方式，并探讨其向三维空间的扩展——3D RoPE。

一、详解旋转位置编码（RoPE）的定义

旋转位置编码（Rotary Position Embedding，RoPE）是一种用于 Transformer 模型的相对位置编码技术，它通过旋转变换将位置信息融入到嵌入向量中，从而增强模型对序列中相对位置信息的捕获能力。

1. 频率参数的定义

假设嵌入向量的维度为 $D$（为偶数）。

首先，定义一组频率参数 ${\omega }_k$，用于确定各个维度上的位置编码频率：

$${\omega }_k=10000^{-2k/D},k=0,1,2,\dots ,\frac{D}{2}-1$$

• 解释：
  • ${\omega }_k$ 是第 $k$ 个频率参数。
  • $D$ 是嵌入维度。
  • $k$ 从 $0$ 到 $\frac{D}{2}-1$，总共有 $\frac{D}{2}$ 个频率参数。
  • 指数中的 $-2k/D$ 确保频率参数在不同的维度上呈指数级变化。

---

2. 旋转矩阵的定义

对于序列中的每个位置 $i$，我们定义一个二维的旋转矩阵。由于嵌入维度为 $D$，我们可以将嵌入向量划分为 $\frac{D}{2}$ 个二维子空间。对每个子空间，都使用一个旋转矩阵。

对于位置 $i$ 上的第 $k$ 个二维子空间，旋转角度为：

$${\theta }_k(i)={\omega }_k\cdot i$$

• 解释：
  • ${\theta }_k(i)$ 是位置 $i$ 上第 $k$ 个子空间的旋转角度。
  • ${\omega }_k$ 是对应的频率参数。
  • $i$ 是位置索引，从 $0$ 开始。

定义位置 $i$ 的旋转矩阵 $R_{{\theta }_k(i)}$：

$$R_{{\theta }_k(i)}=\left(\begin{array}{cc}\cos \left({\theta }_k(i)\right)& -\sin \left({\theta }_k(i)\right)\\ \sin \left({\theta }_k(i)\right)& \cos \left({\theta }_k(i)\right)\end{array}\right)$$

整体的旋转矩阵 $R_{\theta }(i)$ 可以表示为一个块对角矩阵，由 $\frac{D}{2}$ 个 $2\times 2$ 的旋转矩阵构成：

$$R_{\theta }(i)=\left(\begin{array}{cccc}{R}_{{\theta }_0(i)}& & & \\ & R_{{\theta }_1(i)}& & \\ & & \ddots & \\ & & & R_{{\theta }_{D/2-1}(i)}\end{array}\right)$$

• 解释：
  • $R_{\theta }(i)$ 是一个 $D\times D$ 的矩阵，每个 $2\times 2$ 的旋转矩阵对应嵌入向量的一个二维子空间。

---

3. 对嵌入向量进行旋转变换

对于每个位置 $i$ 的嵌入向量 ${\mathbf{x}}_i$，进行旋转变换：

$${\tilde{\mathbf{x}}}_i=R_{\theta }(i)\cdot {\mathbf{x}}_i$$

• 解释：
  • ${\mathbf{x}}_i$ 是位置 $i$ 的嵌入向量，维度为 $D$。
  • ${\tilde{\mathbf{x}}}_i$ 是经过旋转位置编码后的嵌入向量。
  • $R_{\theta }(i)$ 对嵌入向量 ${\mathbf{x}}_i$ 进行旋转，将位置信息编码进向量中。

维度分组：

• 将嵌入向量 ${\mathbf{x}}_i$ 的维度按照每两个维度进行分组。
• 即，第 $k$ 个二维子向量为：

$$\left(\begin{array}{c}{x}_i^{(2k)}\\ x_i^{(2k+1)}\end{array}\right)$$

对每个二维子向量，应用对应的旋转矩阵 $R_{{\theta }_k(i)}$。

---

4. 数学公式的展开

对于位置 $i$ 和第 $k$ 个二维子向量，旋转变换可以具体表示为：
$$\left(\begin{array}{c}{\tilde{x}}_i^{(2k)}\\ {\tilde{x}}_i^{(2k+1)}\end{array}\right)=\left(\begin{array}{cc}\cos \left({\theta }_k(i)\right)& -\sin \left({\theta }_k(i)\right)\\ \sin \left({\theta }_k(i)\right)& \cos \left({\theta }_k(i)\right)\end{array}\right)\left(\begin{array}{c}{x}_i^{(2k)}\\ x_i^{(2k+1)}\end{array}\right)$$

展开后得到：

$$\{\begin{array}{l}{\tilde{x}}_i^{(2k)}=x_i^{(2k)}\cos \left({\theta }_k(i)\right)-x_i^{(2k+1)}\sin \left({\theta }_k(i)\right)\\ {\tilde{x}}_i^{(2k+1)}=x_i^{(2k)}\sin \left({\theta }_k(i)\right)+x_i^{(2k+1)}\cos \left({\theta }_k(i)\right)\end{array}$$

• 解释：
  • 经过旋转变换后，嵌入向量的第 $2k$ 和第 $2k+1$ 维度的分量被重新组合，引入了位置信息。
  • 这相当于将二维向量 $\left(\begin{array}{c}{x}_i^{(2k)}\\ x_i^{(2k+1)}\end{array}\right)$ 绕原点旋转了角度 ${\theta }_k(i)$。

---

5. 旋转变换的意义

通过上述旋转变换，将位置信息以旋转的形式编码进嵌入向量中。

• 相对位置信息：

  • 在后续的自注意力计算中，查询 ${\tilde{\mathbf{q}}}_i$ 和键 ${\tilde{\mathbf{k}}}_j$ 的点积中天然地包含了位置 $i$ 和 $j$ 之间的相对位置信息。
  • 这使得模型能够捕获序列中元素之间的相对关系，而不仅仅是绝对位置。

• 旋转编码的优势：

  • 旋转变换是等距变换，不改变向量的模长。
  • 在保持原始嵌入向量信息的同时，融入了位置信息。

---

6. 实现示例

在实际实现中，通常会采用向量化的方式对整个批次的数据进行处理。以下是实现步骤：

1. 计算频率参数 ${\omega }_k$：

   import torch
   import math
   
   D = embedding_dim  # 嵌入维度，假设为偶数
   half_D = D // 2
   k = torch.arange(half_D, dtype=torch.float32)  # (D/2,)
   omega = 10000 ** (-2 * k / D)  # (D/2,)
   

2. 计算旋转角度 ${\theta }_k(i)$：

   对于输入序列长度为 $L$，位置索引 $i$ 从 $0$ 到 $L-1$：

   positions = torch.arange(L, dtype=torch.float32).unsqueeze(1)  # (L, 1)
   theta = positions * omega.unsqueeze(0)  # (L, D/2)
   

3. 计算 $\sin$ 和 $\cos$：

   sin_theta = torch.sin(theta)  # (L, D/2)
   cos_theta = torch.cos(theta)  # (L, D/2)
   

4. 对嵌入向量进行旋转变换：

   • 假设嵌入向量 $\mathbf{x}$ 的形状为 $(\text{batch\_size},L,D)$：

     x = embeddings  # (batch_size, L, D)
     x = x.to(torch.float32)  # 确保数据类型为浮点型
     

   • 将嵌入向量拆分为偶数和奇数维度：

     x_even = x[:, :, ::2]  # (batch_size, L, D/2)
     x_odd = x[:, :, 1::2]  # (batch_size, L, D/2)
     

   • 调整 $\sin \theta$ 和 $\cos \theta$ 的形状以匹配 $x$：

     sin_theta = sin_theta.unsqueeze(0)  # (1, L, D/2)
     cos_theta = cos_theta.unsqueeze(0)  # (1, L, D/2)
     

   • 应用旋转变换：

     x_rotated_even = x_even * cos_theta - x_odd * sin_theta  # (batch_size, L, D/2)
     x_rotated_odd = x_even * sin_theta + x_odd * cos_theta   # (batch_size, L, D/2)
     

   • 合并旋转后的结果：

     x_rotated = torch.zeros_like(x)
     x_rotated[:, :, ::2] = x_rotated_even
     x_rotated[:, :, 1::2] = x_rotated_odd
     

   • 最终得到旋转编码后的嵌入向量 $\tilde{\mathbf{x}}$：

     • $\tilde{\mathbf{x}}=x_{\text{rotated}}$，形状为 $(\text{batch\_size},L,D)$

---

7. 总结

• 旋转位置编码通过对嵌入向量进行位置相关的旋转变换，将位置信息以乘法的方式融入到向量中，而不是像传统位置编码那样通过加法。

• 这种方法的优点在于：

  • 高效地编码相对位置信息，有助于模型理解序列中的相对关系。
  • 实现简单，仅需少量的张量操作，不引入额外的参数。

---

8. 示例代码

下面提供一个完整的 Python 实现示例：

import torch
import math

def apply_rope(x):
    """
    对输入嵌入向量 x 应用旋转位置编码（RoPE）。

    参数：
    - x: 输入嵌入向量，形状为 (batch_size, seq_len, D)，D 必须为偶数。

    返回：
    - x_rotated: 经过 RoPE 编码的嵌入向量，形状与 x 相同。
    """
    batch_size, seq_len, D = x.shape
    device = x.device
    half_D = D // 2

    # 1. 计算频率参数 omega_k
    k = torch.arange(half_D, dtype=torch.float32, device=device)  # (D/2,)
    omega = 10000 ** (-2 * k / D)  # (D/2,)

    # 2. 计算旋转角度 theta_k(i)
    positions = torch.arange(seq_len, dtype=torch.float32, device=device).unsqueeze(1)  # (seq_len, 1)
    theta = positions * omega.unsqueeze(0)  # (seq_len, D/2)

    # 3. 计算 sin 和 cos
    sin_theta = torch.sin(theta)  # (seq_len, D/2)
    cos_theta = torch.cos(theta)  # (seq_len, D/2)

    # 4. 将嵌入向量拆分为偶数和奇数维度
    x_even = x[:, :, ::2]  # (batch_size, seq_len, D/2)
    x_odd = x[:, :, 1::2]  # (batch_size, seq_len, D/2)

    # 5. 调整 sin_theta 和 cos_theta 形状
    sin_theta = sin_theta.unsqueeze(0)  # (1, seq_len, D/2)
    cos_theta = cos_theta.unsqueeze(0)  # (1, seq_len, D/2)

    # 6. 应用旋转变换
    x_rotated_even = x_even * cos_theta - x_odd * sin_theta
    x_rotated_odd = x_even * sin_theta + x_odd * cos_theta

    # 7. 合并旋转后的结果
    x_rotated = torch.zeros_like(x)
    x_rotated[:, :, ::2] = x_rotated_even
    x_rotated[:, :, 1::2] = x_rotated_odd

    return x_rotated

# 示例使用
batch_size = 2
seq_len = 5
D = 8  # 嵌入维度，必须为偶数
x = torch.randn(batch_size, seq_len, D)  # 随机初始化嵌入向量

x_rope = apply_rope(x)  # 应用 RoPE 编码

---

9. 在自注意力机制中的应用

在 Transformer 模型中，通常在计算自注意力的查询向量 $\mathbf{Q}$ 和键向量 $\mathbf{K}$ 时，对它们应用 RoPE 编码。

• 步骤：

  1. 计算嵌入向量 $\mathbf{X}$ 后，通过线性层得到 $\mathbf{Q}$ 和 $\mathbf{K}$。
  2. 对 $\mathbf{Q}$ 和 $\mathbf{K}$ 应用 RoPE 编码，得到 $\tilde{\mathbf{Q}}$ 和 $\tilde{\mathbf{K}}$。
  3. 使用 $\tilde{\mathbf{Q}}$ 和 $\tilde{\mathbf{K}}$ 计算自注意力得分。

• 注意：

  • RoPE 编码不应用于值向量 $\mathbf{V}$。
  • 在多头注意力的实现中，需要对每一个头的 $\mathbf{Q}$ 和 $\mathbf{K}$ 分别应用 RoPE 编码。

---

10. 结论

• RoPE（旋转位置编码） 通过对嵌入向量进行旋转变换，巧妙地将位置信息融入到向量中，大大增强了模型对相对位置信息的捕获能力。

• 实现简单高效：只需在相关张量上进行几次张量操作，无需额外参数，且兼容现有的 Transformer 架构。

• 实际效果显著：在处理长序列、需要建模相对位置信息的任务中，RoPE 通常能带来性能的提升。

二、扩展到三维：3D RoPE

原始的 RoPE 主要针对一维序列数据，例如文本。然而，在处理三维数据（如三维点云、视频数据或三维图像）时，需要扩展原有的 RoPE 方法，以适应三维空间中的位置编码。这就是 3D RoPE。

---

1. 3D RoPE 的必要性

在处理三维数据时，模型需要捕获输入数据在三维空间中的位置信息。例如，对于点云中的点，需要知道它们在 $x$、$y$、$z$ 三个方向上的位置信息。使用适当的位置信息编码，可以帮助模型更好地理解数据的空间结构，提高模型的性能。

---

2. 3D RoPE 的原理

2.1 三维位置编码

与一维序列中的位置不同，三维数据中的位置由三个坐标表示：$i$, $j$, $k$，对应于空间中的 $x$、$y$、$z$ 轴。因此，需要为每个坐标方向分别引入旋转位置编码。

2.2 多维旋转矩阵

为了在三维空间中应用旋转位置编码，需要使用高维的旋转矩阵。对于每个位置 $(i,j,k)$，我们定义三个旋转角度：

$$\begin{gathered} {\theta }_x={\omega }_x\cdot i \\ {\theta }_y={\omega }_y\cdot j \\ {\theta }_z={\omega }_z\cdot k \end{gathered}$$

其中，${\omega }_x$、${\omega }_y$、${\omega }_z$ 是在三个坐标轴上预定义的频率参数。

2.3 定义旋转操作

由于嵌入向量的维度 $D$ 通常为偶数，我们可以将其划分为三元组，每三个维度为一组。对于每个三维子空间，我们应用三维空间的旋转操作。

三维旋转变换需要使用三维旋转矩阵。三维旋转矩阵可以表示为绕 $x$、$y$、$z$ 轴的旋转矩阵的组合。对于简化计算，我们可以分别在每个二维平面上进行旋转。

---

3. 3D RoPE 的实现方法

3.1 频率参数的定义

首先，定义一组频率参数：

$${\omega }_k=10000^{-2k/D},k=0,1,2,\dots ,\frac{D}{3}-1$$

这里，$D$ 是嵌入向量的维度，假设可以被 3 整除。

3.2 计算旋转角度

对于位置 $(i,j,k)$，旋转角度定义为：

$$\begin{gathered} {\theta }_k^x={\omega }_k\cdot i \\ {\theta }_k^y={\omega }_k\cdot j \\ {\theta }_k^z={\omega }_k\cdot k \end{gathered}$$

3.3 应用旋转变换

将嵌入向量 ${\mathbf{x}}_n$（$n$ 表示样本索引）分为三部分，每部分维度为 $D/3$：

$${\mathbf{x}}_n=[{\mathbf{x}}_n^x,{\mathbf{x}}_n^y,{\mathbf{x}}_n^z]$$

对于每个部分，应用对应的旋转变换。例如，对于 ${\mathbf{x}}_n^x$，应用绕 $x$ 轴的旋转；对于 ${\mathbf{x}}_n^y$，应用绕 $y$ 轴的旋转；对于 ${\mathbf{x}}_n^z$，应用绕 $z$ 轴的旋转。

旋转公式如下（以 $x$ 轴为例）：
$$\left(\begin{array}{c}{\tilde{x}}_n^{(3k)}\\ {\tilde{x}}_n^{(3k+1)}\\ {\tilde{x}}_n^{(3k+2)}\end{array}\right)=R_{{\theta }_k}^x\left(\begin{array}{c}{x}_n^{(3k)}\\ x_n^{(3k+1)}\\ x_n^{(3k+2)}\end{array}\right)$$

其中，$R_{{\theta }_k}^x$ 是绕 $x$ 轴的旋转矩阵，定义为：

$$R_{{\theta }_k}^x=\left(\begin{array}{ccc}1& 0& 0\\ 0& \cos ({\theta }_k^x)& -\sin ({\theta }_k^x)\\ 0& \sin ({\theta }_k^x)& \cos ({\theta }_k^x)\end{array}\right)$$

同理，可以定义绕 $y$ 轴和 $z$ 轴的旋转矩阵 $R_{{\theta }_k}^y$ 和 $R_{{\theta }_k}^z$。

3.4 综合旋转

最终的旋转变换可以表示为旋转矩阵的组合：

$${\tilde{\mathbf{x}}}_n=R_{{\theta }_k}^z\cdot R_{{\theta }_k}^y\cdot R_{{\theta }_k}^x\cdot {\mathbf{x}}_n$$

这种结合方式可以将三个方向上的位置信息编码进嵌入向量中。

---

4. 3D RoPE 的实现示例

4.1 PyTorch 实现

下面是一个使用 PyTorch 实现 3D RoPE 的示例代码：

import torch
import math

def apply_3d_rope(x, positions):
    """
    对输入嵌入向量 x 应用 3D RoPE 编码。

    参数：
    - x: 输入嵌入向量，形状为 (batch_size, seq_len, D)，其中 D 能被 6 整除。
    - positions: 位置张量，形状为 (batch_size, seq_len, 3)，包含 (i, j, k) 位置。

    返回：
    - x_rotated: 经过 3D RoPE 编码的嵌入向量，形状与 x 相同。
    """
    batch_size, seq_len, D = x.shape
    device = x.device
    assert D % 6 == 0, "Embedding dimension must be divisible by 6 for 3D RoPE."
    D_div_6 = D // 6  # 每个坐标轴各自占据两个部分，共6个部分

    # 1. 计算频率参数 omega_k
    k = torch.arange(D_div_6, dtype=torch.float32, device=device)  # (D/6,)
    omega = 10000 ** (-2 * k / D_div_6)  # (D/6,)

    # 2. 获取位置坐标 i, j, k
    pos_i = positions[:, :, 0]  # (batch_size, seq_len)
    pos_j = positions[:, :, 1]
    pos_k = positions[:, :, 2]

    # 3. 计算旋转角度 theta_k^x, theta_k^y, theta_k^z
    theta_x = pos_i.unsqueeze(-1) * omega  # (batch_size, seq_len, D/6)
    theta_y = pos_j.unsqueeze(-1) * omega
    theta_z = pos_k.unsqueeze(-1) * omega

    # 4. 计算 sin 和 cos
    sin_theta_x = torch.sin(theta_x)
    cos_theta_x = torch.cos(theta_x)
    sin_theta_y = torch.sin(theta_y)
    cos_theta_y = torch.cos(theta_y)
    sin_theta_z = torch.sin(theta_z)
    cos_theta_z = torch.cos(theta_z)

    # 5. 将嵌入向量拆分为六部分
    x_x1 = x[:, :, :D_div_6]        # (batch_size, seq_len, D/6)
    x_x2 = x[:, :, D_div_6:2*D_div_6]
    x_y1 = x[:, :, 2*D_div_6:3*D_div_6]
    x_y2 = x[:, :, 3*D_div_6:4*D_div_6]
    x_z1 = x[:, :, 4*D_div_6:5*D_div_6]
    x_z2 = x[:, :, 5*D_div_6:]

    # 6. 对每一对部分应用对应的旋转变换

    # 绕 x 轴旋转
    x_rotated_x1 = x_x1 * cos_theta_x - x_x2 * sin_theta_x
    x_rotated_x2 = x_x1 * sin_theta_x + x_x2 * cos_theta_x

    # 绕 y 轴旋转
    x_rotated_y1 = x_y1 * cos_theta_y - x_y2 * sin_theta_y
    x_rotated_y2 = x_y1 * sin_theta_y + x_y2 * cos_theta_y

    # 绕 z 轴旋转
    x_rotated_z1 = x_z1 * cos_theta_z - x_z2 * sin_theta_z
    x_rotated_z2 = x_z1 * sin_theta_z + x_z2 * cos_theta_z

    # 7. 合并旋转后的结果
    x_rotated = torch.cat([x_rotated_x1, x_rotated_x2, x_rotated_y1, x_rotated_y2, x_rotated_z1, x_rotated_z2], dim=-1)
    return x_rotated

# 示例参数
batch_size = 2
seq_len = 4
D = 24  # 嵌入维度，必须能被 6 整除

# 随机初始化嵌入向量
torch.manual_seed(42)
x = torch.randn(batch_size, seq_len, D)

# 生成位置张量，假设每个位置都有 (i, j, k) 坐标
positions = torch.randint(0, 10, (batch_size, seq_len, 3))

# 应用 3D RoPE 编码
x_rope = apply_3d_rope(x, positions)

# 输出结果
print("输入嵌入向量形状：", x.shape)
print("位置张量形状：", positions.shape)
print("经过 3D RoPE 编码后的嵌入向量形状：", x_rope.shape)

执行上述代码，将输出：

输入嵌入向量形状： torch.Size([2, 4, 24])
位置张量形状： torch.Size([2, 4, 3])
经过 3D RoPE 编码后的嵌入向量形状： torch.Size([2, 4, 24])

这表明输入的嵌入向量经过 3D RoPE 编码后，形状保持不变，可以继续用于后续的模型计算。

---

5. 3D RoPE 在应用中的优势

• 捕获空间关系：3D RoPE 能够有效地捕获三维空间中数据的相对位置关系，增强模型对空间结构的理解。
• 适用于三维数据：在处理点云、三维物体、空间序列等三维数据时，3D RoPE 提供了一种自然且有效的位置编码方法。
• 兼容现有模型：3D RoPE 可以与 Transformer 等模型无缝结合，只需在输入嵌入向量时应用适当的旋转编码。

---

6. 应用场景

• 点云处理：3D RoPE 可用于点云分类、分割和目标检测等任务，帮助模型理解点云中的三维结构。
• 视频理解：在时空 Transformer 中，将时间和空间维度都考虑进去，3D RoPE 可以编码时空位置信息。
• 医学影像：对于三维医学影像，如 MRI、CT 扫描等，3D RoPE 有助于模型捕获空间位置信息，提高诊断和分析的准确性。

---

7. 结论

3D RoPE 是对原始旋转位置编码的有益扩展，适用于需要建模三维位置信息的任务。通过对嵌入向量进行三维旋转变换，模型能够更好地理解三维空间中的数据分布和结构，提高在三维数据处理任务中的性能。

---

参考资料：

• Su, J., et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv preprint arXiv:2104.09864.