RoPE 位置编码
paperreading
本文字数:3.2k 字 | 阅读时长 ≈ 14 min

RoPE 位置编码

paperreading
本文字数:3.2k 字 | 阅读时长 ≈ 14 min

旋转位置编码(简称 RoPE)是一种用于 Transformer 类模型的相对位置编码方法,在像 Llama 等模型中被广泛应用。

1. 1D RoPE

RoPE 是一种旋转式位置编码方法。它通过对查询和键(Query, Key)向量施加一个与位置相关的旋转变换,从而隐式地引入位置信息。

给定位置 p 上的 Query 向量 $\mathbf{q}_p \in \mathbb{R}^{d}$,Key 向量 $\mathbf{k}_p \in \mathbb{R}^{d}$,我们将它们视为一组二维向量的拼接:

$$
\mathbf{q}_p = [q_0, q_1, q_2, q_3, …, q_{d-2}, q_{d-1}] \rightarrow [(q_0, q_1), (q_2, q_3), …, (q_{d-2}, q_{d-1})]
$$

我们把每两个维度 $(q_{2i}, q_{2i+1})$ 看作一个二维坐标或者复数:$z_i^{(p)} = q_{2i}^{(p)} + i \cdot q_{2i+1}^{(p)}$

为了方便下文阅读,我们引入复数的一个简单性质:复数向量 $z_i^{(p)}$ 乘以单位模复数 $e^{i \theta_p}$,其中 $\theta_p$ 是与位置 $p$ 相关的角度。具体来说,复数 $z_i^{(p)}$ 表示在位置 $p$ 上的某个点或信息,而 $e^{i \theta_p}$ 是一个单位模复数,模长为 1,表示一个旋转因子,其极坐标表示为 $e^{i \theta_p} = \cos(\theta_p) + i \sin(\theta_p)$。将复数向量 $z_i^{(p)}$ 乘以旋转因子 $e^{i \theta_p}$,相当于在复平面上对 $z_i^{(p)}$ 进行一个旋转操作,旋转角度为 $\theta_p$。

现在我们继续上文,我们将 embeeding 中的每两维看做一个二维坐标,对每个复数向量 $z_i^{(p)}$,我们乘上一个单位模的复数 $e^{i \theta_p}$,其角度与位置有关:

$$
z_i^{(p)} \rightarrow z_i^{(p)} \cdot e^{i \theta_p} = z_i^{(p)} \cdot \left( \cos(\theta_p) + i \cdot \sin(\theta_p) \right)
$$

介绍到这里,我们知道了 RoPE 的核心思想是通过旋转变换来引入位置信息,那么是如何融入到 token 的 attention 当中的呢?正常来讲 token 的 attention 计算如下所示:

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

为了引入位置信息,我们通过旋转变换来处理 QueryKey 向量。经过 RoPE 变换后的 Query 和 Key 分别为 $Q’_p = \text{RoPE}(Q_p)$ 和 $K’_p = \text{RoPE}(K_p)$。RoPE 的核心思想是将位置信息通过旋转因子融入到每个 token 的表示中。旋转因子的形式为:

$$
R_p = e^{i \theta_p}
$$

在计算点积时,我们将旋转后的 Query 和 Key 进行点积,得到:

$$
Q’_p \cdot K’_q = \left( Q_p \cdot e^{i \theta_p} \right)^T \left( K_q \cdot e^{i \theta_q} \right)
$$

可以将旋转因子分解出来:

$$
Q’_p \cdot K’_q = Q_p^T \cdot e^{i \theta_p} \cdot e^{-i \theta_q} \cdot K_q Q_p^T \cdot e^{i (\theta_p - \theta_q)} \cdot K_q = Q_p^T \cdot R_{q-p} \cdot K_q
$$

$R_{q-p}$ 表示的是 相对位置 的旋转因子。通过这种方式,RoPE 使得每个位置的位置信息在计算时通过旋转因子结合进了 Query 和 Key 的表示,而旋转因子 $R_{q-p}$ 完全由 token p 和 token q 的相对位置决定。

ok,现在我们对 RoPE 的核心思想有了一个大致的了解,上面提到了一个旋转角度的问题,他是如何定义的呢,下面我们来解释一下。

