RoPE

RoPE

RoPE (Rotary Positional Embedding)，旋转位置编码，是目前广泛采用的一种位置编码方式。

其核心思想是，将高维表征向量分解为多组正交分量，并分别对分量根据位置的不同做不同角度的“旋转”，从而使得编码后同一对向量的点积结果只与二者的相对位置有关，而与绝对位置无关。

要理解 RoPE，可以先从二维情况开始考虑：

矩阵视角

设词向量 $\vec{q}=(x^{(1)},x^{(2)})$ 位于第 $m$ 个位置，则RoPE 会将它旋转 $m\theta$ 角：

$$RoPE(q,m)=RoPE(x^{(1)},x^{(2)},m)=\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{c}{x}^{(1)}\\ x^{(2)}\end{array}\right)$$

对位于第 $n$ 个位置的词向量 $\vec{k}=(x^{(3)},x^{(4)})$，有

$$\begin{array}{rl}& ⟨RoPE(\vec{q},m),RoPE(\vec{k},n)⟩\\ & =\left(\begin{array}{c}{x}^{(1)}\cos m\theta -x^{(2)}\sin m\theta \\ x^{(1)}\sin m\theta +x^{(2)}\cos m\theta \end{array}\right)\cdot \left(\begin{array}{c}{x}^{(3)}\cos n\theta -x^{(4)}\sin n\theta \\ x^{(3)}\sin n\theta +x^{(4)}\cos n\theta \end{array}\right)\\ & =x^{(1)}{x}^{(3)}(\cos m\theta \cos n\theta +\sin m\theta \sin n\theta )\\ & +x^{(1)}{x}^{(4)}(\sin m\theta \cos n\theta -\cos m\theta \sin n\theta )\\ & +x^{(2)}{x}^{(3)}(\cos m\theta \sin n\theta -\sin m\theta \cos n\theta )\\ & +x^{(2)}{x}^{(4)}(\sin m\theta \sin n\theta +\cos m\theta \cos n\theta )\\ & =x^{(1)}{x}^{(3)}\cos (m-n)\theta +x^{(1)}{x}^{(4)}\sin (m-n)\theta \\ & -x^{(2)}{x}^{(3)}\sin (m-n)\theta +x^{(2)}{x}^{(4)}\cos (m-n)\theta \\ & =\left(\begin{array}{c}{x}^{(1)}\cos (m-n)\theta +x^{(2)}\sin (m-n)\theta \\ x^{(1)}\sin (m-n)\theta +x^{(2)}\cos (m-n)\theta \end{array}\right)\cdot \left(\begin{array}{c}{x}^{(3)}\\ x^{(4)}\end{array}\right)\\ & =⟨RoPE(\vec{q},m-n),RoPE(\vec{k},0)⟩\end{array}$$

由此可见同一对向量的点积在 RoPE 后确实只与相对位置有关。

复数视角

从复数视角可以更直观地理解 RoPE

复数中旋转 $\theta$ 角相当于乘上一个 $e^{i\theta }$，把二维向量视为一个复数，则 RoPE 相当于对第 $m$ 个位置的 $q$ 乘上了 $e^{im\theta }$

容易发现 $\mathrm{Re}(q\overline{k})=\vec{q}\cdot \vec{k}$，因此

$$\begin{array}{rl}⟨RoPE(\vec{q},m),RoPE(\vec{k},n)⟩& =\mathrm{Re}(q{e}^{im\theta }\cdot \overline{k}{e}^{-in\theta })\\ & =\mathrm{Re}(q\overline{k}{e}^{i(m-n)\theta })\\ & =⟨RoPE(\vec{q},m-n),RoPE(\vec{k},0)⟩\end{array}$$

高维情况

一般来讲 RoPE 会将 $d$ 维向量分解为 $d/2$ 组二维向量，即将 $(x_1,x_2,\dots ,x_d)$ 分为 $(x_1,x_{1+d/2}),\dots ,(x_{d/2},x_d)$，对于第 $i$ 组，${\theta }_i$ 取 $10000^{-2(i-1)/d}$，接着分别对每一组做旋转变换。

RoPE 不同维度中 $\theta$ 的取值

如果每一组的 $\theta$ 取相同值，那么对于两个相同的 token，它们在位置 $L$ 和位置 $L+\frac{2\pi }{\theta }$ 处经位置编码后的结果是完全相同的，这显然不符合我们的目标。

相反，如果对每一组取不同的 $\theta$，不仅可以使得位置编码失去简单的周期性，还可以让不同维度以不同速度“旋转”：

• 高频维度：随着位置的微小变化，旋转角度会迅速改变，有利于模型捕捉精细的位置差异
• 低频维度：即使位置有很大变化，旋转角度变化也很小，使得模型可以在处理长距离的 token 时依然保持稳定

另外，可以数学证明当 ${\theta }_i$ 按照指数形式分布时，Attention Score 会随着相对距离 $m-n$ 的增加而产生自然的衰减，符合语言模型的直觉：物理距离越远的词，相互之间越不相关

实现

class RoPE(nn.Module):
    def __init__(self, n_embed, base=10000):
        super().__init__()
        self.base = base
        self.n_embed = n_embed
        self.register_buffer('cos_cache', None) # (T, C)
        self.register_buffer('sin_cache', None)

可以将 cos 和 sin 矩阵缓存起来避免重复计算：

def _build_cache(self, x):
    seq_len = x.size(-2)
    if self.cos_cache is not None and seq_len <= self.cos_cache.size(-2):
        return
    theta = 1. / (self.base ** (torch.arange(0, self.n_embed, 2).float() / self.n_embed))
    seq_idx = torch.arange(seq_len, device=x.device)
    idx_theta = torch.outer(seq_idx, theta) # (T, C / 2)
    idx_theta = torch.cat([idx_theta, idx_theta], dim=-1) # （T, C）
    self.cos_cache = idx_theta.cos()
    self.sin_cache = idx_theta.sin()

def forward(self, x):
    # x: (B, T, C)
    seq_len = x.size(-2)
    self._build_cache(x)
    cos = self.cos_cache[:seq_len, :]
    sin = self.sin_cache[:seq_len, :]
 
    # [x_front, x_back] -> [-x_back, x_front]
    x_front, x_back = torch.chunk(x, 2, dim=-1)
    x_shift = torch.cat([-x_back, x_front], dim=-1)
 
    x_rope = x * cos + x_shift * sin
    return x_rope

为什么这样分组？

从代码中可以发现，选择将 $x_i$ 与 $x_{i+d/2}$ 分为一组的好处是对于乘 cos 矩阵和乘 sin 矩阵这两个操作，只需将 $x$ 分为两半即可，不需要进行复杂的交错切片，对 GPU 缓存友好。

Reference

LabML