Source-linked AI summary
Cross Learning between Electronic Structure Theories for Unifying Molecular, Surface, and Inorganic Crystal Foundation Force Fields
Ilyes Batatia, Chen Lin, Joseph Hart, Elliott Kasoar, Alin M. Elena, Sam Walton Norwood, Thomas Wolf, Gábor Csányi
TL;DR
Accurate, transferable interatomic potentials remain difficult to unify across chemical domains. This work combines MACE architecture changes with multi-head replay post-training across diverse datasets. The resulting cross-domain learning improves molecular, surface, and crystal benchmarks without degrading materials accuracy.
Problem
Existing MLIPs are fragmented by chemical domain, and foundation models often have limited scope, with datasets and architectures favoring domain-specific optimization over transferability.
Source: §I · p. 1Method
The approach modifies MACE to increase element weight sharing and nonlinear tensor-decomposition factors, then uses multi-head replay post-training across diverse datasets.
Source: §I · pp. 1–2Results
Cross-domain learning improves molecular, surface, and crystal benchmarks without degrading materials accuracy.
Source: §I · p. 2Takeaways & Limitations
The demonstrated principles provide guidance toward foundation MLIPs intended to handle molecular, surface, and materials chemistry with a single model.
Source: §XIV · p. 14
Abstract
from arXiv · showhide
Creating a single unified interatomic potential capable of attaining ab initio accuracy across all chemistry remains a long-standing challenge in computational chemistry and materials science. This work introduces a training protocol for foundation machine-learning interatomic potentials (MLIPs) that bridge molecular, surface, and materials chemistry through cross-domain learning. First, we introduce enhancements to the MACE architecture that improve its performance on chemically diverse databases by increasing weight sharing across chemical elements and introducing non-linear factors into the tensor decomposition of the product basis. Second, we develop a multi-head replay post-training methodology that enables efficient knowledge transfer across diverse chemical domains. By fine-tuning on datasets at different levels of electronic structure theory, including inorganic crystals, molecular systems, surface chemistry, and reactive organic chemistry, we demonstrate that a single unified model achieves state-of-the-art performance across several chemical domains. Comprehensive benchmarking reveals superior cross-domain transferability compared with existing specialised and multi-task models, with notable improvements in molecular and surface properties while maintaining state-of-the-art performance in materials-property prediction.
I. Introduction
Existing MLIPs are fragmented across chemical domains, while foundation models often have limited scope. This work combines MACE enhancements with multi-head cross-learning to build a unified potential for molecular, surface, and materials chemistry.
I. Introduction
- Traditional models are divided among molecular, surface, and bulk-material domains, while foundation models often excel in one domain but perform inadequately in others.The paper identifies dataset constraints and domain-specific architectural choices as contributing factors.Source: §I · p. 1
- The proposed multi-head approach shares chemical and geometrical representations across theory-specific readouts, with cross-learning directed toward a main PBE head.The goal is one continuous potential applicable across chemical contexts.Source: §I · p. 1, A. Theoretical Foundation · p. 2
- The architecture changes increase weight sharing across elements and add nonlinear tensor-decomposition factors, which the authors report improve performance on chemically diverse databases.Source: §I · pp. 1–2
- Cross-domain learning yields measurable generalization across molecular, surface, and crystal benchmarks without degrading materials accuracy.Source: §I · p. 2
A. Theoretical Foundation
- MACE maps atomic coordinates and atomic numbers to total potential energy by summing atomic energies built from invariant descriptors of each atom’s many-body neighborhood.Each atomic energy depends on neighbors within a cutoff radius, their relative vectors, and their atomic numbers.Source: A. Theoretical Foundation · p. 2
B. MACE with non-linear tensor decomposition
The MACE modifications target chemically diverse datasets through greater element sharing and nonlinear features, while retaining the architecture’s many-body message-passing framework.
B. MACE with non-linear tensor decomposition
- Equations (3)–(19) give the complete MACE functional form, with highlighted parts marking this work’s modifications selected through large-dataset grid searches.The authors searched for accuracy with as few changes as possible.Source: B. MACE with non-linear tensor… · p. 2
- The architecture also uses learnable element embeddings for scalar initial node features, with higher-order components initialized to zero.Source: B. MACE with non-linear tensor… · p. 2
- Element-specific source and target embeddings are concatenated with radial Bessel features before the radial MLP, and applying the cutoff outside that MLP enforces smoother decay near the cutoff.Source: B. MACE with non-linear tensor… · p. 3
- The one-particle basis sums over neighborhood atoms regardless of element identity, and the tensor product on edges is identified as MACE’s computational bottleneck.The implementation uses fewer edge channels than the channel limit, with more for the inexpensive first layer.Source: B. MACE with non-linear tensor… · p. 3
- A learnable residual connection adds central-atom information to the neighborhood basis, allowing richer many-body features that depend on the central atom as well as its neighbors.Source: B. MACE with non-linear tensor… · p. 4
- The nonlinear gate permits nonlinear rank-1 factors in the tensor decomposition, while element-agnostic weights are used for message construction and updates.Source: B. MACE with non-linear tensor… · p. 4
III. Multi-Head Replay Post-Training
Multi-Head Replay Post-Training first builds shared representations from inorganic crystals, then fine-tunes domain- and theory-specific heads while replaying foundation data.
III. Multi-Head Replay Post-Training
- Stage 1 pre-trains a unified backbone on a large, diverse inorganic-crystal dataset to capture shared chemical and geometrical patterns.Source: §III · p. 4
- Stage 2 simultaneously fine-tunes shallow heads for different domains or theory levels and continues updating the backbone weights.Source: §III · p. 4
- A replay buffer samples configurations from the original pre-training data, and mixed minibatches combine these with new domain-specific samples to retain foundational chemical knowledge.Source: §III · p. 4
IV. Datasets
The dataset strategy combines large-scale foundation data with targeted fine-tuning data across inorganic crystals, surfaces, and molecular chemistry.
IV. Datasets
- The datasets combine large-scale foundation data with targeted fine-tuning to cover inorganic crystals, surfaces, and molecular chemistry.Source: §IV · p. 4
A. Pre-training Dataset
The workflow pre-trains on a large inorganic-crystal dataset, then fine-tunes across chemical domains using replay. The model descriptions distinguish the proposed multi-head model from its pretrained, single-domain, and comparison baselines.
A. Pre-training Dataset
- OMAT supplies 100 million inorganic-crystal configurations spanning 89 elements at PBE and PBE+U levels for pre-training.The OMAT replay buffer uses a randomly selected 10 million configurations to retain foundational inorganic structures during fine-tuning.Source: A. Pre-training Dataset · p. 4, B. Fine-tuning Datasets · p. 4
B. Fine-tuning Datasets
- Fine-tuning combines OMAT replay with RGD1 reactive-organic data, MPTraj, OC20 surfaces, OMOL configurations, and r2SCAN inorganic-crystal snapshots.The datasets span different electronic-structure levels, including B3LYP, PBE, hybrid DFT, and r2SCAN.Source: B. Fine-tuning Datasets · pp. 4–5
V. Model Descriptions
- The model set includes the proposed multi-head replay model, OMAT and OMOL MACE models, UMA, and specialized or foundation-model comparators.The proposed mace-mh-1 is fine-tuned on the mace-omat-1 backbone; its OMAT head is used for main-text results, while the OMOL head uses 1% of the full OMOL dataset.Source: §V · pp. 5–6
VI. Benchmark Overview
The overview compares foundation potentials across five weighted domains and finds mace-mh-1-omat the top overall model with a balanced profile. Its gains in molecular and surface categories accompany stable materials performance.
VI. Benchmark Overview
- The five-domain aggregation uses normalized per-metric means and a weighted global score, with weights the authors describe as subjective.The authors state that the weighting needs improvement through community discussion.Source: §VI · p. 6, Figure 2 · p. 7
VII. Dispersion corrections (D3)
- The benchmark overview includes PBE and PBE+U evaluations, with D3(BJ) corrections applied to PBE-trained models for systems with significant van der Waals interactions.The OMOL models trained at ωB97M-VV10 were evaluated without additional dispersion correction.Source: §VII · p. 6
VIII. Materials Benchmarks
- Materials benchmarks test elastic moduli, thermal conductivity, and phonon spectra, extending evaluation beyond energy and force errors to potential-energy-surface curvature.Most reference data are computed at PBE+U, matching MPtraj and OMAT specifications.Source: §VIII · p. 6
VI. Benchmark Overview
- Mace-mh-1-omat has the top global score and most balanced five-domain profile, with strong molecular and molecular-crystal scores while remaining competitive elsewhere.The global score weights Materials 0.25, Molecules 0.25, Surfaces 0.20, Molecular Crystals 0.20, and Physicality 0.10.Source: §VI · p. 6, Figure 2 · p. 7
A. Elastic Moduli
- Mace-mh-1-omat has the lowest shear-modulus error, while orb-v3-consv-inf-omat has the lowest bulk-modulus error; the multi-head model remains competitive with OMAT.The comparison uses Materials Project elasticity references.Source: A. Elastic Moduli · p. 8, Table I · p. 8
B. Thermal Conductivity
The materials evaluation reports strong thermal-conductivity and phonon results for the multi-head model, including the best overall dynamical-stability classification. Molecular-crystal formation energies probe intermolecular interactions and solid cohesion.
B. Thermal Conductivity
- On the 103-compound Wigner-transport benchmark, uma-m-1p1-omat achieves the lowest κRMSE, closely followed by mace-omat-1 and uma-s-1p1-omat; mace-mh-1-omat maintains strong performance compared with the baseline.Source: B. Thermal Conductivity · p. 8
C. Phonons
- Mace-mh-1-omat achieves the lowest errors for most phonon and thermodynamic properties and the best overall dynamical-stability classification.It ranks behind other models for CV, ωmax, and entropy, while maintaining or slightly outperforming mace-omat-1.Source: C. Phonons · pp. 8–9, Figure 3 · p. 9
- Mace-mh-1 shows stronger cross-learning than uma-s-1p1 on phonon benchmarks, outperforming it in each domain and surpassing specialized mace-mp-0a.The uma-m-1p1-omol model was excluded because it could not geometry-optimize many structures.Source: C. Phonons · p. 9
IX. Molecular Crystal Benchmarks
- Molecular-crystal formation energies test intermolecular interactions and the cohesive assembly of molecular solids.Source: §IX · p. 9
A. X23 Molecular Crystals
The X23 and water-ice benchmarks show strong formation-energy performance from mace-mh-1-omat-D3, while the surface passages describe the scope of adsorption tests.
A. X23 Molecular Crystals
- mace-mh-1-omat-D3 achieves the lowest MAE among tested models on X23 molecular-crystal formation energies.It outperforms a model trained on a larger molecular-crystal dataset; the authors attribute the gain over mace-omat-1-D3 to molecular-chemistry transfer from OMOL and SPICE heads.Source: A. X23 Molecular Crystals · p. 9, Table IV · p. 9
B. DMC Water Ice Polymorphs
- mace-mh-1-omat-D3 achieves the lowest MAE among tested models for ice-polymorph relative energies.The authors describe this as improved modeling of hydrogen-bonding interactions.Source: B. DMC Water Ice Polymorphs · p. 10, Table V · p. 10
X. Surface Benchmarks
- Surface benchmarks assess adsorption energetics and site-specific binding relevant to catalysis and gas–surface processes.The S24 set covers 24 adsorption systems spanning covalent, ionic, metallic, and porous surfaces.Source: §X · p. 10, A. General Surface Adsorption: S24… · p. 10
B. OC20 Metal Surface Benchmark
Across molecular and surface tests, the reported results include improved adsorption and molecular-property predictions, with the GMTKN55 model performing comparably to specialized molecular models.
B. OC20 Metal Surface Benchmark
- 0.097 eV is the MAE for uma-m-1p1-omat-D3 on the OC20 metal-surface benchmark, compared with 0.138 eV for mace-mh-1-omat-D3.The UMA models were trained on 100 times more surface configurations; the authors also report better cross-learning by the MACE model on OMOL tasks.Source: B. OC20 Metal Surface Benchmark · p. 10, Table VII · p. 11
XI. Molecular Benchmarks
- The molecular benchmarks cover conformer energies, reaction energies, and noncovalent interactions to assess prediction of key chemical properties governing molecular behavior.The section's reported molecular tests include strained conformers, main-group chemistry, and peptide fragments.Source: §XI · p. 11, A. Wiggle 150 · p. 11, B. GMTKN55 Main Group Chemistry · p. 11
A. Wiggle 150
- 4.80 kcal mol−1 is the MAE for mace-mh-1-omat-D3 on Wiggle150, comparable to the underlying PBE-D3 functional performance.Multi-head fine-tuning improves by a factor of 2 over the mace-omat-1-D3 pretrained baseline.Source: A. Wiggle 150 · p. 11, Table VIII · p. 11
B. GMTKN55 Main Group Chemistry
- mace-mh-1-omat-D3 significantly outperforms other PBE-trained foundational models; uma-m-1-p1-omat-D3 records a GMTKN55 neutral-singlet weighted mean absolute error of 11.23 kcal mol−1, comparable to PBE-D3(BJ).mace-mh-1-omat-D3 also performs comparably to specialized molecular models AIMNet2 and MACE-OFF23(L), trained on ωB97M-D3BJ.Source: Table IX · p. 11, B. GMTKN55 Main Group Chemistry · p. 12
C. Protein fragments: PLF547
- 0.626 kcal mol−1 is the lowest MAE for mace-mh-1-omat-D3 on neutral PLF547 peptide-fragment interaction energies.The authors report substantial improvement over the OMAT-only backbone.Source: C. Protein fragments: PLF547 · p. 12, Table X · p. 12
D. S30L Molecular complexes
The S30L results favor mace-mh-1-omat-D3 over its pretrained baseline, while physicality tests reveal model differences in extensivity, additivity, and diatomic-curve smoothness.
D. S30L Molecular complexes
- For S30L [70] host-guest binding-energy benchmarking, mace-mh-1-omat-D3 achieves the best MAE of 10.13 kcal mol−1, outperforming pre-trained mace-omat-1-D3 at 15.35 kcal mol−1, compared with experimental references.Source: D. S30L Molecular complexes · p. 12, Table XI · p. 12
XII. Physicality Benchmarks
- Physicality tests assess size extensivity, additivity, and smoothness of pair interactions because violations can yield unphysical simulations.The slab test checks that separated system energies sum to the isolated energies.Source: §XII · p. 12, A. Size Extensivity and Locality · p. 13
A. Size Extensivity and Locality
- UMA models and ORB models show large energy deviations in the slab test and spurious forces in the distant-hydrogen test.The reported maximum spurious force for uma-s-1p1-omat is 969.2 meV/Å; the authors associate the deviations with unphysical non-local interactions.Source: Table XII · p. 13, A. Size Extensivity and Locality · p. 13, Table XIII · p. 13
B. Homonuclear and Heteronuclear Diatomics
- MACE models show good diatomic-curve behavior with minimal force flips and energy minima, while UMA models show numerous flips, minima, and inflection points.The assessment evaluates curve smoothness and monotonicity across compatible homonuclear and heteronuclear diatomics.Source: Table XIV · p. 13, B. Homonuclear and Heteronuclear… · pp. 13–14
XIII. Computational efficiency
In the tested setup, the models reach competitive idealized speed and run at 1.4–2.2 million MD steps per day in real simulations. The discussion links cross-domain transfer to shared representations and identifies further domains and capabilities for future work.
XIII. Computational efficiency
- 43 steps per second (around 3.8 Megasteps/day) is achieved in a best-case scenario that neglects MD overheads.Speed is assessed for energy and force calculations on 1,000-atom carbon FCC structures; the authors note that MLIP speed comparisons depend on the setup.Source: §XIII · p. 14, Table XV · p. 14
- 1.4–2.2 million steps per day in real simulations is about 2.5 times slower than idealized timing, mainly without torch.compile and partly due to graph-construction and LAMMPS-update overheads.These measurements use the LAMMPS MLIAP KOKKOS interface with cuEquivariance in float32, but not torch.compile.Source: §XIII · p. 14, Table XVI · p. 14
XIV. Discussion and Outlook
- Molecular knowledge transfers to the OMAT head when SPICE and OMOL are included, through improved descriptions of local environments and coordination patterns.The multi-head architecture shares representations while maintaining consistency across electronic-structure levels; the authors report that UMA’s more flexible global level-of-theory embedding shows less transfer.Source: §XIV · p. 14
- The authors identify solid/liquid interfaces, molecular crystals, and amorphous systems as additional domains for future extensions.They also propose uncertainty quantification, hybrid learning with experimental data, and extensions to charged and magnetic systems.Source: §XIV · pp. 14–15
Conflict of Interest
- The paper reports that GC has commercial force-field licensing and equity interests, while SWN has a financial interest and equity stake in Mirror Physics.Source: Conflict of Interest · p. 15
Data and Models Availability
- The MACE code, example scripts, pretrained models, and public training datasets are identified as available; processed benchmark data and analysis scripts are planned for a forthcoming ML-PEG publication.The passage also identifies cuEquivariance kernels as available.Source: Data and Models Availability · p. 15
A. Re-estimation of atomic reference energies (E0s)
- Fine-tuning adapts atomic reference energies to normalize target energies as atomization energies, and the paper proposes a linear-system approach for more robust re-estimation.The proposed corrections are fitted across configurations using least squares because the system has more equations than unknowns; replay-head reference energies remain fixed to their pre-training DFT values.Source: A. Re-estimation of atomic reference… · p. 19
XVI. Hyperparameters
The section specifies the loss terms, optimizer, and pre-training and fine-tuning schedules used for the MACE variants.
XVI. Hyperparameters
- Training uses a weighted sum of Huber losses for energy, forces, and stress, with weights (1, 10, 10) and deltas of 0.01.The force Huber delta adapts to force magnitude, decreasing step-wise from 1.0 to 0.1 as atomic force increases from 0 to 300 eV/Å.Source: §XVI · p. 19
- The models use AMSGrad, an Adam variant, with β1 = 0.9, β2 = 0.999, and ϵ = 10^-8.Source: §XVI · p. 19
- Pre-training runs for 7 epochs and fine-tuning for 20 epochs, both at learning rate 0.001 with EMA decay 0.999 and gradient clipping of 100.Both phases use 48 NVIDIA H100 GPUs across 12 nodes.Source: §XVI · p. 19