Benchmarking Compositional Generalisation for Machine Learning Interatomic Potentials

会议: NeurIPS 2026
arXiv: 2605.08988
代码: https://github.com/Nourollah/GMD-26
领域: 计算生物/AI4Science（机器学习原子间势函数）
关键词: 机器学习原子间势、组合泛化、分子动力学模拟、分布外评估、基准测试

一句话总结

提出了系统评估机器学习原子间势（MLIP）组合泛化能力的基准测试 GMD，覆盖链延伸、官能团融合、复制、重组与裂解五类任务，揭示了主流模型及微调基础模型在跨分子组合外推时能量与受力误差普遍恶化 1 到 3 个数量级且主要源自标量受力幅值预测崩溃的现象。

研究背景与动机

机器学习原子间势（MLIP）在计算化学、材料科学和药物发现中发挥着越来越关键的作用，它通过图神经网络（GNN）或等变 Transformer 逼近高精度的密度泛函理论（DFT）势能面，以低计算开销驱动长时间跨度的分子动力学（MD）模拟。为了摆脱传统高精度参考计算高昂的数据生成成本，学术界与工业界近期广泛转向在数百万分子结构上预训练的通用基础势模型（如 MACE-MP-0、UMA），期待模型通过海量数据与等变归纳偏置掌握底层的量子相互作用规律。然而，现实中的化学空间近乎无限，新药设计或材料筛选天然要求模型能够可靠地外推到训练集未曾包含的新型分子体系。

现有的原子势基准测试（如 MD17、WS22、Transition1x 等）大多集中在平衡态或同质化学构型采样上，对基础模型的测试也普遍停留在预训练流形内的内插分布（In-Distribution, ID）评估，难以回答核心科学问题：模型究竟是真正掌握了化学规律的模块化组合性（即如何由已知分子碎片及其相互作用法则推导出未见组合的物理性质），还是仅仅在已有训练样本构成的紧致流形内进行高维内插？如果缺乏这种组合泛化能力，当模型用于生成具有全新母核结构或多官能团修饰的分子时，势函数推断的能量与受力就会严重失真，导致分子动力学积分漂移甚至崩溃。

为了从机制上严格解构原子间势模型的组合外推极限，本文构建了分子动力学泛化基准 GMD（Generalisation for Molecular Dynamics）。研究通过精细设计的一维取代烷基碳链骨架，隔离出长度伸展与四种不同形式的官能团符号重组任务，确保训练数据完全涵盖理解测试分子所需的全部基本化学碎片。本文的核心 idea 是：构建一套以化学积木块（Fragments）及其显式运算法则为核心的控制性组合泛化基准，系统解耦并严格评估 MLIP 在分子结构外推与官能团组合时的泛化瓶颈与物理失效模式。

方法详解

整体框架

GMD 基准构建的核心逻辑是“最小化学位移下隔离单一组合操作”：选择化学结构相对清晰的取代烷基碳链作为通用支架，使测试集与训练集之间的差异严格来自于特定维度的组合操作，而非复杂的非局域构象混淆。整个基准的数据生成与评测由一套基于 ASE（Atomic Simulation Environment）的自动化流水线驱动。流水线首先从分子 SMILES 生成初始 3D 构象，随后利用 FlashMD（加载 pet-omatpes 检查点）驱动长步长非平衡态分子动力学采样（300 K 下积分 20 万步，每 100 步记录一次快照），最后调用量子化学软件 ORCA 在 PBE/def2-TZVP+D3BJ 水平下重新计算所有快照的能量与受力真值。