对于 RoPE 中第 i 个二维子向量(每两个维度一组),其旋转角度定义为:
$$
\theta_i(p) = p \cdot \omega_i = p \cdot \frac{1}{10000^{\frac{2i}{d}}}
$$

为什么这样设计角度?

现在我们举个例子,假设 token 位置为 p = 5,维度为 d = 8,总共 $\frac{d}{2} = 4$ 个二维子向量。

子向量编号 i 维度对 (2i, 2i+1) 频率 $ω_i$ 旋转角度 $θ_i(p) = p × ω_i$
0 (0, 1) 1 5
1 (2, 3) 1 / 100 0.05
2 (4, 5) 1 / 316.23 ≈ 0.0158
3 (6, 7) 1 / 1778.28 ≈ 0.00281

这里的是绝对位置,即 token 在序列中的位置,当进行 attention 的时候,是两个位置的向量进行计算。那 RoPE 融入 attention 是怎么体现的?直接看下表就很好理解了。

注意我这里说的是绝对位置表达相对位置,RoPE 的 attention score 是纯相对位置驱动的,但因为 RoPE 会对原始向量进行绝对位置旋转,模型下游仍能“感知”到某种弱的绝对位置信息。因此总结一下,RoPE 可以显式编码相对位置的同时,隐式编码绝对位置信息,这就是 RoPE 我感觉最牛逼的地方了。

子向量 i 维度对 (2i,2i+1) 频率 $ω_i$ Query角度 $θ_i(5)$ Key角度 $θ_i(3)$ 相对旋转角 $θ_i(3{-}5)$ Attention Score 形式
0 (0,1) 1 5 3 -2 $Q_5^{(0)\top} R_{-2} K_3^{(0)}$
1 (2,3) 1 / 100 0.05 0.03 -0.02 $Q_5^{(1)\top} R_{-0.02} K_3^{(1)}$
2 (4,5) 1 / 316.23 ≈0.0158 ≈0.00949 ≈-0.00631 $Q_5^{(2)\top} R_{-0.00631} K_3^{(2)}$
3 (6,7) 1 / 1778.28 ≈0.00281 ≈0.00169 ≈-0.00113 $Q_5^{(3)\top} R_{-0.00113} K_3^{(3)}$

下面举个例子说明一下

import torch
import math
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np

# 中文字体设置(macOS)
matplotlib.rcParams['font.family'] = "Heiti TC"
matplotlib.rcParams['axes.unicode_minus'] = False

def apply_rope(x: torch.Tensor, position: int):
    D = x.shape[0]
    assert D % 2 == 0
    x_roped = x.clone()
    for i in range(0, D, 2):
        freq_idx = i // 2
        theta = position / (10000 ** (2 * freq_idx / D))
        cos_theta = math.cos(theta)
        sin_theta = math.sin(theta)
        xi, xi1 = x[i], x[i + 1]
        x_roped[i]     = cos_theta * xi - sin_theta * xi1
        x_roped[i + 1] = sin_theta * xi + cos_theta * xi1
    return x_roped

# 原始 token 向量
# x = torch.arange(12).float()
x = torch.ones(12).float() * 5
positions = [0, 1, 2, 5, 10]
roped_vectors = [apply_rope(x, pos) for pos in positions]

# 展示所有偶数维度对
D = x.shape[0]
num_pairs = D // 2
nrows, ncols = 2, 3
fig, axes = plt.subplots(nrows, ncols, figsize=(6 * ncols, 5 * nrows))
axes = axes.flatten()
colors = cm.viridis(np.linspace(0, 1, len(positions)))

for plot_idx in range(num_pairs):
    i = plot_idx * 2
    ax = axes[plot_idx]
    ax.set_title(f"维度对 ({i},{i+1}) 的旋转效果")
    ax.set_xlim(-15, 15)
    ax.set_ylim(-15, 15)
    ax.axhline(0, color='gray', lw=0.5)
    ax.axvline(0, color='gray', lw=0.5)
    ax.set_aspect('equal')

    handles = []

    # 原始向量
    xi, xi1 = x[i], x[i + 1]
    q = ax.quiver(0, 0, xi.item(), xi1.item(), angles='xy', scale_units='xy', scale=1,
                  color='black', width=0.015, label='原始向量', zorder=10)
    handles.append(q)

    # 各位置旋转向量
    for j, pos in enumerate(positions):
        xi_r, xi1_r = roped_vectors[j][i], roped_vectors[j][i + 1]
        q = ax.quiver(0, 0, xi_r.item(), xi1_r.item(), angles='xy', scale_units='xy', scale=1,
                      color=colors[j], width=0.01, label=f'pos={pos}')
        handles.append(q)

    ax.legend(handles=handles, loc='upper left')

