自旋加权球谐函数实现完整且可扩展的$\mathrm{E}(3)$等变网络

# 自旋加权球谐函数实现完整且可扩展的E(3)等变网络
来源：https://arxiv.org/html/2607.01408
Yuchao Lin¹* Andrii Kryvenko¹* Wendi Yu¹ Chuan Li² Jianwen Xie² Xiaofeng Qian³,⁴,⁵ Shuiwang Ji¹,³,⁶
¹ 德克萨斯A&M大学计算机科学与工程系 ² Lambda, Inc. ³ 德克萨斯A&M大学材料科学与工程系 ⁴ 德克萨斯A&M大学电气与计算机工程系 ⁵ 德克萨斯A&M大学物理与天文系 ⁶ 德克萨斯A&M大学机械工程系

###### 摘要

**摘要：**
E(3)等变网络在三维原子系统建模中前景广阔，但其可扩展性受限于Clebsch-Gordan张量积（CGTP）的O(L⁶)复杂度。最近提出的Gaunt张量积（GTP）降低了复杂度，但无法捕捉反对称路径，导致表达能力不完整。本文提出SpinGTP，通过从标量函数推广到自旋加权球谐函数（SWSH）来克服GTP的不完整性。利用SWSH的代数性质，SpinGTP恢复了缺失的反对称相互作用，同时保持了GTP的渐近效率。它还允许更富表达能力的等变基，自然地考虑了张量积的奇宇称分量。我们在多个基准上评估了SpinGTP，包括Tetris、3BPA、SPICE-MACE-OFF和OC20。结果表明，SpinGTP达到了与完整CGTP相当的精度。值得注意的是，通过显式捕捉反对称路径，SpinGTP在手性材料和非中心对称几何相关任务中表现出优越性能。这项工作为大规模三维原子系统模拟中的高阶等变性提供了一条完整、可扩展且数学严谨的路径。

关键词：人工智能科学，等变网络，自旋加权球谐函数，三维原子模拟

* 这些作者贡献相同

