Token 即群元素：论矩阵李群上的李代数注意力

arXiv: 2606.20547v1

论文信息

标题: The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups

作者: Przemyslaw Musialski

发布日期: 2026-06-18

arXiv ID: 2606.20547v1

PDF 链接: 下载 PDF

引言：当令牌成为变换本身

在大多数处理空间数据的等变模型中，令牌（token）通常是一个特征向量，例如三维点坐标或分子中原子的嵌入。对称群（如旋转、平移群）通过外部表示作用于这些向量，等变性必须借助复杂的数学工具来保证：不可约表示、克莱布施-戈尔丹张量积、可操控核或者几何代数中的夹层积。此类路线虽然有效，但存在结构性限制：不可约表示方法局限于紧致群，无法处理包含缩放和剪切的非紧致仿射群；而依赖满射指数映射的方法同样无法覆盖仿射群。这些共同的瓶颈源于一个根本设定：令牌始终是向量，群作用在外部施加。

《The Token Is a Group Element: On Lie-Algebra Attention over Matrix Lie Groups》一文做出了一次根本性转向：令牌本身就是一个矩阵李群元素。具体而言，一个令牌 gig_i 就是群 GG 中的一个裸变换，不再附载任何特征向量，也不再有外部表示 ρ(g)\rho(g) 作用其上。这一转向使得相对几何量 wij=log⁡(gi−1gj)w_{ij} = \log(g_i^{-1}g_j) 自然成为不变量，而注意力分数也因此拥有了一个封闭形式的代数范数，不再需要学习的核函数。作者将此构造命名为李代数注意力（Lie-Algebra Attention）。本文将对这一核心思想、技术实现、实验结果及应用前景进行深入解析。

核心思想：令牌即群元素

传统等变模型中的令牌 v∈Vv \in V ，群通过表示 ρ(g):V→V\rho(g): V \to V 施加作用。李代数注意力则完全抛弃 VV 和 ρ\rho ，令令牌直接成为群元素 gi∈Gg_i \in G 。例如，在平面仿射群 Aff(2)\mathrm{Aff}(2) 中，一个令牌就是一个带框架的点：包含两个基向量和一个原点，相当于一个完整的仿射坐标系。群运算直接就是矩阵乘法，令牌的相对关系由群运算的商 gi−1gjg_i^{-1}g_j 给出，这一相对位姿（relative pose）是典型的内在不变量。

自然不变量与闭式注意力分数

给定两个令牌 gi,gjg_i, g_j ，计算它们的相对位姿 gi−1gjg_i^{-1}g_j 并将其映射到李代数 g\mathfrak{g} 上，得到对数映射 wij=log⁡(gi−1gj)w_{ij} = \log(g_i^{-1}g_j) 。由于对角群作用下的抵消恒等式 (agi)−1(agj)=gi−1gj(a g_i)^{-1}(a g_j) = g_i^{-1}g_j ，wijw_{ij} 自动是群不变量的。这一性质无需额外设计，纯粹由群结构保证。

注意力分数直接定义为该不变量 wijw_{ij} 的负平方代数范数，除以温度 τ\tau：

sij=−∥wij∥λ2τs_{ij} = -\frac{\| w_{ij} \|_{\lambda}^2}{\tau}

其中 ∥⋅∥λ\| \cdot \|_{\lambda} 是李代数上的加权弗罗贝尼乌斯内积。权重 λ\lambda 按代数在正交群作用下的不可约块分解设置，例如 Aff(2)\mathrm{Aff}(2) 的代数分解为平移、旋转、各向同性缩放和各向异性剪切四个块，每个块有一个独立权重。这使得注意力头可以专注于不同几何属性：比如一个头放大位置差异，另一个头强调方向相似性。该分数是一个封闭形式的邻近核，不需要不可约表示、球谐函数或学习的核网络。

等变输出与余循环自动满足

模型输出一个代数元素 δi∈g\delta_i \in \mathfrak{g} ，表示对输入令牌的矫正。输出令牌为 giexp⁡(δi)g_i \exp(\delta_i) 。由于 δi\delta_i 由不变特征计算而来，因此本身是对角群作用不变的，最终 giexp⁡(δi)g_i \exp(\delta_i) 的群作用为 agiexp⁡(δi)=(agi)exp⁡(δi)a g_i \exp(\delta_i) = (a g_i) \exp(\delta_i) ，自动满足等变性。这种等变性是结构性的，不需要训练信号强加。

此外，所有输出令牌之间的相对位姿自动满足余循环条件（cocycle condition），即 gijgjk=gikg_{ij} g_{jk} = g_{ik} 。因为令牌就是群元素，输出只是左乘一个不变偏移，内部一致性保证。

技术框架：六种群实例化

该构造适用于任何具有忠实有限维矩阵表示的李群。论文给出了六个群的显式封闭形式：SO(2)\mathrm{SO}(2) 、SE(2)\mathrm{SE}(2) 、SO(3)\mathrm{SO}(3) 、SE(3)\mathrm{SE}(3) 、Aff(2)\mathrm{Aff}(2) 和 Aff(3)\mathrm{Aff}(3) 。其中仿射群是核心亮点：它们非紧致、非阿贝尔，包含缩放和剪切，传统不可约表示方法不存在非平凡有限维酉表示，而基于满射指数映射的方法也因仿射群的指数映射非满射而失效。

