词元是一个群元素：矩阵李群上的李代数注意力机制

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

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

3 分钟速览

  • 研究问题：打破 “token = 特征向量” 的默认假设，把 transformer 的 token 直接定义为矩阵李群的元素（一个裸变换），从而让等变性从天生的群结构里自动浮现，不再需要借助表示论去强行施加。

  • 核心方法：用相对姿态的对数 wij=log⁡(gi−1gj)w_{ij} = \log(g_i^{-1}g_j) 作为 token 之间的内禀不变量，注意力分数采用闭式代数范数 sij=−∥wij∥λ2/τs_{ij} = -\|w_{ij}\|_\lambda^2/\tau，不学习任何核函数；输出通过 giexp⁡(δi)g_i \exp(\delta_i) 的方式保持等变。

  • 关键结果：在 SE(2) 序列补全任务中，闭式分数模型 G 用 36 个分数参数（比学习型 MLP 核 C 少 54 倍）取得了更低的姿态误差 0.003（C 为 0.005），且等变性误差仅 4.0×10−104.0\times10^{-10}，而向量 token 基线 A 的等变性误差高达 1.3×10−11.3\times10^{-1}（表 2）。

  • 主要局限：只能在 principal log chart 内工作，临近图边界时数值不稳定；注意力分数对称，丢失方向信息（依赖值通路弥补）；实验仅停留在序列补全玩具任务，尚未在真实三维场景中验证；块权重假设各向同性，未充分消融。

  • 适合读者：从事几何深度学习、等变神经网络、李群方法、以及机器人/分子/蛋白质结构建模的研究者，尤其是需要突破 SE(3) 限制、进入仿射群等非紧非交换域的实践者。

论文背景和研究动机

在几乎所有等变模型中，token 被默认为一个向量 vv，对称群 GG 通过一个外部表示 ρ(g)\rho(g) 作用在其上。为了迫使模型在群作用下保持行为一致，需要一套繁重的 “表示论工具链”：不可约表示分解、Clebsch–Gordan 张量积、球谐函数或可导向核等。这套范式在技术上高度成熟，却也因此将等变性的适用范围锁定在紧致群（例如 SO(3)）附近：非紧致非交换群没有非平凡有限维酉表示，传统路子走不通，仿射群 Aff(n)（包含尺度与剪切变换）长期被排除在外。

论文提出的 “李代数注意力”（Lie‑Algebra Attention）从根基上改变了这一状况。作者不再把 token 看作被群作用的对象，而是让 token 本身就是群元素——一个裸的矩阵李群元素 gig_i，不附带任何特征载荷，也没有外部表示 ρ(g)\rho(g) 去携带它。一旦 token 落在群上，许多结构性质便随之而来：相对几何由 gi−1gjg_i^{-1}g_j 内在地定义；逐对不变量 wij=log⁡(gi−1gj)w_{ij}=\log(g_i^{-1}g_j) 不再是人为设计的，而是直接从群运算中读出；对角 GG–作用下的等变性变成一个同义反复的结论，而不是需要学习的性质；所有成对相对姿态的相容性（cocycle 条件）自动成立。正因为这一切都是群公理的直接推论，作者断言该构造免去了前述所有表示论机械。

该工作的根本动机是，不是去寻找更好的核函数，而是改变 “token 是什么” 这个前提。由此产生的李代数注意力在闭式分数的加持下，不仅统一了旋转、平移等常见对称性，还首次将注意力机制延伸到仿射全标架（包含各向同性缩放和各向异性剪切）的非紧非交换领域。

核心方法和技术细节

从向量 token 到群元素 token

构设的核心极为简单：一个 token 就是群 GG 的一个元素 gig_i。以仿射群 Aff(2)\mathrm{Aff}(2) 为例，它由 2×22\times2 可逆线性部分 AA 和二维平移 tt 组成，gi=(Aiti01)g_i = \begin{pmatrix} A_i & t_i \\ 0 & 1 \end{pmatrix}。群运算（乘法、求逆）是标准矩阵运算，实现了 token 的合成与修正。

不变量 wijw_{ij} 与闭式注意力分数