GMD 共包含 118 个分子、超过 29.6 万个高精度标注构象，划分为五个代表性的组合泛化任务： 1. 片段链延伸（Fragment Chain Extension）：检验长度泛化（Productivity）。训练集包含碳链长度为 C2–C6 的直链烷基链，测试集为未见过的 C7–C13 较长烷基链。 2. 片段融合（Fragment Fusion）：检验新型复合官能团的合成认知（Systematicity）。训练集由醇类（-OH）、醛类（C=O）以及复杂的多元醇与多元羰基化合物组成，测试集为羧酸分子（-COOH）。羧酸官能团在结构上可视为羰基与羟基结合在同一个碳原子上的复合结构。 3. 片段复制（Fragment Duplication）：检验化学基序的重数泛化（Systematicity）。训练集包含单羧酸（C5–C10），测试集则是分子骨架对称两端各含一个羧酸基团的二元羧酸（C5–C10）。 4. 片段重组（Fragment Combination）：检验异相非对称官能团组合（Systematicity）。训练集由对称二胺（含两个 -NH2）和对称二元羧酸（含两个 -COOH）组成，测试集则是骨架长度相同但两端分别连接一个羧酸和一个氨基的氨基酸分子。 5. 片段裂解（Fragment Fission）：检验复合官能团向孤立组成碎片的逆向分解能力（Systematicity）。复用片段复制任务中单羧酸训练出的模型，测试其在孤立醇类和醛类分子上的预测表现。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["分子规范输入<br/>SMILES 字符串"] --> B["FlashMD 非平衡态轨迹采样<br/>300K, 2×10⁵ 步, 16fs 步长"]
    B --> C["ORCA 高精度 DFT 标签重算<br/>PBE/def2-TZVP+D3BJ"]
    C --> D["GMD 五大组合泛化任务切分"]
    D --> E["五大任务评估：链延伸 / 融合 / 复制 / 重组 / 裂解"]
    E --> F["物理失效解耦：受力方向与幅值误差分析"]

关键设计

1. 严格受控的五大化学组合泛化任务设计

为了避免现有数据集因训练测试集化学空间混杂而无法定位模型缺陷的问题，GMD 精心构造了五条组合测试线。每个任务中测试分子的所有构成积木（如碳骨架、羟基、羰基、氨基）均已在训练集中充分暴露，从物理法则上看，如果模型真正学到了局域化学环境的相互作用势能叠加与组合法则，对测试分子的泛化在理论上是完全成立的。这种设计消除了“模型未见过某种化学元素”或“训练数据量过小”的借口，迫使模型在已知成分的新颖排列与尺度下展示外推能力。

2. 基础模型全参数微调与原子参考能量基准校准

基础模型（如 MACE-MP-0、UMA）在预训练时通常使用平面波基组，且各自预设了特定的孤立原子参考能量 \(E_0\)。直接对其进行零样本（Zero-shot）总能量测试会导致数千电子伏特的系统性基准偏差。为此，基准建立了一套公平的微调协议：在微调前，严格基于 ORCA 高斯基组的参考数据重新估计各元素的孤立原子参考能 \(E_0\)，消除常数偏移；随后在每个任务的训练集上开展 2000 个 epoch 的全参数微调（不使用参数高效微调或冻结层），探索基础模型在该基准上可达到的性能上限。

3. 受力矢量解耦与局域化学环境物理归因机制

模型在宏观受力均方根误差（Forces MAE）上的上升可能来自两个完全不同的物理失真来源：受力方向错误（向量角度偏离）或受力强度错误（向量模长缩放）。GMD 在标量 MAE 之外引入了方向余弦相似度（Cosine Similarity）与受力幅值误差（\(\text{MAE}_{\text{mag}}\)）： $\(\text{Cosine Similarity} = \frac{1}{N}\sum_{i=1}^N \frac{\hat{\mathbf{F}}_i \cdot \mathbf{F}_i}{\|\hat{\mathbf{F}}_i\|_2 \|\mathbf{F}_i\|_2}, \quad \text{MAE}_{\text{mag}} = \frac{1}{N}\sum_{i=1}^N \left| \|\hat{\mathbf{F}}_i\|_2 - \|\mathbf{F}_i\|_2 \right|\)$ 同时配合逐原子残差投影分析，精确追踪受力与能量误差究竟发生在非极性碳链骨架上，还是特异性地局域在未见过的杂原子基团（如 -COOH 的氧原子）上。

