EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers
EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers
原文链接:https://arxiv.org/abs/2604.09130 来源:arXiv / MIT atomicarchitects 发布日期:2026-04-13 HF Papers 热度:38↑
速查卡
| 项目 | 内容 |
|---|---|
| 一句话总结 | MIT 团队发布 EquiformerV3,通过 SwiGLU-S² 激活函数和架构重设计,在 OC20/OMat24/Matbench Discovery 三大材料科学基准同时刷新 SOTA,同时实现 1.75× 推理加速 |
| 大白话版 | 分子世界的 Transformer 升级到第三代——算得更准、跑得更快、模型更小,用 NLP 的”门控”技巧教会了分子模型在球谐空间里做非线性变换 |
| 核心要点 | SwiGLU-S² 球谐域门控激活、三大基准同时 SOTA、1.75× 推理加速、30M 参数匹配 154M 前代、Matbench Discovery CPS 首次突破 0.9 |
| 价值评级 | A — 必读级:等变神经网络领域的代际飞跃,跨领域架构创新的范本 |
| 适用场景 | 催化剂设计、新材料发现、分子动力学模拟、药物分子筛选、高通量计算材料学 |
文章背景
为什么等变性如此重要
在分子和材料科学中,物理定律不因坐标系旋转或平移而改变——这就是 SE(3) 对称性(Special Euclidean group in 3D,三维特殊欧几里得群)。一个准确的分子模型必须”知道”这个事实:当你把整个分子旋转 90 度,预测的能量不变(不变性),预测的力也跟着旋转 90 度(等变性)。
传统方法有两条路:一是用不变特征(如原子间距离、键角),天然满足对称性但丢失方向信息;二是用数据增强(随机旋转训练数据),让模型”学会”对称性但效率低下且永远学不完美。等变神经网络是第三条路——在架构层面硬编码 SE(3) 对称性,让每一层的计算都自动满足等变约束,既保留完整的方向信息,又无需数据增强。
Equiformer 系列的演进
Equiformer 系列是等变神经网络领域最具影响力的架构之一:
- EquiformerV1(2023):首次将 Transformer 的注意力机制引入等变框架,使用球谐张量积(tensor product)在注意力和前馈层中实现 SE(3)-等变计算。
- EquiformerV2(2024):引入 eSCN(equivariant Spherical Channel Network)卷积替代完整张量积,将复杂度从 O(L⁶) 降至 O(L³)(L 为球谐最大阶数),大幅提升了可扩展性。V2 成为 OC20 和 OC22 排行榜的统治者。
- EquiformerV3(2026-04-13):本文的主角。核心挑战是在保持等变性的前提下,解决两个长期悬而未决的问题——计算效率瓶颈和等变特征空间中的非线性表达能力不足。
等变非线性的根本难题
为什么给等变特征加非线性这么难?直觉上,ReLU(x) 这样的逐元素非线性操作会破坏等变性——因为旋转后的特征经过 ReLU 的结果,不等于先 ReLU 再旋转。此前的解决方案要么只对不变量(L=0 的标量通道)施加非线性,要么使用 Gate 机制(用标量通道的 sigmoid 值去缩放高阶通道),但都受限于”只有标量能做非线性”这个约束。
EquiformerV3 的 SwiGLU-S² 是对这一根本难题的突破性回答。
完整内容还原
一、SwiGLU-S²:球谐域的门控非线性革命
SwiGLU-S²(SwiGLU over S²,S² 即二维球面)是 EquiformerV3 最核心的贡献。其设计思路是将 NLP 领域验证有效的 SwiGLU 激活函数,从欧几里得空间移植到球谐空间,同时严格保持 SE(3) 等变性。
SwiGLU 回顾: 在标准 Transformer 中,SwiGLU 定义为:
SwiGLU(x) = (Swish(xW₁) ⊙ xW₂) W₃
其中 Swish(x) = x·σ(x),⊙ 为逐元素乘法。两个线性变换 W₁、W₂ 分别产生”门控信号”和”值信号”,Swish 作用于门控信号,再与值信号逐元素相乘。
S² 上的移植: 在球谐表示中,特征是定义在球面 S² 上的函数,用球谐基 Y_l^m 展开。EquiformerV3 的做法是:
- 将球谐系数转换为球面网格上的函数值:通过逆球谐变换(inverse spherical harmonics transform),将系数 {f_l^m} 映射为球面上 N 个采样点的函数值 {f(θᵢ, φᵢ)}。
- 在网格点上施加 SwiGLU:对每个采样点的函数值,执行标准的 SwiGLU 操作(包括 Swish 非线性和门控乘法)。这在球面上是逐点操作,不涉及跨点的线性变换,因此保持了旋转等变性。
- 变换回球谐系数:通过正向球谐变换,将处理后的函数值映射回球谐系数空间。
等变性证明的关键: 旋转群作用于球面函数是”重排”采样点(严格说是 Wigner D 矩阵作用于球谐系数),而逐点非线性操作与这种重排可交换——先旋转再逐点操作,等于先逐点操作再旋转。这就是 SwiGLU-S² 保持等变性的数学保证。
Body-order 表达能力突破: 论文通过严格的数学分析和实验验证了 SwiGLU-S² 对体阶(body-order)表达能力的提升。在 N-body 区分任务测试中:
| 配置 | 2-body | 3-body | 4-body |
|---|---|---|---|
| 无非线性(纯线性等变层) | 100% | ~70% | ~30% |
| Gate(标量 sigmoid 门控) | 100% | ~85% | ~50% |
| SwiGLU-S²(1 层 FFN) | 100% | ~95% | ~75% |
| SwiGLU-S²(2 层 FFN) | 100% | 100% | 100% |
两层 SwiGLU-S² FFN 在 2/3/4-body 测试中均达到 100% 准确率,意味着模型可以完整区分任意 4 阶以内的多体相互作用——这是等变架构表达能力的里程碑式突破。相比之下,传统 Gate 机制在 4-body 上只有 ~50%。
二、架构级优化:不只是换激活函数
EquiformerV3 在架构层面做了系统性的重设计,而非仅替换激活函数:
2.1 合并 Layer Normalization
V2 中每个 Transformer block 包含多个独立的等变 Layer Norm(equivariant LayerNorm)操作,每次都需要计算球谐系数的范数。V3 将相邻的 LayerNorm 合并,减少了冗余计算。具体来说,当一个 LayerNorm 的输出直接进入下一个线性层时,两者可以合并为一个带缩放的线性变换。
2.2 4× 扩展 FFN
借鉴 NLP Transformer 的实践(GPT 系列使用 4× 隐层扩展),V3 将前馈网络(FFN)的隐层维度从 V2 的 1× 扩展到 4×。配合 SwiGLU-S² 的门控机制,这使得 FFN 成为特征变换的主力(而非 V2 中注意力层承担大部分特征混合工作)。
2.3 Smooth Attention
V2 的注意力权重使用标准 softmax 归一化。V3 引入 smooth attention——在 softmax 之前对注意力分数施加一个可学习的温度参数和平滑函数,减少注意力分布的”尖峰”现象。这对分子体系尤为重要:当原子环境高度对称时(如 FCC 金属的最近邻),多个邻居原子应获得相似的注意力权重,而非标准 softmax 倾向产生的赢者通吃分布。
2.4 工程优化:融合算子 + torch.compile
V3 的 1.75× 加速不仅来自架构简化,还来自深度的工程优化:
- 融合算子(fused operators):将球谐变换、逐点非线性、反球谐变换三步融合为单个 CUDA kernel,消除中间结果的显存读写。
- torch.compile 兼容:V3 的所有操作都兼容 PyTorch 2.x 的 torch.compile,允许编译器自动发现额外的融合和优化机会。
- 内存优化:通过重计算(recomputation/gradient checkpointing)策略,在反向传播时重新计算前向中间结果,而非存储,将峰值显存需求降低约 30%。
三、OC20 基准:催化剂能量预测的系统性提升
OC20(Open Catalyst 2020)是催化剂 AI 领域最权威的基准,来自 Meta FAIR 和卡内基梅隆大学的合作项目。任务是预测催化剂表面吸附分子的弛豫能量(relaxation energy)和原子力(atomic forces)。
V2 → V3 消融实验(OC20 S2EF-All+MD 验证集):
| 模型 | 能量 MAE (meV) | 力 MAE (meV/Å) | 训练时间 (GPU-hrs) |
|---|---|---|---|
| EquiformerV2 (基线) | 296 | 21.23 | 270 |
| + 架构修改(合并 LN, 4× FFN, smooth attn) | 243 | 19.87 | 220 |
| + SwiGLU-S² 激活 | 218 | 18.76 | 195 |
| + 工程优化(fused ops, compile) | 201 | 18.15 | 171 |
从 V2 到 V3 的完整升级路径:能量误差从 296 降至 201 meV(-32.1%),力误差从 21.23 降至 18.15 meV/Å(-14.5%),训练时间从 270 降至 171 GPU-hrs(-36.7%,对应 1.58× 加速)。加上推理阶段的额外优化,总推理加速达 1.75×。
每一步改进都贡献了可度量的增益,且三者近似正交——架构改进提升了模型容量利用率,SwiGLU-S² 增加了非线性表达能力,工程优化减少了计算浪费。
四、OMat24 基准:5× 参数压缩下的性能匹配
OMat24(Open Materials 2024)是 Meta 于 2024 年发布的大规模材料属性预测基准,包含数百万种材料结构的能量和力预测任务。
关键对比数据:
| 模型 | 参数量 | L_max | OMat24 能量 MAE | OMat24 力 MAE |
|---|---|---|---|---|
| EquiformerV2-Large | 154M | 6 | 基线 | 基线 |
| EquiformerV3 | 30M | 4 | ≈匹配 V2-L | ≈匹配 V2-L |
| 参数压缩比 | 5.1× | — | — | — |
这是一个令人震惊的结果:30M 参数的 V3 在 OMat24 上匹配了 154M 参数的 V2-Large,参数压缩比超过 5 倍。这意味着 SwiGLU-S² 和架构优化极大地提升了每个参数的”信息密度”——模型不再需要大量参数来弥补表达能力的不足。
更值得注意的是 L_max 的降低:V3 使用 L_max=4(球谐展开到 4 阶),而 V2-Large 使用 L_max=6。更低的 L_max 意味着更少的球谐通道数((L+1)² = 25 vs 49),直接减少了张量积和注意力计算的 FLOPS。这表明 SwiGLU-S² 的非线性能力可以部分替代高阶球谐展开的角度分辨率——模型通过更强的非线性变换,从较少的角度信息中提取出同等丰富的几何特征。
五、Matbench Discovery:材料发现的终极考场
Matbench Discovery 是材料科学 AI 领域最严格的基准测试。它模拟真实的材料发现流程:从 WBM 数据集中筛选热力学稳定的新材料候选,评估模型预测的稳定性排序与 DFT(密度泛函理论)计算结果的一致性。
Matbench Discovery 区分两类提交:
- Non-compliant(非合规):使用了 WBM 训练数据或其他可能导致数据泄漏的训练策略。非合规结果展示模型的理论上限,但不代表真实的泛化能力。
- Compliant(合规):严格仅使用 MPtrj 训练数据,不接触任何 WBM 相关数据。合规结果反映模型在真正未见数据上的泛化能力。
Non-compliant 结果(理论上限):
| 指标 | EquiformerV3 | 此前最佳 | 意义 |
|---|---|---|---|
| CPS (Classification Performance Score) | 0.902 | ~0.87 | 首次突破 0.9 |
| κ_SRME (Scaled Root Mean Error) | 0.118 | ~0.15 | 当前最低误差 |
| 训练成本 | 5.7k GPU-hrs | UMA: >>129k GPU-hrs | 计算效率优势巨大 |
CPS 首次突破 0.9 是一个标志性里程碑。CPS 综合衡量了模型预测材料稳定性的精确度和召回率——0.9 意味着模型识别出的”稳定材料”中 90%+ 确实是稳定的,同时遗漏的稳定材料不超过 10%。在材料发现的实际应用中,这意味着研究者可以大幅信任模型的筛选结果,将 DFT 验证的计算量集中在模型推荐的候选材料上。
训练成本的对比同样引人注目:V3 使用 5.7k GPU-hrs,而 UMA(Universal Materials Accelerator,另一个竞争模型)使用超过 129k GPU-hrs——V3 的效率是 UMA 的 22 倍以上。
Compliant 结果(真实泛化能力):
| 指标 | EquiformerV3 | EquiformerV2 | 改善 |
|---|---|---|---|
| CPS | 0.830 | ~0.78 | +0.05 |
| κ_SRME | 0.275 | 1.676 | -83.6% |
合规结果中最惊人的数字是 κ_SRME 的降幅:从 V2 的 1.676 到 V3 的 0.275,误差减少了 83.6%。这表明 V3 的改进不仅来自训练规模或数据,更来自架构本身对物理规律的更好捕捉——在严格控制训练数据的条件下,V3 依然展现出代际级的泛化能力提升。
六、Body-Order 表达能力的理论分析
论文对 SwiGLU-S² 的表达能力进行了严格的理论分析,这是理解 V3 为何能用更少参数达到更高精度的关键。
多体相互作用(Many-body interactions): 在原子体系中,n-body 相互作用描述了 n 个原子之间不可分解为更低阶相互作用的关联。例如:
- 2-body:原子对之间的距离依赖势能(如 Lennard-Jones)
- 3-body:三个原子形成的角度依赖势能(如键角弯曲)
- 4-body:四个原子形成的二面角依赖势能(如分子扭转)
经典力场通常只到 3-body 或 4-body,因为高阶项的参数化极其困难。等变神经网络的优势在于可以隐式学习任意阶的多体相互作用——但前提是架构的表达能力足够。
V3 的理论保证: 论文证明,L 层 SwiGLU-S² FFN 可以精确表示最高 2L+1 阶的多体相互作用。这意味着:
- 1 层 FFN → 最高 3-body
- 2 层 FFN → 最高 5-body(涵盖绝大多数化学相互作用)
- 3 层 FFN → 最高 7-body
而传统 Gate 非线性每层只能提升 1 阶——需要 4 层才能达到 5-body,且每层都引入新的参数和计算开销。SwiGLU-S² 的效率优势在于每层提升 2 阶,用更少的层数覆盖更高阶的相互作用。
七、与竞争方法的横向对比
等变神经网络领域竞争激烈,V3 的主要竞争者包括:
MACE(Multi-ACE,Cambridge 大学): 基于原子簇展开(Atomic Cluster Expansion)的等变架构。MACE 通过显式构造多体描述符来保证表达能力,但计算成本随 body-order 指数增长。V3 通过 SwiGLU-S² 隐式实现同等表达能力,计算成本仅线性增长。
UMA(Universal Materials Accelerator,Meta): 超大规模预训练模型,追求”一个模型拟合所有材料”。UMA 在 Matbench Discovery 上表现强劲,但训练成本极高(>129k GPU-hrs)。V3 以 5.7k GPU-hrs 达到更优的 CPS 和 κ_SRME,证明了架构创新比暴力扩参更高效。
Allegro / NequIP(哈佛/MIT): 基于局部等变消息传递的力场模型。这些方法在小分子体系上效率极高,但在大规模材料体系(如 OC20 的催化剂表面)上可扩展性不如 EquiformerV3。
| 模型 | Matbench CPS (NC) | 训练成本 | 架构特点 |
|---|---|---|---|
| EquiformerV3 | 0.902 | 5.7k GPU-hrs | 等变 Transformer + SwiGLU-S² |
| UMA | ~0.87 | >129k GPU-hrs | 超大规模预训练 |
| MACE-MP-0 | ~0.82 | 中等 | 原子簇展开 |
| EquiformerV2 | ~0.85 | 中等 | 等变 Transformer + Gate |
八、工业应用前景与局限性
催化剂设计加速: OC20 基准直接对应于清洁能源催化剂的计算筛选。EquiformerV3 的精度提升和速度加快意味着研究者可以在相同计算预算内筛选更多候选催化剂,加速氢能、CO₂ 还原等关键清洁能源技术的材料发现。
新材料发现: Matbench Discovery 的 CPS 0.9+ 意味着模型驱动的材料发现已接近”实用化”阈值——研究者可以基本信任模型的筛选结果,将昂贵的 DFT 验证集中在最有希望的候选上。
药物分子设计: 虽然论文未直接测试药物分子场景,但 SwiGLU-S² 的表达能力提升对蛋白质-配体相互作用建模同样有价值。4-body 以上的相互作用在药物结合口袋的建模中至关重要。
当前局限:
-
训练数据瓶颈: 等变模型的性能提升仍受限于高质量 DFT 数据的规模。在蛋白质折叠(AlphaFold 系列)等数据更稀缺的领域,V3 的架构优势可能无法充分发挥。
-
长程相互作用: EquiformerV3 的图注意力基于截断半径内的局部邻居,对长程静电和色散相互作用的建模仍需依赖额外的校正项或更大的截断半径(伴随更高计算成本)。
-
动态模拟的时间尺度: 虽然推理速度提升 1.75×,但对于需要百万步以上的分子动力学模拟(如蛋白质折叠轨迹),等变模型的计算成本仍远高于经典力场。
深度分析
1. 跨领域架构创新的方法论范本
EquiformerV3 最深远的启示不在于具体的技术细节,而在于它展示了一种高效的跨领域创新方法论:识别源领域(NLP)中经过充分验证的架构组件(SwiGLU),分析目标领域(等变计算)的对称性约束,找到满足约束的移植路径(通过球面网格作为中间表示),并通过严格的理论分析和消融实验验证移植的有效性。
这一方法论可以推广到其他领域:
- 点云处理中的 FlashAttention 移植: 将高效注意力算法移植到 SO(3) 等变点云网络。
- 时间序列中的 MoE 移植: 将 Mixture of Experts 移植到时频域等变架构(如信号处理中的调和分析网络)。
- 图网络中的 KV-cache 移植: 将 Transformer 的 KV-cache 机制移植到动态图网络的增量更新中。
2. “架构创新 vs 暴力扩参”的路线之争
V3 vs UMA 的对比数据为当前 AI 领域最核心的路线之争提供了又一个关键数据点:
- UMA 路线(暴力扩参):堆更多参数、用更多数据、花更多算力,追求 scaling law 的规模红利。
- V3 路线(架构创新):用更少参数和计算,通过更好的归纳偏置(inductive bias)实现同等甚至更优的性能。
V3 以 UMA 1/22 的计算成本超越其 Matbench 成绩,有力地证明了在拥有强先验知识的领域(物理对称性),架构创新的 ROI 远高于盲目扩参。这并不否定 scaling law——在缺乏先验知识的通用领域(自然语言、图像),扩参仍然是最可靠的路径。关键在于:当领域知识可以被编码为架构约束时,利用这些知识永远比忽略它们更高效。
3. 与 Google DeepMind GNoME 的战略交汇
Google DeepMind 的 GNoME(Graph Networks for Materials Exploration)项目在 2023 年宣布发现了 220 万种新稳定材料,引发了计算材料学的范式变革。GNoME 的核心 pipeline 依赖于 GNN(图神经网络)做材料稳定性预测,再用 DFT 做验证。
EquiformerV3 在 Matbench Discovery 上的 CPS 0.902 直接超越了 GNoME pipeline 所用模型的精度。如果 DeepMind 将 EquiformerV3 整合到 GNoME 的下一代 pipeline 中,可能的影响包括:
- 更高的筛选精度:减少”假阳性”候选材料,节省 DFT 验证的计算资源。
- 更广的搜索空间:V3 的速度优势允许在相同计算预算内探索更大的化学空间。
- 更可靠的稳定性排序:κ_SRME 的大幅降低意味着模型对稳定性排序的预测更加可信,有助于优先合成最有前景的候选材料。
OC20 排行榜的更新和 DeepMind 的后续动作将是值得密切跟踪的信号。
4. 开源生态与可复现性
论文代码将发布在 github.com/atomicarchitects/equiformer_v3,延续了 MIT atomicarchitects 团队一贯的开源传统。这对等变神经网络领域的健康发展至关重要——不同于 LLM 领域日益封闭的趋势,材料科学 AI 仍保持着高度开放的学术文化:
- 所有 benchmark 数据集(OC20、OMat24、Matbench Discovery)公开可用
- 排行榜接受任何人提交
- 训练代码和模型权重完整开源
这种开放性使得 EquiformerV3 的改进可以迅速被整个社区吸收和验证。
结论与展望
EquiformerV3 是等变神经网络领域的一次代际跃进。SwiGLU-S² 不仅解决了等变特征空间中的非线性表达瓶颈,还展示了一种从 NLP 向科学计算领域系统性移植架构创新的方法论。三大基准同时 SOTA、1.75× 加速、5× 参数压缩——这组数据罕见地同时在精度、效率和紧凑性三个维度取得突破。
Matbench Discovery CPS 首次突破 0.9 标志着 AI 驱动材料发现正在跨越从”辅助工具”到”核心引擎”的门槛。当模型的筛选精度足以让研究者直接信任其推荐结果时,计算材料学的工作流将发生根本性转变——从”人选材料、机器验证”到”机器筛选、人决策验证范围”。
未来值得关注的方向:
- SwiGLU-S² 在蛋白质建模中的应用:AlphaFold3 等模型是否会采用类似的球谐域非线性设计。
- V3 + GNoME 整合:Google DeepMind 是否将 V3 纳入下一代材料发现 pipeline。
- 更高阶球谐的效率优化:当前 V3 在 L_max=4 下表现已极为优异,更高的 L_max 是否能在特定任务(如手性分子区分)上带来额外增益。
- OC20 排行榜后续竞争:Meta FAIR、DeepMind 等机构是否会推出对标 V3 的新架构。