RoPE旋转位置编码理解

1、概述

在 Transformer 中，注意力本身是位置无关的。

RoPE（Rotary Position Embedding）通过一种非常巧妙的方式，把位置信息耦合到注意力结果中，并且让注意力天然感知相对位置。

把向量的每一对维度当作二维平面，在不同位置上按不同角度旋转。
向量本身不变长
旋转角度由位置决定
点积结果天然会包含相对位置信息，而与绝对位置无关

假设有一段输入中某一个tokens是X，X在注意力机制中有三个向量，分别是Key、Query、Value。

假设 Query 向量的一部分是二维向量：

\begin{matrix} (1) & \begin{matrix} q = [\begin{matrix} x \\ y \end{matrix}] \end{matrix} \end{matrix}

位置为 p，对应一个旋转角度 θ_p。

通过一个旋转矩阵：

\begin{matrix} (2) & \begin{matrix} R (θ_{p}) = [\begin{matrix} \cos θ_{p} & - \sin θ_{p} \\ \sin θ_{p} & \cos θ_{p} \end{matrix}] \end{matrix} \end{matrix}

$q$ :

\begin{matrix} (3) & {\tilde{q}}_{p} = R (θ_{p}) q \end{matrix}

Key向量同理：

\begin{matrix} (4) & {\tilde{k}}_{p} = R (θ_{p 2}) k \end{matrix}

他们的注意力机制结果是：

\begin{matrix} (5) & \tilde{q} p \cdot \tilde{k} p_{2} = (R (θ_{p}) q)^{⊤} (R (θ_{p 2}) k) \end{matrix}

由于旋转矩阵是正交矩阵：

\begin{matrix} (6) & R (θ_{p})^{⊤} R (θ_{p_{2}}) = I, R (θ_{p})^{⊤} = R (- θ_{p}) \end{matrix}

可得：

\begin{matrix} (7) & \tilde{q} p \cdot \tilde{k} p_{2} = (R (θ_{p}) q)^{⊤} (R (θ_{p 2}) k) = q^{⊤} R (θ_{p 2} - θ_{p}) k \end{matrix}

$\theta_{p2} - \theta_{p}$ 相关，且与绝对位置无关。

在 RoPE多维推广中，我们总是把相邻两维当成一组：

\begin{matrix} (8) & (q_{2 i}, q_{2 i + 1}) \end{matrix}

在复数视角下，它可以自然地表示为一个复数：

\begin{matrix} (9) & z_{i} = q_{2 i} + i q_{2 i + 1} \end{matrix}

同理，Key 的第 i 组：

\begin{matrix} (10) & w_{i} = k_{2 i} + i k_{2 i + 1} \end{matrix}

复数的旋转，则是乘以单位复数；

二维旋转矩阵：

\begin{matrix} (11) & \begin{matrix} R (θ) = [\begin{matrix} \cos θ & - \sin θ \\ \sin θ & \cos θ \end{matrix}] \end{matrix} \end{matrix}

在复数里，对应的操作只需要一行：

\begin{matrix} (12) & z ⟶ z \cdot e^{i θ} \end{matrix}

也就是说模长 |z| 不变，只改变相位。

$\theta_p$ 。

于是：

\begin{matrix} (13) & {\tilde{z}}_{p} = z \cdot e^{i θ_{p}} \end{matrix}

Query 和 Key 的复数表示为：

\begin{matrix} (14) & {\tilde{z}}_{m} = z_{q} \cdot e^{i θ_{m}}, {\tilde{w}}_{n} = z_{k} \cdot e^{i θ_{n}} \end{matrix}

分别对应第 m的 Query 和第 n 个位置的Key。

二维实向量点积：

\begin{matrix} (2) & \begin{matrix} {\tilde{z}}_{m} = z_{q} e^{i θ_{m}} = a + b i \\ {\tilde{w}}_{n} = z_{k} e^{i θ_{n}} = c + d i \end{matrix} \end{matrix}

\begin{matrix} (3) & \begin{aligned} {\tilde{q}}_{m} \cdot {\tilde{k}}_{n} = {\tilde{z}}_{m} \cdot \overset{―}{{\tilde{w}}_{n}} & = z_{q} \overset{―}{z_{k}} \cdot e^{i (θ_{m} - θ_{n})} \\ = (a + b i) (c - d i) \\ = a c - a d i + b c i - b d i^{2} \\ = (a c + b d) + (b c - a d) i \end{aligned} \end{matrix}

$(ac+bd)$ 正好是点积的结果，虚部丢弃：

可以看到在复数视角下取共轭，是为了最后可以构造出点积形式，在这里利用复数特性构造共轭复数作为乘数因子：

\begin{matrix} (4) & {\tilde{q}}_{m} \cdot {\tilde{k}}_{n} \end{matrix}

$\Re\left( \tilde{z}_m \cdot \overline{\tilde{w}_n} \right)$ 即可得到ac+bd：

取实部：

\begin{matrix} (5) & ℜ (z_{q} \overset{―}{z_{k}} \cdot e^{i (θ_{m} - θ_{n})}) \end{matrix}

可得到点积（这里的实部为ac+bd，而虚部为bc-ad，正好是叉积式的旋转量）。