# 隐藏多余子图
for idx in range(num_pairs, nrows * ncols):
    axes[idx].axis('off')

plt.tight_layout()
plt.savefig('rope.png', dpi=300)  # 保存为 PNG 文件
plt.show()

可视化结果

这是呈现的结果,在RoPE中,theta的计算方式如下$\theta_j = \frac{\text{position}}{10000^{2j/D}}$

采用“指数衰减”的原因?

将频率(角度步长)做成指数衰减,可以确保:既可以编码“很短距离的”相对位置信息(如某词的邻居是谁)——靠低维做细致区分,也能反映“很长距离的”全局顺序信息——靠高维做粗略编码,这与原始Transformer的正弦位置编码里的“多尺度频率”思路与Motivation是一样的!

为什么同一 pos 的 embedding 在不同维度上旋转角度不同?

因为每个维度的旋转角度 $\theta_j$ 是根据该维度的频率 $j$ 计算的。低频维度(小 $j$)的旋转角度较大,能更细致地捕捉位置变化;高频维度(大 $j$)的旋转角度较小,能更平滑地捕捉全局顺序信息。
这种设计使得模型能够同时感知局部和全局的位置信息。如果所有维度同角度旋转,RoPE的“不同空间编码多尺度位置关系”的能力就没了,模型只能获得极弱的位置感知力,信息容量极低,也失去跨序列长度泛化的意义,这就是必须多频率的本质原因!

2. 2D RoPE

Transformer 在处理图像时,每个 patch/pixel 有二维空间坐标 $(pos_x, pos_y)$。1D RoPE 编码一维序列的位置,而 2D RoPE 则致力于高效表示 二维网格上每个点的空间位置信息

核心思想是将 embedding 按 x, y 两个方向各自“分管一半”,分别用 pos_xpos_y 采用 1D RoPE 的方法进行独立旋转。

设 embedding 维度为 $D$, 分成前后各一半:

每半仍按一维 RoPE 那样,分成若干组二维向量(每组为 $(x_{i}, x_{i+1})$)各自以多尺度频率旋转。对 $D$ 维 embedding $x$,$pos_x, pos_y$ 为横纵位置,有:

前半:$i = 0,2,4,\dots, D/2-2$
$$
\theta^{(x)}_j = \frac{pos_x}{10000^{\frac{2j}/{(D/2)}}}
$$

$$
\begin{bmatrix} x’_{i} \\ x’_{i+1} \end{bmatrix} =
\begin{bmatrix}
\cos \theta_j^{(x)} & -\sin \theta_j^{(x)} \\
\sin \theta_j^{(x)} & \cos \theta_j^{(x)}
\end{bmatrix}
\begin{bmatrix} x_{i} \\ x_{i+1} \end{bmatrix}
$$

其中 $j = i/2$,后半:$i = D/2, D/2+2, …, D-2$
$$
\theta^{(y)}_j = \frac{pos_y}{10000^{\frac{2j}/{(D/2)}}}
$$

$$
\begin{bmatrix} x’_{i} \\ x’_{i+1} \end{bmatrix} =
\begin{bmatrix}
\cos \theta_j^{(y)} & -\sin \theta_j^{(y)} \\
\sin \theta_j^{(y)} & \cos \theta_j^{(y)}
\end{bmatrix}
\begin{bmatrix} x_{i} \\ x_{i+1} \end{bmatrix}
$$

其中 $j = (i - D/2)/2$,这里每一半仍然有自己的多尺度频率,依赖于二维坐标的不同分部。

复数形式

类比 1D RoPE:每个二维小组可用 $z^{(x)}_j = x_{2j} + i x_{2j+1}$,旋转为 $z^{(x)} _{j} \cdot e^{i\theta _{j}{(x)}}$,$z{(y)}_j = x_{D/2+2j} + i x_{D/2+2j+1}$,旋转 $z^{(y)} _{j} \cdot e^{i\theta _{j}^{(y)}}$。

