从李群的视角看 ROPE 旋转位置编码

该博客仍在施工中；This is still a Work in Progress (WIP)

最近比较闲，学了点数学，李群李代数什么的，然后一下就联想到了ROPE，所以这次系统地从数学视角回顾了一下。我也是边写边查漏补缺，所以我尽量写得容易阅读并且适合学习，不仅仅是当成自己的笔记，可以带点科普性质，希望不那么熟悉相关内容的同学也能流畅阅读。

群

学过群论的同学都知道，我们可以用群来描述旋转。我们以三角形的旋转群 $C3$ 作为例子，它包含了旋转 0°、旋转 120° 和 旋转 240° 这三个操作。群需要满足四个性质：

1. 封闭性：先旋转 120°，再旋转 120°，最后等于旋转 240°，这个操作仍然是原操作集合的一员。

2. 单位元：需要有一个无用的操作（旋转 0°）

3. 逆元：任何操作都有逆操作，比如旋转 120°的逆操作是旋转 240°

4. 结合律：$(A\cdot B)\cdot C = A \cdot (B \cdot C)$

但是，这是一种离散的有限群。如果我们考虑一个圆的旋转群呢？我们可以旋转任意小的角度。这种“连续”的群就称为李群

李群

从定义上看，李群是一个集合，它既是一个群、又是一个光滑流形。比如说二维平面上的所有旋转矩阵 $\begin{pmatrix}\cos\theta & -\sin\theta \\ \sin\theta & \cos \theta\end{pmatrix}$ 就构成一个李群 $SO(2)$

到这里大家肯定知道，为什么我会说从李群的视角看 ROPE 了。

李代数

为了研究这个连续的李群，我们选择研究它在单位元附近的微元。毕竟只需要掌握了旋转 0.00…01°，多次应用就可以得到一个更大的旋转。我们将李群 $G$ 在单位元 $e$ 处的切空间 $T_eG$ 称为李代数 $\mathfrak{g}$。李代数就像是李群的微分版本，我们以旋转群 $SO(N)$ 为例来说明。

任取一个李群元素 $R(t) \in SO(N)$，我们现在考虑李代数的元素 $X$。如果说李群元素 $R(t)$ 是旋转角度，那么李代数元素 $X$ 就可以理解为角速度。有方程：

$$\dot{R}(t) = R(t) \cdot X$$

其中 $X$ 代表的是相对于单位元的变化率，对于 $R(t)$ 而言，它的变化率需要把 $X$ 迁移到自己的坐标系上。

解微分方程，我们得到 $R(t) = e^{tX}$。

所以，对于 $SO$ 这个李群，它的李代数 $\mathfrak{so}$ 中的元素的矩阵指数，就是李群中的元素。有一种积分的感觉。我们可以对 $R(t)^TR(t) = I$ 关于时间 $t$ 求导：

$$\begin{aligned} \dot{R}(t)^TR(t) + R(t)^T\dot{R}(t) = \dot{I} = 0 \end{aligned}$$

代入 $t=0$，又因为我们知道 $R(0) = I$，$X = \dot{R}(0)$，所以：

$$\begin{aligned} \dot{R}(0)^T \cdot I + I^T \cdot \dot{R}(0) &= 0 \\ X^T + X &= 0 \\ X^T &= -X \end{aligned}$$

所以，旋转群 $SO(N)$ 的正交性可以导出它的李代数 $\mathfrak{so}(N)$ 的元素必须是反对称的。

李括号

李代数是一个向量空间，基础配置只有向量元素间的加法。但对于群而言，加法不是自然的，因为加法总满足交换律（ $+1+2$ 与 $+2 +1$ 相同），群上的操作更像是不可交换的乘法，例如先翻转再旋转，和先旋转再翻转，会得到不同的结果——这就相当于是先乘翻转算子 $X$ 还是先乘旋转算子 $Y$ 的区别。因此，我们希望给李代数也安装上乘法。不过挪威数学家索菲斯 · 李告诉我们，我们只需要衡量乘法操作不可交换的程度即可，称为李括号：

$$[X, Y] = XY - YX$$

怎么理解这个李括号呢？按照我的理解，它其实只是群操作在微观的切空间的投影。

