EquiformerV3：高效、富有表达力且通用的SE(3)-等变图注意力Transformer的规模化 EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers
第三代等变图注意力Transformer，三重提升效率、表达力与通用性
前置知识
SE(3)-等变神经网络
SE(3)指三维空间中的特殊欧氏群（旋转+平移）。一个SE(3)-等变网络的输出会随输入坐标的旋转/平移做对应的变换，使物理量（如力、能量、向量）的几何对称性被内建到模型中。实现方式是把特征表达为SO(3)不可约表示（irreps），用张量积（tensor product）作为核心等变运算。
本文所有改进（合并LN、SwiGLU-S2激活、平滑截断）都建立在严格保持SE(3)等变性的前提下，是讨论的前提基础。
不可约表示（irreps）与张量积
SO(3)的不可约表示由非负整数L（度）刻画，维度2L+1，按Wigner-D矩阵变换。L越大，能编码越精细的角向结构（如力的方向）。张量积通过Clebsch-Gordan系数把不同L的特征耦合起来，复杂度约为 $O(L_{\max}^6)$，是等变GNN的主要瓶颈。
本文SwiGLU-S2激活和eSCN卷积都围绕如何更便宜地做张量积来设计，是理解本文方法论的关键。
S2激活与球面网格采样
Cohen等人提出的S2激活先把irreps特征投影到单位球面 $S^2$ 上的网格（用 $\phi, \theta$ 离散化），在网格上做普通非线性，再投影回irreps空间。这种网格-非线性-反投影流程相当于在频域做点积，可同时混合多个度的信息。
EquiformerV2已使用S2激活，本文SwiGLU-S2是对S2的关键改造，必须先理解S2的工作机制才能体会改进点。
eSCN卷积与SO(2)线性层
Passaro & Zitnick提出的eSCN把张量积拆成两步：先用相对位置 $\vec{r}_{ij}$ 对特征做旋转 $D_{ij}$，再施加SO(2)线性层。SO(2)线性层只作用在每个L对应的2L+1通道上，大幅降低复杂度。配合截断高阶M（$M_{\max}$）效果更佳。
EquiformerV3的1.75×加速主要靠把SO(2)中的排列矩阵预先融合进 $D_{ij}$，属于eSCN实现层的优化。
SwiGLU激活
SwiGLU是Shazeer提出的门控线性单元：$\text{Swish}(x_1) \odot x_2$。PaLM、LLaMA、DeepSeek等大模型都用它替代ReLU/GELU，能稳定提升性能。一般解释是引入了输入特征的二阶交互。
本文将SwiGLU从一维标量推广到S2网格特征上，是命名SwiGLU-S2的由来。
研究动机
在3D原子尺度建模中，SE(3)-等变图神经网络虽然通过内建旋转/平移对称性提高了数据效率，但仍存在三重局限。其一，计算成本：传统张量积复杂度 $O(L_{\max}^6)$，迫使实际只能用 $L_{\max}\le 4$ 这样的低阶表示，限制了对方向性强的物理量（如原子受力）的拟合能力。其二，理论表达力：现有S2激活在球面网格上做标量非线性，从图论角度看等价于2-body标量化，无法区分只差角度不差距离的几何图（Joshi等2023已证明），而真实势能面（PES）天然包含多体相互作用。其三，泛化性：前代Equiformer、EquiformerV2主要针对单点DFT的能量和力优化，缺少对PES连续性、能量守恒的考虑，因此在分子动力学、热导率等需要高阶导数的下游任务中表现不佳，典型例子是EquiformerV2在Matbench Discovery的 $\kappa_{\text{SRME}}$ 高达1.676，远差于GRACE-2L-MPtrj的0.525。
本文的目标是本文希望从效率、表达力、通用性三个维度同时推进等变GNN。具体目标包括：在OC20 S2EF-2M上获得 $\ge 1.5\times$ 训练加速；引入能严格保持等变性的多体交互激活，使模型能区分3-body及以上几何图；用平滑截断让PES连续可微，从而胜任能量守恒模拟；在OC20、OMat24、Matbench Discovery三大主流基准上同时达到SOTA。理想结果是不靠堆数据、不靠扩大模型，仅靠架构+激活+工程协同就能在小模型（<100M）下追平甚至超过700M级UMA-L，训练时间也压缩到接近1/20。
与已有工作不同的是，与UMA（依赖跨化学域预训练+大模型）、eSEN（强调截断但未用合并LN）、MACE（侧重多体但用球谐通道乘积）等已有路线不同，本文选择从架构+实现+激活三处协同下手：把SO(2)排列矩阵预先融合进旋转矩阵以降低eSCN开销；用单个共享RMS统一所有度做LN让训练动力学更稳；提出只在标量上做sigmoid、在网格上做元素乘的SwiGLU-S2激活，把二阶交互带回到S2路径中，同时把S2网格点数减半仍保持严格等变。整套改进几乎不需要新数据，仅靠算法和工程协同就把通用原子势推到了新的帕累托前沿。
核心方法
可以把EquiformerV3想成一辆在原子尺度建模赛道上的第三代赛车：底盘（eSCN卷积+Transformer）沿用EquiformerV2，但本文做了三件事——把车架上的冗余零件（排列矩阵）焊死以减重提速；换上更合理的悬挂（合并LN、加大FFN、平滑注意力截断）以获得更好的过弯稳定性；最关键的是换了一台新引擎SwiGLU-S2：它在标量分支上用Sigmoid控制两个网格特征相乘，模拟自张量积的多体交互，同时让等变性误差不再依赖密集网格。整套架构对外保持SE(3)严格等变，但内部运算更稀疏、更快、更有表达力。
核心创新是把SwiGLU从一维标量推广到S2网格：$\text{SwiGLU-S2}(x_{\text{scalar}}, x_{\text{grid1}}, x_{\text{grid2}}) = \text{Sigmoid}(x_{\text{scalar}}) \odot x_{\text{grid1}} \odot x_{\text{grid2}}$，其中 $x_{\text{scalar}}$ 是标量门控、$x_{\text{grid1}}, x_{\text{grid2}}$ 是网格特征。收益有四点：网格乘法等价于irreps空间中的自张量积，能引入3-body以上标量化，弥补S2只能区分2-body的缺陷；Sigmoid只作用在标量上，避免向网格注入高频成分，使严格等变所需网格数从 $324$ 降到 $160$，节省50.6%；与平滑半径截断组合后PES对位置连续可微，能量守恒性能大幅提升。这是首例在等变GNN中显式引入门控S2激活的工作。
方法步骤详情
EquiformerV3沿用EquiformerV2的Transformer块。**嵌入**：原子one-hot映射节点特征；边经 $\tilde{D}_{ij}=S\cdot D_{ij}$ 旋转的SO(2)层聚合。**合并LN**：先算每度RMS $\sigma^{(L)}$，合并为共享 $\sigma$ 统一标准化，保留度间相对幅值。**注意力**：拼 $(x_i,x_j)$ 后旋转切出标量与全度 $f^{(L)}_{ij}$；前者经LN+LeakyReLU+线性得logit $z_{ij}$，envelope乘softmax分子分母得平滑权重；后者经第二条SO(2)层+SwiGLU-S2得 $v_{ij}$，$m_{ij}=a_{ij}\cdot e(\|\vec{r}_{ij}\|)\cdot v_{ij}$ 反旋转。**FFN**：隐层 $4\times d_{\text{embed}}$，用SwiGLU-S2。**训练**：OC20用 $\lambda_E=4,\lambda_F=100$；OMat24叠加DeNS去噪（$p=0.5,\sigma=0.025$）。
技术新颖性
技术新颖性主要体现在三点。第一，SwiGLU-S2把NLP里的SwiGLU在S2路径上做了标量门控+网格值的分解，既保留了S2激活多度混合的优势，又通过让非线性只走标量分支把严格等变所需的网格数砍掉一半。这是首例在等变GNN中显式引入门控S2激活的工作。第二，合并LN把跨度的相对幅度一次性保留，区别于EquiformerV2的标量独立+非标量共享两段式LN，是一类新的等变归一化方案。第三，把eSCN的SO(2)排列矩阵融合到 $D_{ij}$ 中以及修复torch.compile动态shape问题，是少见的等变实现工程化工作，1.75×的加速完全来自对同一数学运算的更优编排，没有改数学也没有改精度上限。
实验结果
**OC20 S2EF-2M消融（Table 1）**：从EquiformerV2基线能量MAE 296 meV/力MAE 21.23 meV/Å出发，叠加七项改进后误差降至201/18.15 meV，能量-32%、力-14.5%，训练时间从270→171 GPU-h，对更大EquiformerV2省5.9×。**OMat24（Table 2）**：$L_{\max}=4$ 梯度微调后能量10.1 meV/atom、力41.6 meV/Å、应力2.5 meV/ų，49M参数即可与700M UMA-L可比，比EquiformerV2-L小5×仍可比。**Matbench Discovery（Table 3, 4）**：仅MPtrj训练时EquiformerV3 $\kappa_{\text{SRME}}$ 从EquiformerV2的1.676骤降到0.275，CPS 0.830；OAM大设置下F1=0.931、CPS=0.902成为首个CPS>0.9模型，比UMA-M-1.1省22.6×训练时间，热导率比eSEN改善18–31%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| OC20 S2EF-2M（直接预测） | 能量MAE (meV) ↓ | 201 | EquiformerV2基线 296 | −32% |
| OC20 S2EF-2M（直接预测） | 力MAE (meV/Å) ↓ | 18.15 | EquiformerV2基线 21.23 | −14.5% |
| OC20 S2EF-2M（直接预测） | 训练时间 (GPU-h) ↓ | 171 | EquiformerV2基线 270 | 1.58× 加速（5.9× 对更大EquiformerV2） |
| OMat24 validation | 力MAE (meV/Å) ↓ | 41.6 (Lmax=6, grad. ft.) | UMA-L 43.5（700M参数） | 可比较但小23× |
| OMat24 validation | 能量MAE (meV/atom) ↓ | 10.5 (Lmax=4) | UMA-L 9.7 | 略高，但参数仅34M |
| Matbench Discovery (compliant MPtrj) | CPS ↑ | 0.830 | EquiformerV2 0.522; eSEN 0.797 | +59% (相对EquiformerV2) |
| Matbench Discovery (compliant MPtrj) | κ_SRME ↓ | 0.275 | EquiformerV2 1.676; eSEN 0.340 | −84% (相对EquiformerV2) |
| Matbench Discovery (non-compliant OAM) | CPS ↑ | 0.902 | UMA-M-1.1 0.889 | 首个CPS>0.9模型 |
| Matbench Discovery (non-compliant OAM) | F1 ↑ | 0.931 | PET-OAM-XL 0.924; UMA-M-1.1 0.929 | +0.007 vs PET-OAM-XL |
| Matbench Discovery (non-compliant OAM) | κ_SRME ↓ | 0.118 | eSEN 0.170; NequIP-OAM-XL 0.125 | −31% vs eSEN |
局限与改进
作者在结论与Matbench Discovery讨论里明确指出两个瓶颈。其一，在OMat24上 $L_{\max}=4$ 与 $L_{\max}=6$ 的性能差距小，提升最大度带来的收益已经饱和，说明此时数据/训练流程而非模型容量是瓶颈。其二，在Matbench上训练 $L_{\max}=6$ 几乎无收益，提示数据筛选（如过滤rattled结构）和扩展训练分布（加声子、双原子数据）才是进一步提升的关键。此外还有几个值得关注的点：消融中平滑截断在直接预测下对单点误差无影响（Index 6），其价值主要体现在下游PES连续性，单看单点指标会低估其意义；合并LN虽在2M子集上有效，但作者承认该方案最初是EquiformerV2代码库中未发表的边角发现，缺少跨数据集的独立消融；SwiGLU-S2的50.6%网格削减是在 $L_{\max}=6, M_{\max}=2$ 的OC20默认配置下测得，更高 $M_{\max}$ 设置下的等变性误差与加速比尚未给出。
独立分析的弱点
第一，SwiGLU-S2对网格几何依赖较强：作者只在 $L_{\max}=2,4,6$ 做了离散扫描，未给出采样点位置（高斯-勒让德vs等距）系统分析，用户想换 $L_{\max}=8$ 等体系时 $R_\phi,R_\theta$ 仍需重搜索。改进方向：自适应soft-grid或切比雪夫基展开。第二，合并LN用单 $\sigma$ 统一所有度，对度间差异大的特征是否最优未做对照。第三，最大模型仅91M，UMA-L 700M仍未被打败，提示模型容量与表达力非同一回事——可考虑把SwiGLU-S2扩展到MoE。第四，作者未研究 $L_{\max}=2,3$ 等极简配置在低算力场景的可用性。第五，注意力envelope贡献未与logit上envelope做对照消融。
未来方向
作者明确指出的方向有三：（1）数据筛选与扩展——过滤rattled结构、加入声子（Koker等2026）、双原子（Wood等2025）以缓解OMat24/Matbench数据瓶颈；（2）把SwiGLU-S2推广到其他高阶等变GNN（如MACE）以验证其通用性；（3）用定制CUDA内核或cuEquivariance/OpenEquivariance进一步加速eSCN。基于论文成果自然延伸的方向至少包括：把SwiGLU-S2与多任务预训练（UMA式跨域训练）结合，看是否能用更小模型达到UMA的精度；探索在分子动力学模拟里用SwiGLU-S2+合并LN做端到端能量守恒训练，验证 $\kappa_{\text{SRME}}$ 提升能否转化为真实材料热导率预测精度的提升；以及把SwiGLU思想拓展到更高阶S3（4D球面）流形上，应对周期性或角动量守恒体系。
复现评估
本文可复现性做得相当充分。代码已开源在GitHub（atomicarchitects/equiformer_v3），预训练权重同步发布到HuggingFace（mirror-physics/equiformer_v3），下游任务可直接加载微调或推理。训练数据全部公开（OC20、OMat24、MPtrj、sAlex、Matbench协议），无私有数据。算力方面：OC20 S2EF-2M约170 H100 GPU-h，OMat24直接预训练+梯度微调需 $2188+2594$ GPU-h（$L_{\max}=4$）或 $3320+5877$ GPU-h（$L_{\max}=6$），Matbench需5.7k H100 GPU-h——工业实验室尚可承受，中小机构难以从零复现，建议直接用预训练checkpoint做下游任务。Table 8–11列了完整超参与DeNS配方。
论文图表
