[Paper Review] Taming Multi-Domain, -Fidelity Data: Towards Foundation Models for Atomistic Scale Simulations

Tomoya Shiota, Kenji Ishihara+3 more|arXiv (Cornell University)|Dec 17, 2024

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper introduces Total Energy Alignment (TEA), a method to unify multi-domain, multi-fidelity quantum chemical datasets—spanning organic molecules and inorganic crystals—without costly recalculations. Using TEA, the authors train MACE-Osaka24, a single open-source neural network potential that achieves state-of-the-art accuracy for both molecular reaction barriers and inorganic crystal properties, enabling democratized development of foundation models in atomistic simulations.

ABSTRACT

Machine learning interatomic potentials (MLIPs) are changing atomistic simulations in the field of chemistry and materials science. However, constructing a single universal MLIP that can accurately model molecular and crystalline systems remains challenging. A central obstacle is the integration of diverse datasets generated under different computational conditions. We present Total Energy Alignment (TEA), which is an approach that enables the seamless integration of heterogeneous quantum chemical datasets without redundant calculations. Using TEA, we trained MACE-Osaka24, the first open-source MLIP model based on a unified dataset covering molecular and crystalline systems. This universal model displays strong performances across diverse chemical systems, exhibiting similar or improved accuracies in predicting organic reaction barriers compared to those of specialized models, while effectively maintaining state-of-the-art accuracies for inorganic systems. These advancements pave the way for accelerated discoveries in the fields of chemistry and materials science via genuine foundation models for chemistry.

Motivation & Objective

• To address the challenge of integrating heterogeneous quantum chemical datasets from different computational methods (e.g., DFT functionals, basis sets) used in molecular and crystalline systems.
• To overcome the computational barrier that limits access to universal machine learning interatomic potentials (MLIPs) for researchers without large-scale computational resources.
• To develop a unified, open-source neural network potential capable of accurately modeling both organic molecules and inorganic crystals within a single model.
• To enable the creation of foundation models in chemistry that transcend traditional domain boundaries between molecular and solid-state systems.

Proposed method

• Apply Total Energy Alignment (TEA), a two-step framework: first aligning inner-core reference energies across datasets, then scaling atomization energies to harmonize total energies.
• Use a two-stage calibration: align reference energies of isolated atoms across datasets using a common reference point, followed by scaling atomization energies to correct for functional and basis set differences.
• Train a single MACE-based neural network potential (MACE-Osaka24) on a unified dataset combining MPtrj (inorganic) and OFF23 (organic) datasets, preserving domain-specific accuracy.
• Leverage the MACE architecture’s equivariance and message-passing design to ensure physical consistency across diverse chemical environments.
• Validate the model using standardized benchmarks: reaction barrier predictions, lattice constant optimization, and classical MD simulations with TIP3P and TIP4P/2005 water models.
• Implement ML-driven MD using a modified OpenMM-ML interface with D3(BJ) dispersion corrections to ensure accuracy in long-term simulations.

Experimental results

Research questions

• RQ1 Can heterogeneous quantum chemical datasets from different DFT functionals and basis sets be unified without extensive recalculations?
• RQ2 Can a single, universal machine learning interatomic potential achieve state-of-the-art accuracy for both molecular and crystalline systems?
• RQ3 Does the proposed TEA method enable researchers with limited computational resources to contribute to and benefit from foundation models in atomistic simulations?
• RQ4 How does the performance of a unified model compare to specialized models in predicting reaction barriers and lattice constants across diverse chemical systems?

Key findings

• MACE-Osaka24 achieves mean absolute error (MAE) of 0.0153 eV/atom for inorganic crystal lattice constants and 0.0205 eV/atom for molecular reaction barriers, outperforming specialized models in the latter.
• The model maintains state-of-the-art accuracy for inorganic systems (MAE < 0.02 Å for non-BCC crystals) and achieves comparable or improved performance on organic reaction barriers compared to domain-specific models.
• For BCC crystals (K, Rb, Cs), prediction errors increase with lattice constant due to the 4.5 Å cutoff radius in MACE-Osaka24, which limits the model’s ability to capture long-range interactions in large-unit-cell structures.
• TEA enables seamless integration of MPtrj (inorganic) and OFF23 (organic) datasets without recalculating total energies, reducing computational overhead and democratizing access to unified datasets.
• Classical MD simulations with MACE-Osaka24 using TIP4P/2005 as reference show good agreement in radial distribution functions and diffusion coefficients, confirming long-term stability and accuracy.
• The model’s performance is robust across diverse chemical environments, including water clusters and tripeptides, demonstrating its generalization capability beyond training data.