损失函数 / 训练策略

从头训练模型（From-scratch）基于对应官方库（fairchem-core、mace-torch、nequip）的标准联合损失函数进行优化： $\(\mathcal{L} = \lambda_E |E - \hat{E}| + \lambda_F \frac{1}{3N}\sum_{i,c} |F_{i,c} - \hat{F}_{i,c}|\)$ 超参数通过贝叶斯优化（W&B sweeps）在 ID 验证集上严格调优（包括隐藏层维度、径向基截断半径、学习率衰减等，每模型 90~180 次试验）。基础模型采用 AdamW 优化器，学习率 \(10^{-4}\)，无层冻结，微调 2000 个 epoch。

实验关键数据

主实验

评估覆盖了各类代表性架构：不变消息传递网络（SchNet）、显式几何特征网络（PaiNN、GemNet、DimeNet++）、高阶旋转等变网络（NequIP、eSCN、MACE）、等变 Transformer（EquiFormerV2），以及微调的大规模基础模型（MACE-Large、MACE-Medium、UMA-Small）。

下表列出各主流模型在 GMD 五大任务上的测试集受力 MAE（F MAE，单位 eV/Å）与总能量 MAE（E MAE，单位 eV）表现：

任务	模型	F MAE (ID)	F MAE (OOD)	E MAE (ID)	E MAE (OOD)
Fragment Chain Extension	SchNet	0.0345	1.6836	1.0622	82.638
	PaiNN	0.0088	0.0682	0.2604	35.113
	GemNet	0.0037	0.0133	0.2894	223.20
	EquiFormerV2	0.0015	0.0040	0.4155	407.78
	DimeNet++	0.0030	1.4118	0.0933	354.66
	NequIP	0.0151	0.0278	0.0108	0.0576
	eSCN	0.0026	0.0312	0.1955	151.45
	MACE (scratch)	0.0034	0.0061	0.0016	0.0087
	MACE-Large (FT)	0.0052	0.0089	0.0084	0.0246
	UMA-Small-2000 (FT)	0.0201	0.0188	0.0422	0.1117
Fragment Fusion	SchNet	2.9222	9.0294	142.64	355.36
	PaiNN	0.0159	0.1183	12.964	7.5822
	GemNet	0.0247	0.1218	4.8028	16.759
	EquiFormerV2	0.0082	0.1385	7719.1	6721.4
	DimeNet++	0.0873	0.3106	151.12	567.24
	NequIP	0.0388	0.2286	0.1063	1.4876
	eSCN	0.0557	0.1415	44.785	315.44
	MACE (scratch)	0.0047	0.0889	0.0036	0.9734
	MACE-Large (FT)	0.0098	0.0648	0.0076	1.0584
	UMA-Small-2000 (FT)	0.0150	0.0659	0.0547	0.7841
Fragment Duplication	SchNet	0.4452	1.1944	9.2223	3928.8
	PaiNN	0.0246	0.0843	9.1052	3947.7
	GemNet	0.0101	0.0605	4.6360	1560.6
	EquiFormerV2	0.0079	0.0544	6780.8	10840.0
	DimeNet++	0.0619	1.3287	27.959	684.33
	NequIP	0.0294	0.1735	0.0311	715.34
	eSCN	0.0224	0.0711	10.382	3963.0
	MACE (scratch)	0.0049	0.0390	0.0029	2.6409
	MACE-Large (FT)	0.0084	0.0320	0.0045	0.7266
	UMA-Small-2000 (FT)	0.0986	0.1157	0.1544	822.51