以 Aff(2)\mathrm{Aff}(2) 为例，其李代数 aff(2)\mathfrak{aff}(2) 的块分解为：

aff(2)=R平移2⊕so(2)旋转⊕R各向同性缩放⊕Sym0(2)各向异性+剪切\mathfrak{aff}(2) = \mathbb{R}^2_{\text{平移}} \oplus \mathfrak{so}(2)_{\text{旋转}} \oplus \mathbb{R}_{\text{各向同性缩放}} \oplus \mathrm{Sym}_0(2)_{\text{各向异性+剪切}}

每个块对应不同几何语义。注意力分数因此计算为：

sij=−λt∥uij∥2+λθδθij2+λsηij2+λq∥qij∥2τs_{ij} = -\frac{ \lambda_t \|u_{ij}\|^2 + \lambda_\theta \delta\theta_{ij}^2 + \lambda_s \eta_{ij}^2 + \lambda_q \|q_{ij}\|^2 }{\tau}

通过正交基的归一化，各块权重可学习，从而在不同几何敏感度上进行调控。这种分解在 Aff(3)\mathrm{Aff}(3) 中直接推广到 12 维，包含 5 维的各向异性剪切块。

实验验证

论文使用序列补全任务测试了三个代表组：SE(2)\mathrm{SE}(2) （两块）、SO(3)\mathrm{SO}(3) （一块）和 Aff(2)\mathrm{Aff}(2) （四块）。每种情况下，给定一个由固定步长生成的序列，去掉一个中间令牌并将剩余令牌打乱，模型需要根据无序集合重建缺失的位姿。对比模型包括：

  • G：本文的代数范数得分；
  • C：使用相同不变量 wijw_{ij} 但得分用一个可学习的 MLP 替代；
  • A：标准向量令牌，输入为绝对坐标特征，得分为缩放点积。

在 SE(2)\mathrm{SE}(2) 上，G 以仅有 36 个得分参数达到位姿误差 0.003，C 拥有 1932 个参数（约 54 倍）却略差（0.005）；两者等变误差均在浮点精度水平（10−1010^{-10} 以下），而向量令牌 A 的等变误差高出 11 个数量级，位姿误差高出 23 倍。在 SO(3)\mathrm{SO}(3) 上，G 用 24 个参数与 C 打平，等变误差落在 10−1410^{-14}，而 A 的等变误差达 7.4×10−27.4 \times 10^{-2}。在 Aff(2)\mathrm{Aff}(2) 上，G 与 C 表现相当（0.007 位姿误差），但 G 用 60 个参数 vs C 的 3084 个（51 倍）。A 模型完全崩塌，等变误差达到 1.29，位姿误差 117 倍于 G。

这些结果表明：在简单但严格测试不变性的任务上，封闭形式的代数范数得分与容量大得多的学习核效果相当甚至更优，而向量令牌本体论完全无法维持等变性。仿射群的实验是首次在注意力机制中处理非紧致非阿贝尔群的展示。

实践应用建议与未来方向

李代数注意力的最大优势在于它提供了对旋转、平移、缩放和剪切等变换的统一处理，且无需复杂等变机制。这为实际空间推理任务开辟了新路径：

  • 机器人学与具身智能：机器人的末端执行器位姿、物体包围盒经常用小规模仿射变换描述。使用 SE(3)\mathrm{SE}(3) 或 Aff(3)\mathrm{Aff}(3) 令牌可以自然处理姿态估计、轨迹预测等任务。
  • 蛋白质结构预测：AlphaFold 类方法中使用的帧可替换为严格的李群令牌，统一旋转和平移的注意力得分，并可能扩展到局部缩放柔性。
  • 计算机视觉：物体检测中的锚框或点云配准可以通过仿射令牌捕获不同尺度和剪切形变，提升对复杂变形物体的鲁棒性。
  • 分子动力学：分子的旋转与平移对称性天然符合 SE(3)\mathrm{SE}(3) ，使用李代数注意力可大幅减少参数，并保证物理等变性。

然而，该方法也有局限：分数依赖于对数映射，仅在主对数图表上有效，当相对位姿接近图表边界（如旋转 π\pi ）时数值不稳定。对此可通过图表切换或直接使用全局形式 giΔig_i \Delta_i 解决，但会失去封闭范数得分。另外，得分是对称的，对于需要非对称关系的任务需要补充方向信息。

未来工作可探索更复杂的序列或图任务，利用不同块权重学习更丰富的几何注意力模式；在 Aff(3)\mathrm{Aff}(3) 上进行完整验证；并尝试将此类令牌应用到大规模真实数据中，检验其参数效率。

总结

本文重新定义了注意力令牌的本质：令牌不再是向量，而是李群上的一个变换。由此诞生了李代数注意力，它仅凭群结构就获得了自然不变量 wijw_{ij} 和闭式邻近核得分，等变性和余循环约束自动成立。构造覆盖了从紧致旋转群到非紧致仿射群的六大群，首次在注意力中实现了对包含缩放和剪切的全帧处理。实验表明，封闭形式的代数范数得分在参数效率上显著优于学习的核，同时彻底碾压了破坏等变性的向量令牌基线。这一 “令牌即群元素” 的哲学有望简化等变建模，并推动空间推理任务向更丰富的对称性扩展。