跳转至

How Do Transformers Learn to Represent Symmetries?

会议: NeurIPS 2026
arXiv: 2610.10305
领域: 学习理论
关键词: 几何深度学习、Transformer归纳偏置、对称性学习、有限数据增强、等变性表征

一句话总结

本文系统探究了标准原生 Transformer 在有限数据增强下学习对称性不变性与等变性的机制,揭示了对称群可学习性的清晰层级结构,并剖析了注意力层、残差流与归一化协同实现旋转、平移和缩放不变性的可解释内部机制。

研究背景与动机

在分子建模、材料发现、计算化学与三维视觉等几何深度学习领域,将物理对称性(如欧几里得群 \(\mathrm{SE}(3)\)、正交群 \(\mathrm{O}(3)\) 等)融入神经网络至关重要。长期以来,领域内主要依赖硬编码等变架构(如 GCNN、EMLP、Equiformer 等)来强制满足几何约束。然而,这类硬等变模型往往伴随着极高的实现复杂度、计算开销大以及受限的假设空间等瓶颈,特别是在大规模真实世界数据中,系统常常存在测量噪声、离散化效应或轻微破缺的对称性。近年来,以 AlphaFold3、TransIP、MCFs 为代表的前沿系统开始转向软对称策略——使用标准 Transformer 架构配合有限数据增强(Finite Data Augmentation)来诱导近似对称性。

尽管基于有限增强的 Transformer 在实证上取得了巨大成功,但现有的理论分析大多假设无限次增强的理想渐近状态,这与实际计算预算受限的真实场景严重脱节。在有限样本视角下,一个核心基础问题始终未被回答:标准 Transformer 是否能无差别地学习所有对称变换群?其自身架构是否存在对特定几何变换的内生归纳偏置(Inductive Bias)?

本文正是针对该核心空白展开系统探究。研究切入点是在可控的点云几何任务与真实基准下,对比 Transformer 面对不同连续李群变换时的样本效率、分布外(OOD)变换幅度外推表现,以及模型内部表征的演变。核心 idea:Transformer 对保持几何角度的对称群具有天然的强归纳偏置,通过揭示对称性可学习性从「非保角群」到「复合保角群」再到「基础保角群」的三层递增层级,本文阐明了自注意力矩阵内积定界、残差流基准点相对坐标减法以及层归一化齐次性是 Transformer 自发诱导不变性与等变性的核心物理机制。

方法详解

整体框架

本文构建了一套基于李群与李代数重参数化的有限数据增强评估体系,并在拓扑合成数据集(3D Topology)、经典三维几何形状分类(ModelNet10)、三维俄罗斯方块(3D Tetris)以及分子属性预测(QM9)等多维基准上进行诊断。评估涵盖两大类别共九种对称群:非保角群(仿射群 \(\mathrm{Aff}(3)\)、特殊线性群 \(\mathrm{SL}(3)\)、一般线性群 \(\mathrm{GL}(3)\))与保角对称群(平移群 \(\mathrm{T}(3)\)、正交群 \(\mathrm{O}(3)\)、特殊正交旋转群 \(\mathrm{SO}(3)\)、缩放群 \(\mathrm{Scale}(3)\)、刚体变换群 \(\mathrm{SE}(3)\) 与相似变换群 \(\mathrm{Sim}(3)\))。

模型以极简双层单头原生 Transformer(去除位置编码与掩码以消除额外结构干扰)为骨干,系统观测训练过程中增强视角数量从 1 扩展到 32 时的下游分类准确率以及基于类概率分布的对称不变性 KL 散度误差。随后,针对可解释性内部表征展开剖析,追踪自注意力矩阵稳定性、残差向量余弦相似度与范数比值,最终将不变性机制进一步推广至点云法向量与部分质心预测的等变任务。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入点云 X ∈ ℝⁿˣᵈ 与有限李代数增强采样"] --> B["对称群可学习性分层诊断<br/>分类精度与不变性 KL 误差综合评估"]
    B --> C["旋转与反射不变机制<br/>注意力分数双线性内积差分定界 (Lemma 1)"]
    B --> D["平移不变机制<br/>基准点注意力聚焦与残差流逆向相减 (Lemma 2)"]
    B --> E["缩放不变机制<br/>LayerNorm 正齐次性与残差流非齐次抑制 (Lemma 3)"]
    C --> F["从不变性向等变性表征推广<br/>ShapeNet 法向量与部位质心预测等变重构"]
    D --> F
    E --> F

关键设计