对于任意两个 token，它们的相对姿态 gi−1gjg_i^{-1}g_j 在群对角作用 gi↦agig_i\mapsto a g_i 下不变，因为 (agi)−1(agj)=gi−1gj(a g_i)^{-1}(a g_j) = g_i^{-1}g_j。在群的主 log chart U\mathcal{U} 上，取对数得到李代数元素

wij=log⁡(gi−1gj)∈g.w_{ij} = \log(g_i^{-1}g_j) \in \mathfrak{g}.

wijw_{ij} 是内禀不变量，不需要专门设计。注意力分数采用 wijw_{ij} 的加权 Frobenius 范数的负值除以温度，即

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

其中的加权范数尊重李代数的几何块结构。以 Aff(2)\mathrm{Aff}(2) 为例，其李代数按 O(2)\mathrm{O}(2) 作用可分解为四个不可约块：平移（R2\mathbb{R}^2）、旋转（R\mathbb{R}）、各向同性缩放（R\mathbb{R}）和各向异性剪切（Sym0(2)≅R2\mathrm{Sym}_0(2)\cong \mathbb{R}^2）。因此加权范数具有形式

∥wij∥λ2=λt(δtx2+δty2)+λθδθ2+λss2+λq(q12+q22).\|w_{ij}\|_{\lambda}^2 = \lambda_t(\delta t_x^2+\delta t_y^2) + \lambda_\theta \delta\theta^2 + \lambda_s s^2 + \lambda_q (q_1^2+q_2^2).

每一块拥有独立的可学习权重 λ\lambda，既可让不同头关注不同几何维度，又保持 O(n)\mathrm{O}(n) 各向同性。整个分数参数量极少：例如 SE(2)\mathrm{SE}(2) 每头只有 2 个块权重 + 1 个温度，4 头 3 层总共 36 个参数。

等变性输出与自动相容性

等变性在该构设中是同义反复的。任何满足等变性的映射必然具有 gi↦g^i=giΔig_i \mapsto \hat{g}_i = g_i \Delta_i 的形式，其中 Δi\Delta_i 是对角作用下的不变量。在局部 chart 上，可写为 Δi=exp⁡(δi)\Delta_i = \exp(\delta_i)，δi∈g\delta_i\in\mathfrak{g}，而 δi\delta_i 由模型从不变量特征中预测。最终输出

g^i=giexp⁡(δi)\hat{g}_i = g_i \exp(\delta_i)

自然满足等变性 g^i(a⋅)=a⋅g^i(⋅)\hat{g}_i(a\cdot) = a\cdot \hat{g}_i(\cdot)。此外，任意两个输出之间的相对姿态满足 cocycle 条件 g^ijg^jk=g^ik\hat{g}_{ij}\hat{g}_{jk} = \hat{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) 六个矩阵李群的闭式实现。对于非紧致非交换的仿射群，尽管不存在 Ad‑不变正定度量，但 Frobenius 范数（及其块加权形式）仍然有效，使得该构造绕开了传统表示论与满射指数映射的限制，首次将注意力覆盖到包含缩放和剪切的仿射全标架。

创新点和贡献

论文的核心贡献并非设计了一个新的注意力核函数，而是重新定义了 token 的本体论。具体贡献可概括为三点：

  1. 裸群元素 token token 仅是一个群元素，没有任何附属特征向量，也无外部表示作用。等变性、逐对不变量 wijw_{ij} 以及 cocycle 相容性均来自群结构的一行证明。与此形成对照的是，一旦给 token 加上特征载荷（基线 A），等变性误差立即放大五到十二个数量级（表 2–4）。

  2. 闭式规范分数 分数纯粹是相对姿态的代数范数：sij=−∥wij∥λ2/τs_{ij} = -\|w_{ij}\|_\lambda^2/\tau。它不使用不可约表示、球谐函数、Clebsch–Gordan 积或可学习核函数。论文在 SE(2)、SO(3)、Aff(2) 三个任务上比较了闭式分数 G 与学习型 MLP 核 C。C 使用相同不变量 wijw_{ij} 并具有 50~80 倍多的分数参数，但性能无法超越 G，甚至在 SE(2) 上明显落后（G 误差 0.003，C 误差 0.005，表 2）。

  3. 进入仿射全标架 该构造自然地延伸至 Aff(2)\mathrm{Aff}(2) 和 Aff(3)\mathrm{Aff}(3)，包含尺度与剪切。这些非紧致非交换群此前无法被基于不可约表示或满射指数映射的注意力方法处理。论文的实验在 Aff(2)\mathrm{Aff}(2) 序列补全中展示了闭式分数的有效性，并给出了 Aff(3)\mathrm{Aff}(3) 的完整闭式，为后续三维仿射 token 应用铺路。

