Equivariant Spherical Transformer for Efficient Molecular Modeling¶

会议: NeurIPS 2026
arXiv: 2505.23086
代码: https://github.com/anjunyi/EST
领域: 图学习
关键词: 等变神经网络, 球面傅里叶变换, 分子表征学习, 球面注意力, 混合专家系统

一句话总结¶

针对传统等变图神经网络依赖 Clebsch-Gordan 张量积卷积导致非线性表达受限与高次表征计算代价高昂的问题，本文提出 Equivariant Spherical Transformer (EST)，通过基于斐波那契晶格与分子动力学校准的均匀球面傅里叶采样将群表征投射至连续空间域，利用跨方向球面注意力与混合专家 FFN 显著提升表征能力并严格保持 SO(3) 等变性。

研究背景与动机¶

三维等变图神经网络（SE(3)-equivariant GNN）已成为模拟分子物理相互作用、加速量子化学计算与材料设计的核心范式。与早期依赖原子间标量距离的不变性模型（如 SchNet）相比，现代等变模型引入高次可导向向量（steerable vectors）作为节点嵌入，并通过 Clebsch-Gordan（CG）张量积实现等变张量卷积，从而直接端到端拟合高阶原子几何相互作用。

然而，基于 CG 张量积的消息传递机制存在内在的表达瓶颈。首先，CG 张量积运算天然缺乏高阶非线性，其计算复杂度随群表征最高阶数呈 $\mathcal{O}(L^6)$ 激增；尽管后续工作引入 Gate 机制或张量注意力，但非线性变换仍被局限在 $l=0$ 的标量不变分量上，对 $l>0$ 的高阶不可约表征仅能进行标量放缩，严重限制了多原子高阶构型的交互表达力。其次，若转而在傅里叶空间域处理球面函数（如 SCN、eSCN），已有方法通常采用点卷积（point-wise FFN），仅对同一方向上的通道进行特征提取，完全忽略了球面不同方向之间的几何依赖关系；更严峻的是，传统基于 e3nn 经纬度网格的采样方式极不均匀，极点附近采样点密集而赤道稀疏，在引入跨方向注意机制时会剧烈破坏旋转等变性。

本文的切入角度是：如果在群表征的傅里叶空间域进行高均匀度的离散采样，使得采样点集近似在任意 SO(3) 旋转下自闭合，便可以在球面离散点序列上引入全连接 Transformer 机制来建模任意方向对之间的相互作用，从数学上突破张量积的阶数表达上限。核心 idea：提出即插即用的 Equivariant Spherical Transformer (EST)，利用基于斐波那契晶格与斥力动力学校准的均匀球面傅里叶采样建立空间点序列，以相对方向感知的跨方向球面注意力机制捕获任意空间方向间的非线性交互，并结合混合专家结构在保证严谨 SO(3) 等变性的同时打破传统 CG 张量积的表达力上限。

方法详解¶

整体框架¶

EST 作为一个即插即用的等变模块，能够直接替代等变 GNN 中的消息传递块（Message block）或更新块（Update block）。整体数据流分为三个阶段：首先，通过截断阶数 $L$ 的正向球面傅里叶变换（SFT），将节点或边的多阶可导向向量嵌入映射到连续球面函数，并利用经过斥力动力学校准的斐波那契晶格采样为 $S$ 个均匀分布的空间点特征序列；接着，在该点序列上施加带相对方向嵌入（Relative Orientation Embedding）的多头球面注意力（Spherical Attention），建模不同空间取向之间的深层几何相关性；最后，通过由不可约可导向 FFN 与空间点 FFN 组成的混合专家网络（MoE）自适应融合特征，并通过逆球面傅里叶变换无损恢复回任意指定阶数的不可约表征空间。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["输入不可约群表征<br/>Type-(0→L) Steerable 向量"] --> B["均匀斐波那契晶格采样<br/>空间域离散化 S 个方向点"]
    B --> C["相对方向嵌入<br/>Q, K 拼接 3D 单位球面向量"]
    C --> D["跨方向球面注意力<br/>Softmax 建模任意两方向交互"]
    D --> E["混合专家前馈网络<br/>可导向 FFN 与空间 FFN 门控融合"]
    E --> F["逆球面傅里叶变换<br/>投影回不可约群表征空间"]

关键设计¶