1. 旋转与反射不变性:标量内积与行内注意力差分定界 根据古典不变量理论(First Fundamental Theorem of Invariant Theory),正交变换 \(\mathrm{O}(d)\) 的标量多项式可完全由点云成对内积所构成的 Gram 矩阵 \(\bm{G}(\bm{X})=\bm{X}\bm{X}^\top\) 参数化。在 Transformer 的自注意力计算中,预 Softmax 注意力分数天然表现为双线性配对: $\(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d}} = \frac{\bm{X}\bm{W}_Q\bm{W}_K^\top\bm{X}^\top}{\sqrt{d}} = \frac{\bm{X}\bm{B}\bm{X}^\top}{\sqrt{d}}, \quad \bm{B}=\bm{W}_Q\bm{W}_K^\top\)$ 当参数矩阵满足完全各向同性 \(\bm{B} = \alpha \bm{I}_d\) 时,注意力分数直接正比于 Gram 矩阵,对任意 \(\bm{R}\in\mathrm{O}(d)\) 均严格不变。本文引理 1(Lemma 1)进一步证明,由于 Softmax 具有沿行平移不变性,模型在有限增强轨道上并不需要严格的全局各向同性,而只需行内分数差分满足不变性约束: $\(\bm{X}(\bm{R}\bm{B}\bm{R}^\top - \bm{B})\bm{X}^\top \bm{P}_n = \bm{0}\)$ 其中 \(\bm{P}_n = \bm{I}_n - \frac{1}{n}\bm{1}_n\bm{1}_n^\top\) 为中心化矩阵。实证分析证实,训练后的 Transformer 在不同旋转角度下的注意力图展现出极其高度的数值稳定性。

2. 平移不变性:基准点注意力锚定与残差流逆向相减机制 为了消除全局刚体平移 \(\bm{X} \mapsto \bm{X} + \bm{1}_n\bm{t}^\top\) 的影响,网络必须构建相对坐标系。本文引理 2(Lemma 2)表明,若注意力分配算子 \(\bm{\Lambda} = \Lambda(\bm{X})\) 满足平移不变且行随机,则变换 \(f(\bm{X}) = g(\bm{X} - \bm{\Lambda}^\top \bm{X})\) 对任意函数 \(g\) 均严格平移不变。实证激活分析揭示了一个引人注目的计算现象: - 注意力图中的权重高度聚集在单个特定点(Landmark 锚点); - 自注意力输出 \(\mathrm{SA}(\bm{X})\) 与残差流输入 \(\bm{X}\) 之间的点对点余弦相似度在整个测试集上严格为负(集中在 -0.8 到 -1.0 之间)。 这证实自注意力层实际在执行负向加权,通过残差连接 \(\bm{X} + \mathrm{SA}(\mathrm{LN}(\bm{X}))\) 精确完成了从绝对坐标中扣除基准点坐标的操作。消融实验进一步证实,若移除残差连接,模型将彻底丧失平移不变性的学习能力。

3. 缩放不变性:正齐次性保持与残差流非齐次贡献抑制 在全局均匀缩放 \(\bm{X} \mapsto c\bm{X}\)(\(c > 0\))变换下,注意力矩阵若保持对尺度缩放稳定,则纯注意力层 \(\mathrm{SA}(\bm{X}) = \mathrm{A}(\bm{X})\bm{X}\bm{W}_V\) 天然具备一阶正齐次性 \(\mathrm{SA}(c\bm{X}) = c\,\mathrm{SA}(\bm{X})\)(Lemma 3)。在分类任务中,只要 logits 输出被缩放正系数 \(c\),由于单调性不变,\(\mathrm{argmax}\) 判决结果便严格不变。 然而,在含有残差流和偏置项的完整架构中,输入投影偏置与残差流动破坏了齐次性。特征探针分析揭示了模型的自我补偿策略:在第二层模块中,自注意力输出向量的 \(L_2\) 范数显著大于残差流特征向量范数(比值高达 5-10 倍),模型通过主动抑制残差流的相对幅度,将残差流带来的尺度不一致性压制到接近背景噪声的微弱水平。

4. 从不变性到等变性:不变注意力作为等变计算的组合基元 等变任务要求输出伴随输入变换发生等变映射。基于 Villar 等人的物理等变张量理论,任意 \(\mathrm{O}(d)\) 等变向量函数均可表示为标量不变系数对输入向量的加权线性组合。由于自注意力输出本质上即为加权和:\(\mathrm{SA}(\bm{X})_i = \sum_j \mathrm{A}(\bm{X})_{ij} (\bm{X}\bm{W}_V)_j\),若注意力图 \(\mathrm{A}(\bm{X})\) 保持不变,同时值投影 \(\bm{W}_V\) 在数据流形支撑集上近似等变(通过 Gram 矩阵 \(\bm{V}\bm{V}^\top\) 的高度旋转不变性得到验证),自注意力层便自发涌现出近似的几何等变性。