实验结果分析

实验在三个群上执行序列补全任务：每个序列有 8 个等步群元素，随机移除一个内部元素并打乱剩余 7 个的顺序，模型需从无序集合中重建缺失的姿态。对比三种模型：

  • G（闭式范数分数）
  • C（在相同 wijw_{ij} 上用 MLP 学习核）
  • A（传统向量 token + 缩放点积注意力，作为等变性破坏对照）

所有模型均使用相同的 Transformer 架构（3 层、4 头、隐藏维度 32），仅 token 类型和评分机制不同。

SE(2) 结果（表 2）：G 以 36 个分数参数实现姿态误差 0.0030.003，C 用 1932 个参数（54×）得到 0.0050.005，G 的相对改善约为 34%。G 和 C 的等变性误差均在浮点精度附近，而 A 的等变性误差高达 0.130.13，姿态误差为 0.0690.069，比 G 差一个量级以上。

SO(3) 结果（表 3）：G（24 个分数参数）和 C（1932 个参数，80×）的误差均在 10−410^{-4} 量级，统计上难以区分；等变性误差 10−1410^{-14} 级别，几乎是浮点运算的底线。A 的等变性误差 7.4×10−27.4\times10^{-2}，姿态误差 0.0670.067，是 G 的约 370 倍。

Aff(2) 结果（表 4）：G（60 个参数）和 C（3084 个参数，51×）的姿态误差均为 0.0070.007，在种子噪声范围内持平。A 的姿态误差达到 0.790.79（是 G 的约 117 倍），等变性误差 1.291.29，五到九个数量级高于 G 和 C。

整体来看，闭式范数分数在三个块数级别（单块、两块、四块）上以极少的参数匹配或优于学习核，并始终保持结构等变性。这些实验虽然简单，却清晰展示了 “token 落群” 带来的威力。

局限与待解决问题

尽管该构设的数学优美性令人印象深刻，但它目前的形态仍存在若干限制：

chart 依赖。整个构造依赖于主对数图 U\mathcal{U} 的存在，当相对姿态靠近图边界时（例如旋转角接近 π\pi，或仿射部分有负实轴特征值），闭式分数不再可靠甚至未定义。论文在实验中通过约束步长回避了这一问题，但如果将该机制部署到无先验约束的通用场景中，必须配套 chart 切换或回退到全局形式（此时闭式分数失效）。

对称的分数缺乏方向感。注意力分数 sijs_{ij} 对 i,ji,j 对称（因为 ∥wij∥λ=∥wji∥λ\|w_{ij}\|_\lambda = \|w_{ji}\|_\lambda），不对称的方向信息只由值通路 WV[hj;wij]W_V[h_j; w_{ij}] 携带。对于需要区分方向的任务（如自回归预测），可能需要额外加入非对称的分数项。

块各向同性未充分验证。当前加权范数在每个块内使用各向同性 Frobenius 度量，这等效于只用一个标量权重连接块内各分量。由于实验中的不变量呈共线，各块的权重实际上退化为单一标量，无法评估各向异性分数的真实优势。未来需要在数据的各个块方向充分分散的任务上消融块权重的作用。

玩具任务规模。实验仅涉及 8 个 token 的序列补全，远未达到实际应用的数据规模和复杂性。论文虽然给出了 SE(3)\mathrm{SE}(3) 和 Aff(3)\mathrm{Aff}(3) 的闭式公式，但尚未在分子、蛋白质或机器人等典型 3D 场景中验证。这些方向的实证评估被留给后续工作。

参数敏感性。块权重与温度均通过 softplus 参数化并初始化为约 0.69, 论文未报告对超参数的敏感性分析；在更复杂任务上是否需要手动调整近端权重仍有待观察。

尽管有这些限制，李代数注意力依然是向 “token 即变换” 这一直觉迈出的重要一步。它用极其简洁的群论武器，替换了庞大而受限的表示论军械库，并在仿射群上打开了新的可能性。