1. 均匀斐波那契晶格与排斥动力学校准采样：保障离散傅里叶变换下的严格 SO(3) 等变性 针对现有方法使用 e3nn 经纬度网格在极点和赤道密度剧烈变化、导致旋转后离散采样点集不对称并严重破坏等变性的缺陷，本文从理论上证明（定理 1）：当且仅当采样点集 $\mathcal{P} = \{\vec{\mathbf{p}}_1, \dots, \vec{\mathbf{p}}_S\}$ 在旋转群 $\text{SO}(3)$ 作用下近似自闭合（即任意点经旋转后仍能在点集中找到对应邻近点）时，空间 Transformer 的输出才严格保持等变性。为此，EST 首先采用基于黄金分割比 $\lambda = (1+\sqrt{5})/2$ 的斐波那契晶格（Fibonacci Lattice, FL）初始化球面点坐标： $$\vec{\mathbf{p}}_s = \left[\sqrt{1-z_s^2}\cos(2s\pi/\lambda),\; \sqrt{1-z_s^2}\sin(2s\pi/\lambda),\; z_s\right], \quad z_s = \frac{2s-1}{S-1}$$ 该螺旋式分布在不同纬度具有均一的面密度。为进一步消除局域各向异性，EST 在模型初始化阶段引入类分子动力学的排斥势能模拟，通过多轮点间斥力弛豫与球面重投影获得极高对称性的点分布，并重构逆变换的离散共轭函数系数 $\lambda^{(l,m)} = 1/\sum_s |Y^{(l,m)}(\vec{\mathbf{p}}_s)|^2$。离线固化采样点后，离散傅里叶变换可以在常数开销下精确逼近连续空间积分，将旋转误差降低 2 个数量级。

2. 相对方向感知的跨方向球面注意力机制：打破 CG 张量积阶数上限 传统空间域方法（如 SCN）仅对单个采样方向独立做 MLP，丢失了跨方向的几何构型信息；而传统不可约张量积运算在表征阶数为 $L$ 时，理论上无法区分超过 $L$ 重对称的几何结构。EST 将离散化后的 $S$ 个采样点展平为长度为 $S$、特征维度为 $C$ 的 token 序列，在空间域直接计算全连接注意力权重： $$a_{s_i, s_j} = \frac{\exp\left((\mathbf{Q}_{s_i}\mathbf{K}_{s_j}^T + \vec{\mathbf{p}}_{s_i}^T \vec{\mathbf{p}}_{s_j})/\sqrt{C}\right)}{\sum_{s_k=1}^S \exp\left((\mathbf{Q}_{s_i}\mathbf{K}_{s_k}^T + \vec{\mathbf{p}}_{s_i}^T \vec{\mathbf{p}}_{s_k})/\sqrt{C}\right)}, \quad \tilde{\mathbf{f}}^{*}_{s_i} = \sum_{s_j=1}^S a_{s_i, s_j} \mathbf{V}_{s_j}$$ 为了使注意力不仅能识别局部通道特征，还能感知点之间的相对几何取向且保持旋转对称性，EST 将采样点的 3D 笛卡尔坐标向量直接拼接到 Query 与 Key 向量中。当体系整体旋转时，内积项 $(\mathbf{R}\vec{\mathbf{p}}_{s_i})^T (\mathbf{R}\vec{\mathbf{p}}_{s_j}) = \vec{\mathbf{p}}_{s_i}^T \vec{\mathbf{p}}_{s_j}$ 保持严格旋转不变，赋予模型显式感知相对空间夹角的能力。本文在理论上证明，该结构生成的函数空间严格包含 CG 张量积函数空间，使单层低阶（$L=1$）EST 即可完全区分 100 重对称构型。

3. 双域混合专家前馈网络（Hybrid MoE）：平衡严谨等变约束与高容量非线性表达 纯空间域前馈网络虽然具备强非线性容量，但在极小样本或极高对称性约束下可能因采样离散化带来轻微的泛化漂移；而纯不可约可导向 FFN 虽严格解析等变，但表达能力受限于张量阶数。EST 设计了跨谐波域与空间域的双域混合专家模块。由输入消息的标量不可约分量 $\mathbf{m}^{(0)}$ 经过门控网络学习路由权重： $$G(\mathbf{m}) = \text{Softmax}(\mathbf{m}^{(0)} \mathbf{W}_G)$$ 专家库由两类不同物理属性的专家并行构成：一类是作用于不可约可导向向量的 Steerable FFN，由两层度逐项线性映射（DW-Linear）配合 Gate 激活函数构成，保证严谨的数学等变性；另一类是作用于傅里叶离散空间域的 Spherical FFN，由两层点逐项线性映射（PW-Linear）配合 SiLU 激活函数构成，提供强大的高阶非线性拟合能力。门控网络根据局部化学环境自适应调节两类专家的激活权重，在严谨物理对称性与非线性表征容量之间取得最优权衡。

