🔬 applied physics

Fast and Accurate Foundation Models for Equivariant Machine-Learned Interatomic Potentials

This paper presents a family of fast, scalable, and accurate equivariant foundation models for interatomic potentials that overcome speed-accuracy trade-offs through optimized NequIP and Allegro architectures, while emphasizing that future improvements in materials discovery rely more on dataset diversity and consistent transition metal descriptions than on model complexity alone.

Original authors: Seán R. Kavanagh, Chuin Wei Tan, Menghang Wang, Marc L. Descoteaux, Gabriel de Miranda Nascimento, Ulrik Unneberg, Laura Zichi, Francesco Libbi, Norma Rivano, Austin Glover, Vivek Bharadwaj, Anders Jo

Published 2026-07-31

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Seán R. Kavanagh, Chuin Wei Tan, Menghang Wang, Marc L. Descoteaux, Gabriel de Miranda Nascimento, Ulrik Unneberg, Laura Zichi, Francesco Libbi, Norma Rivano, Austin Glover, Vivek Bharadwaj, Anders Johansson, William C. Witt, Albert Musaelian, Boris Kozinsky

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict how a giant, invisible Lego castle will behave. Will it crumble if you shake it? Will it conduct heat like a metal pan or insulate like a wool sweater? In the world of materials science, scientists use super-computers to simulate these tiny building blocks (atoms) to design new batteries, medicines, and super-strong metals. But there's a catch: the most accurate way to calculate how these atoms interact is like trying to solve a million-piece puzzle by hand—it takes forever and costs a fortune in computer time. To speed things up, scientists created "machine-learned interatomic potentials" (MLIPs). Think of these as a super-smart, fast-forwarding video game engine. Instead of calculating every single physics rule from scratch, the engine has "learned" from millions of previous calculations how atoms usually behave, allowing it to guess the outcome in a blink. Recently, scientists started training these engines on massive datasets to create "foundation models"—a sort of universal cheat code that can be tweaked for specific jobs. But here's the big question: Can we make these cheat codes both incredibly fast and incredibly accurate, or do we have to choose one over the other?

This paper is the story of a team that decided to build the ultimate cheat code, proving that you don't have to pick a side. They developed a new family of "foundation models" using two specific architectures called NequIP and Allegro. Imagine these models as two different types of super-athletes: one is a meticulous architect who checks every angle (NequIP), and the other is a lightning-fast sprinter who sees the big picture without getting bogged down in details (Allegro). The researchers trained these athletes on a massive library of over 100 million atomic structures, a dataset so huge it would have taken weeks to train on older systems. By using some clever software tricks—like "graph compilation" (which organizes the data before the race starts) and "mixed-precision math" (using a slightly less detailed calculator for the easy parts of the run)—they managed to train these models 5 to 10 times faster than before, slashing the cost from weeks of computing time down to just a few hundred GPU hours.

The results are a game-changer. The team found that these new models are not just fast; they are also the most accurate they've ever seen for predicting how materials behave. They tested them on three major challenges: finding the most stable shape for new materials (like finding the perfect Lego structure), predicting how well a material conducts heat (crucial for electronics), and calculating how materials stretch or squish. The models passed with flying colors, often beating previous records. However, the paper also points out a specific "blind spot." While the models are great at general chemistry, they still struggle a bit with transition metals (like iron, chromium, and manganese) because the training data for these tricky elements is a bit inconsistent. It's like the models are fluent in most languages but get confused by a few dialects where the grammar rules change unexpectedly.

The authors also looked at how these models scale up. They showed that these models can run on thousands of computer chips working together, simulating systems with over 100 million atoms. This is huge because it means scientists can now simulate entire grains of sand or tiny droplets of liquid with atomic precision, something that was previously impossible. The paper suggests that the future of better materials isn't just about making bigger models, but about cleaning up the training data—especially for those tricky transition metals—and making sure the "rules" used to generate the data are consistent. In short, they've built a faster, sharper, and more reliable tool for discovering the materials of tomorrow, proving that with the right software upgrades, we can have our cake and eat it too: speed and accuracy, all in one package.

Technical Summary: Fast and Accurate Foundation Models for Equivariant Machine-Learned Interatomic Potentials

Problem Statement
Machine-learned interatomic potentials (MLIPs) have become essential for computational materials science, particularly as "foundation models" trained on large, diverse datasets (e.g., MPtrj, Alexandria, OMat24) for downstream fine-tuning. However, a critical bottleneck exists: scientific applications like molecular dynamics (MD) require both high accuracy and high inference/training speeds. While equivariant architectures (which encode physical symmetries like rotation and inversion) offer superior accuracy and data efficiency, training them on ultra-large datasets (hundreds of millions of structures) has historically been computationally prohibitive, often requiring weeks of training time. Furthermore, it remains unclear how to balance the competing demands of model size, inference speed, and accuracy, particularly when scaling to massive datasets where data efficiency is less critical than raw performance.