设我们有一个群 $G$，并选定一个群元素 $g \in G$。我们定义一个映射 $\Psi_g: G \rightarrow G$，它的作用规则是：

$$\Psi_g(x) = gxg^{-1}$$

它被称为共轭作用。该映射固定了单位元，$\Psi_g(e) = geg^{-1} = e$，我们可以研究单位元附近加上微扰的变化，也就是单位元 $e$ 处的切空间，即李代数 $\mathfrak{g}$。

不妨将这个微元当作一条经过单位元 $e$ 的曲线 $x(t)$，其中 $x(0) = e$，它的导数速度向量是 $Y = \dot{x}(0) \in \mathfrak{g}$

考虑它的共轭作用

$$\Psi_g(x(t)) = g \cdot x(t) \cdot g^{-1}$$

对 $t$ 求导并令 $t=0$：

$$\frac{d}{dt} \big|_{t=0} (g \cdot x(t) \cdot g^{-1}) = g \cdot \dot{x}(0) \cdot g^{-1} = gYg^{-1}$$

这个新向量 $gYg^{-1}$ 也在李代数 $\mathfrak{g}$ 中，因为李代数对共轭封闭。

我们可以将这个映射视为一个线性算子，实际上它被称为伴随表示:

$$\text{Ad}_g(Y) = gYg^{-1}$$

伴随表示 $\text{Ad}_g$ 描述了群元素 $g$ 如何旋转或拉伸整个李代数空间。

由于李群是光滑的，我们可以进一步求导，考虑 $g$ 的微扰如何影响伴随表示算子。同样，我们视 $g$ 为随时间 $s$ 变化的曲线 $g(s)$，且 $g(0) = e$，速度向量为 $X= \dot{g}(0) \in \mathfrak{g}$

对 $\text{Ad}_{g(s)}(Y) = g(s)Yg(s)^{-1}$ 关于 $s$ 求导，并令 $s=0$

$$\begin{aligned} \frac{d}{ds}\big|_{s=0}(g(s)Yg(s)^{-1}) &= \dot{g}(0)Yg(0)^{-1} + g(0)Y(\frac{d}{ds}g(s)^{-1})\big|_{s=0} \\ &= X\cdot Y \cdot I + I \cdot Y \cdot (-g(0)^{-1}\dot{g}(0)g(0)^{-1}) \\ &=XY-YX \\ &= [X, Y] \end{aligned}$$

所以，李括号正是共轭作用的微分的微分。

我可以举几个实际的例子来帮助理解。以 $SO(2)$ 为例，若 $g$ 是李群的元素，代表旋转 $\theta$ 度；$Y$ 是李代数的元素，代表了角速度； 那么 $\text{Ad}_g(Y)$ 就是这个角速度经过了 $g$ 的变换，是从 $\theta$ 开始的角速度。$\text{ad}_XY$ 就是当 $\theta$ 无穷小时，这个角速度的变化率。

你会发现在 $SO(2)$ 中，好像 $\text{Ad}_g(Y)$ 这步是无意义的，不管从哪个 $\theta$ 开始，角速度都一样。因为 $SO(2)$ 是一个交换群（阿贝尔群），我们先转 $30°$ 再转 $60°$ 与先转 $60°$ 再转 $30°$ 是一样的。对于交换群而言，共轭作用等于恒等映射：

$$\Psi_g(h) = h$$

从而，$\text{Ad}_g$ 变成了恒等算子；$\text{ad}_X(Y) = 0$，李括号为 0，交换操作无影响。

但在 $SO(3)$ 中就不一样了。我们假设 $g$ 是绕 Z 轴旋转 $\theta$， $Y$ 是绕 $X$ 轴的角速度，那么 $\text{Ad}_g(Y)$ 就会把这个绕 $X$ 轴的角速度绕 $Z$ 轴旋转 $\theta$。如果 $\theta=90°$，那么就变成了绕 $Y$ 轴的角速度，角速度矢量旋转了 $90°$。$\text{ad}_X(Y)$ 衡量的就是，当 $\theta$ 无限小时，这个角速度矢量的变化，也可以称为角加速度。