损失函数 / 训练策略

模型在分类任务中仅采用标准交叉熵损失进行端到端优化,在等变回归任务中采用均方误差(MSE)损失: $\(\mathcal{L}_{\text{task}} = \frac{1}{NK} \sum_{i=1}^N \sum_{j=1}^K \ell(\Psi(g_j x_i), y_i)\)$ 其中对称性不变性 KL 散度与等变误差仅作为诊断与监测指标,完全不参与反向传播。优化器统一使用 Adam(学习率 0.001),隐藏层维度为 64。在分布外(OOD)评估中,训练时仅使用小变换区间(如旋转 \([0^\circ, 90^\circ)\)、平移 \([0, 1.5)\)、缩放 \([1.0, 1.5)\)),测试时则评估不相交的大变换区间(如旋转 \([90^\circ, 180^\circ)\)、平移 \([1.5, 3.0)\)、缩放 \([1.5, 2.0)\)),以严格考察几何泛化外推能力。

实验关键数据

主实验

在 ModelNet10 数据集上,不同对称群在相同数据增强预算(16 次增强/样本)及测试集上的分类准确率与不变性 KL 散度表现如下表所示(对比保角变换与非保角变换):

对称群类别 对称群 变换性质 测试准确率 (%) ↑ 不变性误差 (KL 散度) ↓ 学习难度评估
基础保角群 \(\mathrm{Scale}(3)\) 尺度变换 78.4 ± 1.2 0.04 ± 0.01 极易,小预算即收敛
基础保角群 \(\mathrm{T}(3)\) 平移变换 76.8 ± 1.5 0.06 ± 0.02 极易,基准点减法驱动
基础保角群 \(\mathrm{SO}(3)\) 三维三轴旋转 72.6 ± 1.8 0.12 ± 0.03 易,需适度增强预算
基础保角群 \(\mathrm{O}(3)\) 旋转与镜像反射 71.9 ± 1.7 0.14 ± 0.03 易,注意力差分不变
复合保角群 \(\mathrm{SE}(3)\) 刚体变换 (\(\mathrm{T}\rtimes\mathrm{SO}\)) 64.2 ± 2.1 0.28 ± 0.05 中等,多自由度优化冲突
复合保角群 \(\mathrm{Sim}(3)\) 相似变换 61.5 ± 2.4 0.35 ± 0.06 中等,空间维度增加
非保角群 \(\mathrm{SL}(3)\) 保体积剪切变换 42.1 ± 3.2 0.88 ± 0.12 困难,准确率显著受损
非保角群 \(\mathrm{GL}(3)\) 可逆线性变换 38.5 ± 3.6 1.15 ± 0.15 极难,拓扑几何失真
非保角群 \(\mathrm{Aff}(3)\) 仿射群 (\(\mathrm{T}\rtimes\mathrm{GL}\)) 35.2 ± 4.1 1.34 ± 0.18 极难,有限增强难以收敛

消融实验

在 3D Topology 数据集上,针对平移不变性与缩放不变性关键架构模块(残差连接与 LayerNorm)的消融对比(均在 16 次增强训练下):

模型配置 目标任务/对称性 测试分类准确率 (%) ↑ 不变性误差 (KL 散度) ↓ 结论与机制验证
完整残差模型 (\(\bm{X} + \mathrm{SA}(\mathrm{LN}(\bm{X}))\)) 平移 \(\mathrm{T}(3)\) 拓扑分类 98.4 ± 0.6 0.03 ± 0.01 残差相减实现基准点相对坐标化
移除残差流 (\(\mathrm{SA}(\mathrm{LN}(\bm{X}))\)) 平移 \(\mathrm{T}(3)\) 拓扑分类 34.2 ± 4.8 1.28 ± 0.22 无法构造坐标减法,平移不变性崩溃
纯注意力归一化模型 (\(\mathrm{SA}(\mathrm{LN}(\bm{X}))\)) 缩放 \(\mathrm{Scale}(3)\) 拓扑分类 99.2 ± 0.4 0.01 ± 0.01 满足严格一阶正齐次性,分类天然不变
移除 LayerNorm / 含未归一化残差 缩放 \(\mathrm{Scale}(3)\) 拓扑分类 82.5 ± 2.1 0.19 ± 0.04 需第二层大幅放大注意力输出抑制残差
单头注意力 (1 Head) 刚体变换 \(\mathrm{SE}(3)\) 64.2 ± 2.1 0.28 ± 0.05 基线多自由度联合学习
多头注意力 (4 Heads) 刚体变换 \(\mathrm{SE}(3)\) 66.8 ± 1.9 0.25 ± 0.04 多头并未分别解耦平移与旋转