损失函数 / 训练策略¶

在 OC20 S2EF 能量与受力联合预测任务中，模型采用多任务联合损失函数： $$\mathcal{L} = \lambda_E \cdot |\hat{E} - E| + \lambda_F \cdot \frac{1}{3N}\sum_{i=1}^N \|\hat{\mathbf{F}}_i - \mathbf{F}_i\|_1$$ 其中 $\lambda_F = 100$，$\lambda_E \in \{2, 4\}$。优化器选用 AdamW，采用带线性热身的余弦退火学习率调度策略（初始最大学习率 $4\times 10^{-4}$，Batch Size 为 256，权重衰减 $10^{-3}$，EMA 平滑衰减率 0.999）。在小规模分子数据集 QM9 上，模型以平均绝对误差（MAE）为损失函数，采用 6 层 EST 架构与 Cosine 衰减调度（Epochs 300~700），针对能量类标量与几何极化率等物理量分立调整超参数以确保最优收敛。

实验关键数据¶

主实验¶

在 Open Catalyst 2020 (OC20) S2EF 验证集（所有结构 S2EF-All）上，8 层 EST 仅用标准训练集即可全面超越包含 20 层深层网络及引入分子动力学额外数据的 SOTA 模型：

模型	参数量	吞吐量 (Samples/s)	训练数据	能量 MAE (meV) ↓	受力 MAE (meV/Å) ↓
SchNet	9.1M	-	S2EF-All	548.5	56.8
DimeNet++-L-F+E	10.7M	4.6	S2EF-All	515.5	32.8
GemNet-OC	39.0M	18.3	S2EF-All	244.3	21.7
SCN (20 layers)	271.0M	-	S2EF-All	244.0	17.7
eSCN (20 layers)	200.0M	2.9	S2EF-All	242.0	17.1
EScAIP-Small	83.0M	-	All + MD	229.0	15.1
EquiformerV2 ($\lambda_E=2$, 20 layers)	153.0M	1.8	S2EF-All	236.0	15.7
EquiformerV2 ($\lambda_E=4$, 8 layers)	31.0M	7.1	All + MD	232.0	16.26
EST ($\lambda_E=4$, 8 layers, 本文)	45.0M	6.8	S2EF-All	229.0	15.8
EST ($\lambda_E=2$, 8 layers, 本文)	45.0M	6.8	S2EF-All	232.0	15.0

在量子化学小分子基准 QM9 的 12 项理化性质预测中，EST 在 8 项性质上取得最优结果：

模型	$\alpha$ (bohr$^3$) ↓	$\Delta\varepsilon$ (meV) ↓	$\varepsilon_{\text{HOMO}}$ (meV) ↓	$\varepsilon_{\text{LUMO}}$ (meV) ↓	$\mu$ (D) ↓	$C_v$ (cal/mol K) ↓	$G$ (meV) ↓	$H$ (meV) ↓	$U_0$ (meV) ↓
DimeNet++	0.044	33	25	20	0.030	0.023	8.00	7.00	6.00
PaiNN	0.045	46	28	20	0.012	0.024	7.35	5.98	5.85
SEGNN	0.060	42	24	21	0.023	0.031	15.00	16.00	15.00
Equiformer	0.046	30	15	14	0.011	0.023	7.63	6.63	6.59
EquiformerV2	0.050	29	14	13	0.010	0.023	7.57	6.22	6.17
EST (本文)	0.042	28	13	12	0.011	0.022	7.03	5.94	5.64
EST with GA (本文)	0.041	29	14	13	0.011	0.021	7.18	6.17	6.32

消融实验¶

针对采样策略与等变性保持能力的测试，在随机旋转 1000 个 QM9 分子构型的未训练网络上评估旋转绝对误差（$\times 10^{-3}$）：

采样策略与点数配置	1 层网络误差	2 层网络误差	4 层网络误差	6 层网络误差	说明
e3nn 网格 / $S=210$ (无优化)	8.4	100.6	59.3	19.9	传统网格极不均匀，误差随层数剧烈发散
斐波那契晶格 (FL) / $S=64$ (无优化)	1.0	1.1	1.8	4.0	螺旋均匀采样显著降低等变损失
斐波那契晶格 (FL) / $S=256$ (无优化)	0.2	0.2	0.2	0.3	提高点密度可逼近完全对称
FL + 动力学斥力优化 / $S=64$	0.2	0.1	0.3	0.2	经斥力校准后在仅 64 点下实现最强等变保真度