import torch
import math
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np

# 中文字体设置(macOS,可注释掉)
matplotlib.rcParams['font.family'] = "Heiti TC"
matplotlib.rcParams['axes.unicode_minus'] = False

def apply_2d_rope(x: torch.Tensor, pos_x: int, pos_y: int):
    """
    x: [D], D 为 4 的倍数
    pos_x, pos_y: 坐标
    前一半用 pos_x, 后一半用 pos_y
    """
    D = x.shape[0]
    assert D % 4 == 0, "embedding dim must be a mul of 4"
    x_roped = x.clone()
    half_D = D // 2
    # 前半: 用pos_x
    for i in range(0, half_D, 2):
        freq_idx = i // 2
        theta = pos_x / (10000 ** (2 * freq_idx / half_D))
        cos_theta = math.cos(theta)
        sin_theta = math.sin(theta)
        xi, xi1 = x[i], x[i + 1]
        x_roped[i]     = cos_theta * xi - sin_theta * xi1
        x_roped[i + 1] = sin_theta * xi + cos_theta * xi1
    # 后半: 用pos_y
    for i in range(half_D, D, 2):
        freq_idx = (i - half_D) // 2
        theta = pos_y / (10000 ** (2 * freq_idx / half_D))
        cos_theta = math.cos(theta)
        sin_theta = math.sin(theta)
        xi, xi1 = x[i], x[i + 1]
        x_roped[i]     = cos_theta * xi - sin_theta * xi1
        x_roped[i + 1] = sin_theta * xi + cos_theta * xi1
    return x_roped

# 设置 embedding
x = torch.ones(12).float() * 4
# 测试的 (pos_x, pos_y) 组合
positions = [(0, 0), (2, 0), (0, 2), (2, 2), (4, 4), (8, 4)]
roped_vectors = [apply_2d_rope(x, px, py) for px, py in positions]

# 展示所有维度对
D = x.shape[0]
num_pairs = D // 2
nrows, ncols = 2, 3     # 可根据 embedding 更大可增大
fig, axes = plt.subplots(nrows, ncols, figsize=(6 * ncols, 5 * nrows))
axes = axes.flatten()
colors = cm.viridis(np.linspace(0, 1, len(positions)))

for plot_idx in range(num_pairs):
    i = plot_idx * 2
    ax = axes[plot_idx]
    if i < D // 2:
        title = f"维度({i},{i+1}) (受 pos_x 影响)"
    else:
        title = f"维度({i},{i+1}) (受 pos_y 影响)"
    ax.set_title(title)
    ax.set_xlim(-15, 15)
    ax.set_ylim(-15, 15)
    ax.axhline(0, color='gray', lw=0.5)
    ax.axvline(0, color='gray', lw=0.5)
    ax.set_aspect('equal')

    handles = []

    # 原始向量
    xi, xi1 = x[i], x[i + 1]
    q = ax.quiver(0, 0, xi.item(), xi1.item(), angles='xy', scale_units='xy', scale=1,
                  color='black', width=0.015, label='原始向量', zorder=10)
    handles.append(q)

    # 各 (pos_x, pos_y) 旋转后的向量
    for j, (px, py) in enumerate(positions):
        xi_r, xi1_r = roped_vectors[j][i], roped_vectors[j][i + 1]
        pos_label = f'({px},{py})'
        q = ax.quiver(0, 0, xi_r.item(), xi1_r.item(), angles='xy', scale_units='xy', scale=1,
                      color=colors[j], width=0.01, label=pos_label)
        handles.append(q)

    ax.legend(handles=handles, loc='upper left', fontsize=8)

# 隐藏多余子图
for idx in range(num_pairs, nrows * ncols):
    axes[idx].axis('off')

plt.tight_layout()
plt.savefig('2d_rope.png', dpi=300)
plt.show()

可视化结果

图中每一子图展示embedding一个二维对(如第0/1维),不同颜色箭头代表不同(x, y)格的旋转结果。

和1D RoPE的区别/类比

1D RoPE 2D RoPE
输入数据 一维序列 二维网格
坐标功能 单坐标pos 横纵坐标 $(pos_x,pos_y)$
旋转规则 全部分组用pos 一半用pos_x,另一半用pos_y
能力 顺序相关(文本/序列) 空间相关(图像/表格/patch)