![[无标题图像]](https://arxiv.org/html/2607.01408v1/x1.png)GitHub仓库：https://github.com/divelab/SWSH-GNN

###### 目录

1. 1 引言 (https://arxiv.org/html/2607.01408#S1)
2. 2 相关工作 (https://arxiv.org/html/2607.01408#S2)
3. 3 自旋加权球谐函数的Gaunt张量积 (https://arxiv.org/html/2607.01408#S3)
   1. 3.1 Gaunt张量积效率与反对称性差距 (https://arxiv.org/html/2607.01408#S3.SS1)
   2. 3.2 用于恢复表达能力的自旋加权球谐函数 (https://arxiv.org/html/2607.01408#S3.SS2)
   3. 3.3 实基自旋加权球谐函数 (https://arxiv.org/html/2607.01408#S3.SS3)
   4. 3.4 宇称等变自旋加权球谐函数 (https://arxiv.org/html/2607.01408#S3.SS4)
   5. 3.5 SpinGTP的高性能实现 (https://arxiv.org/html/2607.01408#S3.SS5)
   6. 3.6 专用SWSH等变层 (https://arxiv.org/html/2607.01408#S3.SS6)
4. 4 实验 (https://arxiv.org/html/2607.01408#S4)
   1. 4.1 手性Tetris分类 (https://arxiv.org/html/2607.01408#S4.SS1)
   2. 4.2 3BPA性能 (https://arxiv.org/html/2607.01408#S4.SS2)
   3. 4.3 SPICE-MACE-OFF手性子集性能 (https://arxiv.org/html/2607.01408#S4.SS3)
   4. 4.4 OC20 IS2RE直接 (https://arxiv.org/html/2607.01408#S4.SS4)
5. 5 局限性与总结 (https://arxiv.org/html/2607.01408#S5)
6. 参考文献 (https://arxiv.org/html/2607.01408#bib)
7. A 自旋加权球谐函数（SWSH） (https://arxiv.org/html/2607.01408#A1)
   1. A.1 显式公式 (https://arxiv.org/html/2607.01408#A1.SS1)
   2. A.2 等变性与自旋加权变换 (https://arxiv.org/html/2607.01408#A1.SS2)
   3. A.3 非零自旋加权Gaunt路径 (https://arxiv.org/html/2607.01408#A1.SS3)
8. B 最佳渐近运行时成本 (https://arxiv.org/html/2607.01408#A2)
9. C 实基自旋加权球谐函数 (https://arxiv.org/html/2607.01408#A3)
   1. C.1 酉矩阵Q (https://arxiv.org/html/2607.01408#A3.SS1)
   2. C.2 示例：ℓ=1的Q矩阵 (https://arxiv.org/html/2607.01408#A3.SS2)
10. D 宇称等变基的推导 (https://arxiv.org/html/2607.01408#A4)
11. E SpinGTP实现与时间比较 (https://arxiv.org/html/2607.01408#A5)
    1. E.1 张量积实现 (https://arxiv.org/html/2607.01408#A5.SS1)
    2. E.2 运行时比较 (https://arxiv.org/html/2607.01408#A5.SS2)
12. F 训练细节 (https://arxiv.org/html/2607.01408#A6)
    1. F.1 3BPA (https://arxiv.org/html/2607.01408#A6.SS1)
    2. F.2 SPICE-MACE-OFF手性子集 (https://arxiv.org/html/2607.01408#A6.SS2)
       1. F.2.1 手性分类与宇称泛化 (https://arxiv.org/html/2607.01408#A6.SS2.SSS1)
       2. F.2.2 能量与力预测 (https://arxiv.org/html/2607.01408#A6.SS2.SSS2)
    3. F.3 OC20 IS2RE (https://arxiv.org/html/2607.01408#A6.SS3)

## 1 引言

将物理对称性融入深度学习模型已成为现代物理科学人工智能的基石[zhang2025artificial, bronstein2021geometric, villar2021scalars, kondor2025principles, fei2024rotation, fu2025augmenting]。通过强制执行等变性，神经网络能够建模三维原子环境，与标准架构相比具有更优越的数据效率和泛化能力。这些模型的核心是Clebsch-Gordan张量积（CGTP），这是一种基本操作，能够实现不同角频率特征（称为不可约表示，irreps）之间的相互作用[khersonskii1988quantum, thomas2018tensor, anderson2019cormorant]。尽管CGTP表达能力强大，但其计算复杂度为O(L⁶)，其中L是最大角度度数。这种复杂度常常迫使从业者将L限制在较小值，牺牲了建模复杂分子相互作用和原子间势所需的高阶几何信息。

最近克服这一瓶颈的工作集中于通过替代数学公式来加速张量积[passaro2023reducing, lin2025tensor]。一个显著的进展是Gaunt张量积（GTP）[luo2024enabling]，它将不可约表示的张量积映射为二维傅里叶基中球谐函数的逐点乘法。通过利用卷积定理和快速傅里叶变换（FFT），GTP将渐近复杂度从O(L⁶)降低到O(L³)。然而，这种效率伴随着表达能力损失。GTP仅保留不可约表示度数ℓ₁+ℓ₂+ℓ₃之和为偶数的耦合路径，移除了奇数度数之和的路径[xie2024price]。这些缺失的路径包括反对称相互作用，例如[ℓ₁=1, ℓ₂=1, ℓ₃=1]向量叉积路径，以及形成赝标量和轴向量所需的耦合。这限制了标量GTP在宇称敏感任务中的表现，特别是涉及手性几何的任务。我们将此局限性称为反对称性差距。

为了弥补反对称性差距，最近的工作[xie2026asymptotically]表明，通过将标量球谐信号提升为不可约表示值信号，可以恢复标量GTP中的反对称路径。他们利用向量球谐函数（VSH），引入了向量信号张量积（VSTP），通过包含Wigner 9j符号的广义Gaunt积来恢复缺失的反对称路径。

在本文中，我们寻求一条更直接的数学路径来解决这个问题，提出了自旋加权Gaunt张量积（SpinGTP）。SpinGTP并非将标量球谐信号提升为向量值信号，而是直接利用自旋加权球谐函数（SWSH）及其包含Wigner 3j符号的广义Gaunt积分。自旋权重索引用带符号的自旋选择规则替代了标量零阶耦合，从而允许在存在所需自旋扇区时，实现标量GTP移除的奇数度数之和路径。我们在一个实值、带宇称标签的SWSH基中实现该算子，以获得丰富且具体的表示。

我们在包括Tetris和大型原子数据集（如3BPA[3BPA]、SPICE-MACE-OFF[kovacs2025mace]和OC20[chanussot2021open]）的多个基准上评估了SpinGTP。我们的结果表明，基于SpinGTP的网络实现了与完整O(L⁶) CGTP模型相当的精度，同时保持了与原始GTP相同的计算效率。此外，我们证明，我们的模型在手性材料性质预测方面优于以前的张量积方法，表明通过SWSH包含反对称路径对手性几何的三维几何深度学习是有益的。

我们的贡献总结如下：

- • 我们提出了SpinGTP，一种基于SWSH的等变操作，恢复了Gaunt张量积框架的数学完整性。
- • 我们提供了数学推导，表明SpinGTP公式有效地恢复了先前缺失的、通用E(3)等变性所需的奇宇称相互作用（反对称路径）。
- • 我们在Tetris、3BPA、SPICE-MACE-OFF和OC20等一系列任务上评估了SpinGTP，证明其精度与完整CGTP相当，同时通过恢复反对称路径在手性几何建模中实现了有针对性的改进。

## 2 相关工作

**等变网络架构。** E(3)等变建模领域已从早期的不变描述符（如SchNet[SchNet]和DimeNet[DimeNet++]）发展到利用不可约表示（irreps）的可操控框架。开创性的架构如NequIP[nequip]和Allegro[musaelian2023learning]展示了Clebsch-Gordan（CG）张量积优越的样本效率，而Equiformer系列[equiformer, equiformerv2]将这些原理扩展到基于注意力的机制。为了减轻边级乘积的计算负担，e2former[li2026eformer]引入了节点级消息传递方案。虽然替代性标量化方法（如PaiNN[schutt2021equivariant]和NewtonNet[haghighatlari2022newtonnet]）提供了更高的吞吐量，但它们常常绕过完整的张量积空间，可能牺牲复杂几何建模所需的正式通用性。

**张量积加速。** Clebsch-Gordan张量积（CGTP）是E(3)等变神经网络中的标准交互算子[khersonskii1988quantum]。然而，直接CGTP的复杂度随最大角度度数L按O(L⁶)增长，这限制了高阶特征的使用。最近的一些方法通过利用等变架构中的结构来降低这一成本。eSCN[passaro2023reducing]在消息传递期间将特征与边方向对齐，这将等变卷积稀疏化，并在局部边框架中简化为SO(2)计算。[lin2025tensor]通过低秩张量分解结构降低了张量积成本。[luo2024enabling]提出了Gaunt张量积（GTP），它通过球谐变换和逐点乘积来评估相互作用，将复杂度降低到O(L³)。GTP效率高，但标量GTP移除了奇数ℓ之和的路径，包括诸如向量叉积等反对称相互作用[xie2024price]。为了在不回到O(L⁶)缩放的情况下恢复表达能力，引入了向量球谐函数（VSH）基和9j重耦来恢复缺失的反对称路径[xie2026asymptotically]。该框架提供了完整的算法，实现了真正的渐近加速，通过快速谱变换达到O(L⁴ log² L)复杂度。

## 3 自旋加权球谐函数的Gaunt张量积

本节介绍自旋加权Gaunt张量积（SpinGTP）的方法。我们首先回顾标准的Gaunt张量积，并指出其主要表达能力限制，即奇数ℓ之和反对称路径的丢失。然后介绍自旋加权球谐函数（SWSH）及其自旋加权Gaunt积分。为了在实值神经网络中使用，我们构建了实SWSH基，并用宇称标签对其扩展，使得模型能够在存在所需自旋通道时表示反对称路径。最后，我们总结了实现策略和架构中的专用等变层。

### 3.1 Gaunt张量积效率与反对称性差距

为了降低标准Clebsch-Gordan收缩的O(L⁶)成本，标量Gaunt张量积（GTP）[luo2024enabling]将等变相互作用表示为球谐信号的逐点乘积。由此产生的耦合由Gaunt系数给出，该系数定义为两个输入球谐函数与一个输出球谐函数在S²上的积分，如下：

G^(ℓ₃,m₃)_(ℓ₁,m₁)(ℓ₂,m₂) = ∫_{S²} Y_{ℓ₁m₁} Y_{ℓ₂m₂} Y*_{ℓ₃m₃} dΩ = √((2ℓ₁+1)(2ℓ₂+1))/(4π(2ℓ₃+1)) C^(ℓ₃,m₃)_(ℓ₁,m₁)(ℓ₂,m₂) C^(ℓ₃,0)_(ℓ₁,0)(ℓ₂,0). (1)

如公式(1)所示，GTP保留了标量耦合C^(ℓ₃,0)_(ℓ₁,0)(ℓ₂,0)非零的Clebsch-Gordan路径子集。这个限制简化了相互作用，而空间乘积公式使得通过快速傅里叶变换实现O(L³)的算法成为可能。最终的计算可以扩展到高谱分辨率，这对直接Clebsch-Gordan收缩来说是代价高昂的。

尽管效率高，但标量GTP受到因子C^(ℓ₃,0)_(ℓ₁,0)(ℓ₂,0)的约束，该因子在ℓ₁+ℓ₂+ℓ₃为奇数时为零。这移除了所有奇数ℓ之和的耦合路径，包括诸如[ℓ₁=1, ℓ₂=1, ℓ₃=1]向量叉积路径等反对称相互作用。因此，标量GTP可能无法区分信号依赖于赝标量或手性相互作用的构型。这一局限性促使了下面引入的自旋加权扩展。

### 3.2 用于恢复表达能力的自旋加权球谐函数

自旋加权球谐函数（SWSH）[goldberg1967spin]通过添加整数自旋权重s（以及度数ℓ和阶数m，满足|s|≤ℓ）来推广普通球谐函数。s=0的情况恢复普通球谐函数。当s≠0时，SWSH表示S²上的自旋加权场，并在局部切向架的旋转下通过一个相位进行变换。详细信息见附录A。

SWSH在我们的张量积中的作用源于其广义Gaunt积分。对于三个SWSH，若自旋权重满足s₃=s₁+s₂，则Gaunt系数为：

G^(ℓ₃,m₃,s₃)_(ℓ₁,m₁,s₁)(ℓ₂,m₂,s₂) = ∫_{S²} s₁Y_{ℓ₁m₁} s₂Y_{ℓ₂m₂} (s₃Y_{ℓ₃m₃})* dΩ