关键发现

  • 三层可学习性层级(Finding 1):实验呈现出明确的可学习性优劣顺序:基础保角群(\(\mathrm{Scale}, \mathrm{T}, \mathrm{SO}, \mathrm{O}\)) > 复合保角群(\(\mathrm{SE}, \mathrm{Sim}\)) > 非保角群(\(\mathrm{SL}, \mathrm{GL}, \mathrm{Aff}\))。非保角变换破坏了内积与几何角度,导致注意力机制的相似度度量失效。
  • 分布外外推表现坚挺(Finding 2):在 ModelNet10 与 QM9 上,仅见过小幅度变换的模型在大幅度 OOD 变换集上仍能保持高度不变性。与 PointNet(带显式 T-Net 对齐)和 DeepSets 对比,Transformer 的外推性能显著击败 DeepSets,且在平移不变性外推上甚至优于 PointNet。
  • 复合群学习瓶颈并非表达能力限制,而是高维采样覆盖度不足:\(\mathrm{SE}(3)\) 拥有 6 个自由度,在固定采样预算下,联合空间覆盖极为稀疏;此外,平移所需的残差相减与旋转所需的各向同性在共享权重通道内存在多任务约束竞争。增加注意力头数(2头、4头)并未带来质的跃升,各头并未发生「一头专攻旋转、一头专攻平移」的功能解耦。

亮点与洞察

  • 注意力图的双线性内积不变性定理:将经典的古典不变量理论与现代 Softmax 注意力结合,证明无需权重全局满足 \(\bm{B} = \alpha\bm{I}\),只要满足行内差分不变即可达成注意力图的完全旋转不变性。
  • 残差流充当相对坐标减法器:揭示了 Transformer 应对平移时并非依靠卷积式局部感受野,而是自发让注意力聚焦在单个锚点上,再利用残差连接完成类似 \(x_i - x_{\text{landmark}}\) 的基准点减法,直观优美。
  • 为软等变设计范式奠定理论注脚:AlphaFold3、MCFs 等前沿模型放弃硬等变转投 Transformer+增强,此前常被视作工程妥协;本文从几何归纳偏置角度指明,原生 Transformer 恰好对物理世界中最核心的保角基础群具有天然极高的学习亲和力。

局限与展望

  • 实验受限于轻量级模型与纯点云模态:为隔离机制干扰,本文采用了 2 层单头架构和无序坐标 token。在包含稠密多头、深层网络、复杂位置编码(如 RoPE、ALiBi)以及图像/网格/高分子图数据中,高阶交互对机制的影响仍需进一步验证。
  • 等变性分析仍处探索阶段:对等变性的探讨主要集中在首层自注意力与经验 Gram 矩阵分析,深层网络中各层特征在隐空间中如何以高维群表示实现全局等变流动尚缺乏完整闭环证明。
  • 缺少因果干预实验:目前主要依赖激活分析与模块消融;未来可通过模型编辑(如人为正交化 \(\bm{B}\) 矩阵或钳制注意力锚点)进行精确的因果干预研究。

相关工作与启发

  • vs. 硬等变架构(PointNet [51], SE(3)-Transformers [20], Equiformer [37]):硬等变通过球谐函数张量积或显式 T-Net 严格限制假设空间,但在大算力与噪声数据下拓展性差;本文证明原生 Transformer 仅靠有限增强即可高效内生近似等变性,计算更加轻量高效。
  • vs. 不变量理论与标量化网络(Villar et al., 2021 [52]):Villar 等人证明了标量内积与基准点减法是构造等变多项式的万能形式;本文首次实证并理论证明了 Transformer 的注意力层与残差连接在受训后恰恰自发演化出了这套标量化与基准点减法机制。
  • vs. 有限增强泛化理论(Tahmasebi & Weber, 2026 [14]):后者证明了达成近似对称性比严格对称性在样本复杂度上指数级容易;本文则从具体的神经网络架构切入,明确指出了 Transformer 架构在何种对称群上能够以最低预算兑现这种近似对称性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从微观机制与可解释性角度系统解构 Transformer 学习对称性的内部工作机理。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 个数据集、9 大对称群、OOD 外推测试、模块消融及多种位置编码/多头验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨,引理陈述与实证激活图谱高度呼应,叙事逻辑极其清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为科学计算与几何深度学习中从硬等变向数据增强 Transformer 的范式转移提供了坚实的理论支撑。