针对高阶旋转对称结构区分能力（GWL 表达力测试）的消融结果表明，在 $n$-fold 对称区分任务中，传统 TFN/MACE 在最高阶数 $L=1$ 时对于 $n \ge 2$ 的对称构型分类准确率全部跌至 50.0%（无法区分）；当 $L=2$ 时对 $n \ge 3$ 失效。而 EST 在 $L=1$ 时对 2-fold、3-fold、10-fold 乃至 100-fold 的对称结构均保持 100.0% 的完美区分能力，从实证上验证了其对张量积表达上限的彻底突破。

关键发现¶

采样均匀度是空间域等变网络成败的核心分水岭：使用 e3nn 经纬度采样时，网络层数增加导致旋转等变误差快速累积至 $10^{-1}$ 量级，引发严重泛化退化；而 FL 配合排斥动力学校准可将离线采样误差稳定控制在 $10^{-4}$ 以内，确保了深层堆叠的数值自洽。
空间注意力机制兼备高吞吐与高表征力：在 OC20 S2EF 上，8 层的 EST（45M 参数）仅耗时 6.8 Samples/s 的高吞吐，即在受力预测上达到 15.0 meV/Å，超越了参数量高达 153M 的 20 层 EquiformerV2（1.8 Samples/s）与 200M 参数的 eSCN（2.9 Samples/s）。
混合专家系统的密集协同效果最佳：QM9 上的 MoE 消融显示，设置 10 个专家且采用 Dense 路由（而非语言模型中常用的稀疏 Top-K）能够取得最佳的理化预测精度，因为分子体系整体数据规模较小，严苛的稀疏路由与负载均衡损失反而会削弱网络对细微几何构型变化的捕获。

亮点与洞察¶

将群表征变换与现代 Transformer 机制无缝衔接：通过正逆球面傅里叶变换充当桥梁，将晦涩繁重的不可约群表征张量积卷积简化为连续球面上的自注意力建模，彻底绕开了阶数爆炸问题。
离线排斥动力学校准技巧：在模型初始化阶段通过类分子动力学模拟优化球面采样点分布并持久化缓存，零推理运行时开销地解决了连续球面积分的离散几何畸变问题。
双域协同的 Hybrid MoE 范式：不仅将 MoE 引入物理等变建模，更创造性地将专家拆分为物理等变保底流（Steerable FFN）与空间高维表征流（Spherical FFN），兼顾了第一性原理的严谨性与深度模型的表达力。

局限与展望¶

采样点集离散化带来的极微弱等变损失：虽然实测旋转误差低于 $10^{-4}$，但与解析可导向张量积的严格数学等变相比，有限点离散化仍存在微小近似误差，未来可探索准连续或解析闭合的球面注意力形式。
超大体系的球面采样点计算开销：随着分子规模与图节点数量扩大，为满足高阶 Nyquist 采样定理所需的采样点数量 $S \ge (2L)^2$ 会使注意力的内存开销呈平方级上升，未来可结合稀疏线性注意力或球谐多重极展开进行进一步加速。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙利用空间傅里叶变换与斐波那契均匀采样解决了等变空间注意力的破称问题，理论证明与架构设计极为漂亮。
实验充分度: ⭐⭐⭐⭐⭐ OC20、QM9、GWL 表达力与旋转误差消融全面扎实，对比基准广泛且指标极具说服力。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，动机阐述层层递进，图表表达清晰。
价值: ⭐⭐⭐⭐⭐ 作为高效通用的即插即用等变组件，对材料模拟、药物化学与等变大模型架构具有重大参考意义。

采样策略与点数配置	1 层网络误差	2 层网络误差	4 层网络误差	6 层网络误差	说明
e3nn 网格 / \(S=210\) (无优化)	8.4	100.6	59.3	19.9	传统网格极不均匀，误差随层数剧烈发散
斐波那契晶格 (FL) / \(S=64\) (无优化)	1.0	1.1	1.8	4.0	螺旋均匀采样显著降低等变损失
斐波那契晶格 (FL) / \(S=256\) (无优化)	0.2	0.2	0.2	0.3	提高点密度可逼近完全对称
FL + 动力学斥力优化 / \(S=64\)	0.2	0.1	0.3	0.2	经斥力校准后在仅 64 点下实现最强等变保真度