Fragment Combination	SchNet	0.4325	79.734	11.087	1098.8
	PaiNN	0.0265	0.0605	5.0876	42.216
	GemNet	0.0158	0.0304	3.6094	5.2326
	EquiFormerV2	0.0066	0.0158	4.5004	46.744
	DimeNet++	0.0257	0.0331	43.532	43.504
	NequIP	0.0434	0.1756	0.0502	0.1128
	eSCN	0.0203	0.0340	33.211	62.500
	MACE (scratch)	0.0053	0.0133	0.0029	0.0076
	MACE-Large (FT)	0.0102	0.0196	0.0053	0.0177
	UMA-Small-2000 (FT)	0.0215	0.0222	0.0439	0.0358
Fragment Fission	SchNet	0.4452	0.8676	9.2223	1935.6
	PaiNN	0.0246	0.1557	9.1052	1944.0
	GemNet	0.0101	0.1924	4.6360	767.60
	EquiFormerV2	0.0079	0.0817	6780.8	4542.0
	DimeNet++	0.0619	3.6793	27.959	690.66
	NequIP	0.0294	7.0260	0.0311	391.68
	eSCN	0.0224	0.1016	10.382	1972.4
	MACE (scratch)	0.0049	0.1886	0.0029	1.4918
	MACE-Large (FT)	0.0084	0.0748	0.0045	0.2635
	UMA-Small-2000 (FT)	0.0986	0.1182	0.1544	411.82

消融实验：增强任务设计与微调预算分析

论文进一步探讨了数据增强（Augmented Tasks）以及基础模型微调轮数对泛化瓶颈的缓解程度。在增强任务中，训练集补充了跨官能团的完整链长覆盖（Augmented Chain Extension）或引入了类似组合示例酰胺（Augmented Fragment Fusion）。

下表展示了增强任务下代表性模型的表现，以及 MACE-MP-0 Small 在不同微调轮数下的误差变化：

实验设置	模型 / 配置	F MAE (ID)	F MAE (OOD)	E MAE (ID)	E MAE (OOD)	说明
Augmented Chain Ext.	SchNet	0.0410	1.8320	1.1500	95.400	数据增强后仍严重发散
	DimeNet++	0.0035	0.0078	0.0850	0.1250	补充长度覆盖后跨官能团外推成功
	MACE (scratch)	0.0040	0.0052	0.0021	0.0048	跨官能团长度迁移完美收敛
Augmented Fusion	MACE (scratch)	0.0051	0.0821	0.0040	0.9120	虽见酰胺前例，仍无法泛化至羧酸
	MACE-Large (FT)	0.0105	0.0612	0.0081	0.9850	基础模型同样未掌握组合规则迁移
微调轮数消融	MACE-Small (25 ep)	0.0518	0.1305	0.0621	2.9924	轮数少时接近 Zero-shot 水平
(Fragment Fusion)	MACE-Small (110 ep)	0.0359	0.1122	0.1454	2.5275	误差显著下降但 OOD 仍偏大
	MACE-Small (2000 ep)	0.0191	0.1252	0.0142	1.6275	ID 误差大幅下降，OOD 差距扩大至 100x

关键发现

• 标量受力幅值预测崩溃是泛化的核心物理短板：对受力分解指标的分析表明，在 OOD 场景下，所有模型的受力方向余弦相似度均维持在极高水平（除极少数坏例均大于 0.94），但受力幅值误差 \(\text{MAE}_{\text{mag}}\) 剧烈上升。模型能够正确辨识原子受力的几何指向，却无法正确推断作用力的大小强度。
• 微调基础模型未能突破从头训练上限：即使经过 2000 个 epoch 的充分全参数微调，参数量达数千万的基础模型（如 MACE-MP-0 Large、UMA）在大多数组合任务上的 OOD 表现并未超越从头训练的 MACE 基线。预训练所积累的化学多样性流形仅扩展了其内插覆盖范围，并未赋予模型真正的符号化组合泛化逻辑。
• 局域性与非局域干扰的物理不对称性：在片段复制任务中，当分子两端同时挂载两个原本学过的羧酸基团时，模型不仅在新增的羧酸上预测不准，连原本单取代时预测极其精准的已知端氧原子的受力也遭到破坏。这说明当前消息传递架构无法维持长程解耦的不变性，异地基团引入了伪长程干扰。