综上，李括号可以看作是群性质在微分切空间上的投影。例如，李括号还要求具有反对称性，即 $[X, Y] = -[Y, X]$，这个也很好理解，就来源于群乘法的不可交换性嘛。自然如果李括号恒为零，那就说明是个交换群。

李括号还要求双线性。这是因为李括号所操作的李代数的元素本来就在切空间上，作为一个向量空间自然有线性，李括号的定义 $[X, Y] = XY - YX$ 自然继承这种线性结构。

最后，回顾群的性质，还有一条结合律。群的结合律能够导出，共轭作用 $\Psi_g$ 是群的自同构，即：

$$\Psi_g(xy) = g(xy)g^{-1} = (gxg^{-1})(gyg^{-1}) = \Psi_g(x)\Psi_g(y)$$

并且，它的微分 $\text{Ad}_g$ 也是李代数的自同构：

$$\begin{aligned} \text{Ad}_g([Y, Z]) &= g[Y, Z]g^{-1} = g(YZ -ZY)g^{-1} \\ &= g(YZ)g^{-1} - g(ZY)g^{-1} \\ &= gYg^{-1}gZg^{-1} - gZg^{-1}gYg^{-1} \\ &= \text{Ad}_g(Y)\text{Ad}_g(Z) - \text{Ad}_g(Z)\text{Ad}_g(Y) \\ &= [\text{Ad}_g(Y), \text{Ad}_g(Z)] \end{aligned}$$

我们对上式关于 $g$ 求导，其速度向量为 $X$，则左式为：

$$\frac{d}{dt}\text{Ad}_{g}([Y, Z]) = [X, [Y, Z]]$$

右式为：

$$\begin{aligned} \frac{d}{dt}[\text{Ad}_g(Y), \text{Ad}_g(Z)] &= [\frac{d}{dt}\text{Ad}_g(Y), \text{Ad}_g(Z)] + [\text{Ad}_g(Y), \frac{d}{dt} \text{Ad}_g(Z)] \\ &= [[X, Y], Z] + [Y, [X, Z]] \end{aligned}$$

即

$$\begin{aligned} & [X, [Y, Z]] = [[X, Y], Z] + [Y, [X, Z]] \\ \Longleftrightarrow &[X, [Y, Z]] + [Y, [Z, X]] + [Z, [X, Y]] = 0 \end{aligned}$$

上式被称为雅可比恒等式。

总结一下，李括号的双线性、反对称性、雅可比恒等式，都是群、光滑流形的属性在微观尺度下的投影罢了。

ROPE

ROPE 旋转位置编码是现在大部分大模型位置编码的选型。它用一种绝对编码的方式为隐状态赋予了相对位置编码的效果。我们以一个隐状态维度 $d=2$ 为例，对于一个在位置 $m$ 的 token，其 $q, k$ 向量会被旋转 $m\theta$ 角度：

$$\begin{pmatrix} q_0' \\ q_1' \end{pmatrix} = \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \end{pmatrix}$$

这其实就是复数乘法 $q’ = q \cdot e^{i m\theta}$。由此，当我们计算 Attention Score 时，有：

$$\begin{aligned} q^T_mk_n &= (\mathcal{R}_mq)^T(\mathcal{R}_nk) = q^T\mathcal{R}_m^T\mathcal{R}_nk \\ &= q^T\mathcal{R}_{-m}\mathcal{R}_nk \\ &= q^T\mathcal{R}_{n-m}k \end{aligned}$$

也就是说，我们施加的绝对位置编码，在 QK 交互时的影响却仅限于相对位置关系，自动获得了一种平移不变性。

推广到 $d$ 维，RoPE 选择的是一种 两维 为一组的构造方式，形成了了一个巨大的块对角旋转矩阵 $\Lambda_m$：

$$\Lambda_m = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 \\ \sin m\theta_1 & \cos m\theta_1 & 0 & 0 \\ 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 \\ 0 & 0 & \sin m\theta_2 & \cos m\theta_2 & \dots \end{pmatrix}$$

这里就很值得说道说道了，这其实是一个非常强的归纳偏置，主要包含两个部分：一个是将旋转限定在了李群 $SO(N)$ 的最大环面上；另一个是限制了每个子群的旋转角度。

ROPE 有很多变体，接下来我会用一个统一的数学视角去解释。To be continue