Technical Summary: Benchmarking Compositional Generalisation for Machine Learning Interatomic Potentials

Problem Statement

Machine Learning Interatomic Potentials (MLIPs) have become essential for computational chemistry and materials science, offering a computationally efficient alternative to Density Functional Theory (DFT) for Molecular Dynamics (MD) simulations. While recent foundation models and architectures (e.g., Graph Neural Networks, Equivariant Transformers) demonstrate high precision on in-distribution (ID) data, their ability to generalise to previously unseen molecules remains unclear. Specifically, it is uncertain whether these models learn the underlying compositional structure of chemistry—understanding how molecular fragments and their combinations determine properties—or if they merely interpolate patterns specific to their training examples. Existing benchmarks often evaluate transfer within the training manifold or focus on broad chemical coverage without isolating specific generalisation mechanisms, failing to address whether models can truly extrapolate to novel molecular structures.

Methodology

The authors introduce the Generalisation for Molecular Dynamics (GMD) benchmark, designed to systematically evaluate the compositional generalisation capabilities of MLIPs. The benchmark consists of four distinct tasks, each constructed to test specific aspects of generalisation (length generalisation and systematicity) using MD trajectories generated via a novel toolkit.

Dataset and Task Design

The dataset comprises 118 molecules (primarily H/C/N/O systems) with approximately 2,000 off-equilibrium snapshots per molecule, labelled using the PBE density functional approximation. The training and test sets are constructed such that they contain trajectories from different molecules, ensuring that generalisation requires understanding chemical principles rather than memorising specific configurations.

The four core tasks are:

1. Fragment Chain Extension: Models are trained on linear alkanes with carbon chain lengths $C_2$–$C_6$ and tested on longer, unseen alkanes ($C_7$–$C_{13}$). This tests length generalisation.
2. Fragment Fusion: Models are trained on alcohols and aldehydes and tested on carboxylic acids. The carboxylic acid functional group is a composition of the alcohol and aldehyde groups. This tests the ability to compose known fragments into novel functional groups.
3. Fragment Duplication: Models are trained on monocarboxylic acids and tested on dicarboxylic acids (two occurrences of the functional group). This tests the ability to generalise from a single motif to multiple occurrences.
4. Fragment Combination: Models are trained on symmetrically functionalised molecules (e.g., diamines or dicarboxylic acids) and tested on asymmetrically functionalised molecules (e.g., amino acids). This tests the ability to recombine distinct functional groups on a familiar scaffold.

Experimental Setup

The benchmark evaluates a diverse set of state-of-the-art MLIPs:

• From-scratch models: SchNet, GemNet, DimeNet++, PaiNN, NequIP, eSCN, EquiFormerV2, and MACE.
• Foundation models: UMA (Small) and MACE-MP-0 (Small, Medium, Large). These were evaluated both zero-shot and after fine-tuning on the specific task training splits (using a fixed budget of 2,000 epochs to establish an upper bound).

Performance is measured using Mean Absolute Error (MAE) for forces (eV/Å) and total energy (eV), alongside auxiliary metrics like cosine similarity and per-atom energy error to disentangle directional and magnitude errors.

Key Contributions

1. The GMD Benchmark: A systematic evaluation framework comprising four controlled tasks that isolate compositional generalisation challenges (length, fusion, duplication, combination) distinct from standard in-manifold transfer.
2. Novel Dataset Generation: A toolkit for generating off-equilibrium MD trajectories with PBE labels, specifically curated to create rigorous train-test splits where training and test molecules differ structurally.
3. Comprehensive Evaluation: A large-scale empirical analysis of both from-scratch and fine-tuned foundation models across these challenging generalisation tasks.
4. Diagnostic Analysis: Detailed error decomposition (magnitude vs. direction) and per-atom analysis revealing that failures are often concentrated in the prediction of force magnitudes and the handling of novel local chemical environments (heteroatoms).

Results

The empirical analysis reveals significant limitations in current MLIPs regarding compositional generalisation:

• Task-Specific Performance:
  • Fragment Chain Extension: Many models (including EquiFormer, GemNet, and MACE) generalise reasonably well to longer chains, with errors often close to in-distribution levels.
  • Fragment Fusion, Duplication, and Combination: These tasks prove highly challenging. For these tasks, out-of-distribution (OOD) errors are frequently one to three orders of magnitude higher than in-distribution errors.
• Foundation Models vs. From-Scratch: Surprisingly, fine-tuned foundation models do not consistently outperform models trained from scratch. In several cases (e.g., Fragment Chain Extension), from-scratch MACE outperformed fine-tuned foundation variants. While foundation models showed slightly better robustness in some OOD scenarios, they did not close the generalisation gap.
• Error Decomposition: Analysis of force errors indicates that while models often preserve the direction of forces (high cosine similarity), they fail significantly in predicting the correct magnitude. This suggests models capture local geometric features but struggle to extrapolate the extensive properties (energy scaling) and force magnitudes for novel compositions.
• Augmented Variants: Experiments with augmented training sets (e.g., providing examples of long chains on different functional groups) showed that increasing coverage can close gaps caused by lack of data (as in Chain Extension). However, providing "worked examples" of composition (e.g., training on amides to help with carboxylic acids in Fragment Fusion) did not significantly reduce OOD errors, suggesting models memorise specific functional groups rather than learning abstract composition rules.

Significance and Claims

The paper claims that the GMD benchmark isolates a generalisation challenge distinct from those addressed by existing benchmarks. The primary finding is that despite the impressive accuracy of MLIPs on standard benchmarks, current models do not inherently deliver compositional generalisation in the evaluated regime.

The authors argue that:

• Pre-training on large, diverse corpora broadens the interpolation manifold but does not, by itself, impart the compositional structure required to handle novel molecular combinations.
• Standard evaluation on in-distribution examples is insufficient for identifying models with genuine physical generalisation capabilities.
• The failure mode is often not a lack of training data for specific elements, but a failure to transfer learned local environments to novel chemical contexts and a failure to decompose or compose functional groups abstractly.

The work aims to encourage a shift towards developing models with better generalisability, highlighting that the "softening" of errors in high-energy, off-equilibrium configurations observed in universal potentials is a symptom of a deeper inability to reason compositionally about chemical structure. The authors provide the dataset and toolkit to support further investigation into this gap.