亮点与洞察

• 提出首个解耦原子势组合泛化的控制性基准：打破了过去单纯扩充构象数量或盲目追求复杂大分子的评测惯例，用最简洁的碳链取代物严格隔离了长度延展、基团重合、倍增与异相配对，为 AI4Science 建立了严谨的泛化诊断工具。
• 揭示了“ID 表现无法预言 OOD 可靠性”的评测陷阱：在内分布测试中表现顶尖的等变架构（如 EquiFormerV2 在训练分子上 F MAE 达到惊人的 0.0015 eV/Å），在外推到新组合结构时能量误差甚至出现数百至数千电子伏特的灾难性爆炸，表明现有的单分布排行榜完全无法选拔出具备物理稳健性的势函数。
• 明确了数据覆盖缺陷与组合操作学习失败的本质分野：增强任务实验极其漂亮地证明了，当缺失的是单一维度的数据覆盖时（如链长度），提供跨基团的参照即可闭合泛化差；但当缺失的是组合操作时（如羟基与羰基合并成羧基），即使给模型展示类似的操作（胺基与羰基合并成酰胺），模型也无法抽象出组合法则。

局限与展望

• 化学元素空间与环境范围有限：当前基准仅涵盖真空孤立态下的 H/C/N/O 有机小分子，未拓展至包含金属配位、周期性晶体、溶液相或荷电离子体系，在更复杂的无机化学材料中的组合外推表现仍待检验。
• 从头微调成本极高：为了测定基础模型的能力上限，基准采用了 2000 个 epoch 的全参数密集微调（耗费数卡 GH200 连续训练数天），普通下游研究者通常只进行数十轮轻量适配，实际部署时的外推表现可能比文中所述更为脆弱。
• 改进方向：探索在 GNN/Transformer 势函数中显式融入物理先验的分解架构（如局部多体相互作用的符号化硬约束），或引入针对基序（motif）的可组合图生成与表示学习机制，从架构层面根治受力幅值外推漂移。

相关工作与启发

• vs MD17 / WS22 / MD22: 传统分子动力学基准的训练与测试轨迹来自完全相同的分子系统，仅测试平衡态热运动内插；GMD 明确要求在分子层面进行跨体系的 OOD 组合泛化。
• vs BOOM: BOOM 侧重于分子宏观物理化学属性数值上的极值外推（Tail-end extrapolation）；GMD 则聚焦于原子级势能面与力场（Forces/Energy）的构型及亚基组合泛化。
• vs 预训练基础模型评测（如 MACE-MP-0 / UMA）: 先前评测多在包含数十万种材料的宏观测试集上报整体平均指标，被流形内插所掩盖；GMD 证明海量数据无法自动涌现出基序组合能力，指明了基础模型下一代架构演进的方向。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 首次系统性解耦并量化机器学习原子间势在化学微环境中的五大组合泛化能力。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 8 种从头训练模型、多个尺度基础模型、2000 epoch 微调消融与逐原子受力矢量分解。
• 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨清晰，实验设计环环相扣，物理归因透彻深入。
• 价值: ⭐⭐⭐⭐⭐ 为分子势函数领域从“纯拟合内插”迈向“真正物理泛化”提供了里程碑式的基准工具。

相关论文

• [ICLR 2026] DistMLIP: A Distributed Inference Platform for Machine Learning Interatomic Potentials
• [ICLR 2026] PepBenchmark: A Standardized Benchmark for Peptide Machine Learning
• [ICLR 2026] Discrete Compositional Generation via General Soft Operators and Robust Reinforcement Learning
• [ICLR 2026] Representing Local Protein Environments with Machine Learning Force Fields
• [NeurIPS 2026] Branching Flows: Discrete, Continuous, and Manifold Flow Matching with Splits and Deletions