Methodology
The authors present a family of foundation potentials based on the NequIP and Allegro equivariant graph neural network architectures, trained on increasingly large inorganic materials datasets: DIRECT (~~186k), MPtrj (~~1.58M), MPA (~~10.5M), and OAM (~~113M structures).

Key methodological advancements include:

Infrastructure Accelerations: The work leverages recent developments in the NequIP framework, including train-time graph compilation (torch.compile), distributed data-parallel (DDP) strategies for multi-GPU training, and optimized tensor product kernels (OpenEquivariance for NequIP, cuEquivariance for Allegro).
Mixed-Precision Training: The authors utilize mixed-precision strategies (float64 for embeddings, float32 for tensor products, and automatic mixed precision for intermediate computations) to accelerate training by up to 3x without significant loss in accuracy.
Two-Stage Training: An optimal training protocol involves an initial run with high weighting on force errors, followed by a continuation run with increased energy error weighting, reduced learning rates, and full precision.
Model Variants: Models of varying sizes (Small, Medium, Large, Extra-Large) were trained to explore the trade-off between parameter count, accuracy, and speed. The OAM dataset training involved a specific protocol: pre-training on OMat24 (non-equilibrium geometries) followed by fine-tuning on MPA to ensure compatibility with Materials Project settings.

Key Contributions

Accelerated Training Pipeline: The integration of graph compilation and accelerated kernels enables the training of high-accuracy foundation models on datasets exceeding 100 million structures with a modest computational cost of ~100–750 GPU hours (using H100/H200 GPUs), a significant reduction from previous multi-week requirements.
Benchmarking Suite: The models were rigorously evaluated against community benchmarks covering:
- Materials Discovery: Predicting ground-state geometries and convex hull energies (Matbench Discovery).
- Thermal Conductivity: Predicting off-equilibrium properties via phonon force constants.
- Near-Equilibrium Properties: Evaluating bulk/shear moduli, heat capacity, and force magnitude ratios (MatCalc).
Scalability Analysis: The paper demonstrates that the NequIP architecture, previously limited by iterative message-passing, now supports efficient multi-GPU inference in LAMMPS via the Kokkos-accelerated ML-IAP interface, enabling simulations of systems with up to 100 million atoms.

Results

Accuracy: The NequIP-OAM-XL model achieves state-of-the-art performance across benchmarks, with a convex hull energy MAE of 19.7 meV/atom and a thermal conductivity Symmetric Relative Mean Error ( $\kappa_{SRME}$ ) of 0.129.
Scaling Laws:
- Model Size: Accuracy generally follows a power-law improvement with model size for small-to-large models. However, the XL model shows signs of saturation, suggesting diminishing returns for specific architectures at extreme scales.
- Dataset Size: Thermal conductivity accuracy scales ideally with dataset size (power law exponent $\alpha \approx 0.24$ for NequIP), confirming the value of non-equilibrium geometries. Materials discovery accuracy, however, shows limited improvement when scaling from MPA to OAM.
Error Analysis: The lack of improvement in materials discovery accuracy with larger datasets is attributed to compositional diversity and inconsistent descriptions of transition metal compounds. Specifically, multi-valent first-row transition metals (V, Cr, Mn, Fe) and actinides (Np, Pu) exhibit significantly higher errors. The authors attribute this to the inconsistent application of Hubbard $U$ corrections in the training data (applied to oxides/fluorides but not metals) and the inherent complexity of correlated electronic systems.
Speed vs. Accuracy: The trained models define the Pareto front for the accuracy-speed trade-off. For instance, NequIP-OAM-XL achieves comparable combined performance scores to other leading models (e.g., eSEN-30M-OAM) but with approximately an order of magnitude faster inference speeds.
Over-softening: Large models trained on diverse datasets (including off-equilibrium structures) exhibit minimal "over-softening" (underestimation of forces in high-energy states), with force magnitude ratios ( $f/f_{DFT}$ ) reaching 0.98.

Significance and Claims
The paper claims that the primary barrier to training high-accuracy foundation potentials is no longer computational cost but rather data quality and diversity. The authors argue that future progress in materials discovery MLIPs depends less on raw model scaling and more on:

Consistent Data: Eliminating systematic errors caused by inconsistent computational settings (e.g., Hubbard $U$ application) across training datasets.
Chemical Diversity: Expanding datasets beyond ionic substitution to include broader chemical and structural diversity, particularly for complex transition metal systems.
Architectural Adaptability: Developing models that can dynamically adjust radial and angular resolution based on local chemical environments to handle the disparate complexity of different elements without uniformly inflating model capacity.

The work establishes that with current infrastructure, it is feasible for individual research groups to train and deploy foundation models capable of simulating ~100 million atoms, democratizing access to high-fidelity, large-scale atomistic simulations. The provided models serve as robust starting points for fine-tuning in targeted chemical spaces.

More like this