# High-performance training and inference for deep equivariant interatomic potentials

arXiv: https://arxiv.org/abs/2504.16068 · HTML: https://mlsanity.com/paper/2504.16068/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 39 references that are in the map.
- First posted on arXiv 22 Apr 2025.

## Builds on (most cited first, 10 of 10)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials](https://mlsanity.com/paper/2205.06643/) (2022)
- [SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials](https://mlsanity.com/paper/2209.10702/) (2022)
- [Cartesian atomic cluster expansion for machine learning interatomic potentials](https://mlsanity.com/paper/2402.07472/) (2024)
- [SchNetPack 2.0: A neural network toolbox for atomistic machine learning](https://mlsanity.com/paper/2212.05517/) (2022)
- [Nutmeg and SPICE: Models and Data for Biomolecular Machine Learning](https://mlsanity.com/paper/2406.13112/) (2024)
- [Equivariant Machine Learning Interatomic Potentials with Global Charge Redistribution](https://mlsanity.com/paper/2503.17949/) (2025)

## Built on by (most cited first, 17 of 17)

- [EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers](https://mlsanity.com/paper/2604.09130/) (2026)
- [Training a Foundation Model for Materials on a Budget](https://mlsanity.com/paper/2508.16067/) (2025)
- [Spectral/Spatial Tensor Atomic Cluster Expansion with Universal Embeddings in Cartesian Space](https://mlsanity.com/paper/2509.14961/) (2025)
- [PFT: Phonon Fine-tuning for Machine Learned Interatomic Potentials](https://mlsanity.com/paper/2601.07742/) (2026)
- [A Cartesian-3j Framework for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2512.16882/) (2025)
- [Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2607.02499/) (2026)
- [Smooth Dynamic Cutoffs for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2601.21147/) (2026)
- [Mixture of Experts Framework in Machine Learning Interatomic Potentials for Atomistic Simulations](https://mlsanity.com/paper/2604.26143/) (2026)
- [Polarizable atomic multipoles for learning long-range electrostatics](https://mlsanity.com/paper/2605.05746/) (2026)
- [CrystalREPA: Transferring Physical Priors from Universal MLIPs to Crystal Generative Models](https://mlsanity.com/paper/2605.08960/) (2026)
- [How accurate are foundational machine learning interatomic potentials for heterogeneous catalysis?](https://mlsanity.com/paper/2512.16702/) (2025)
- [Performance of universal machine-learned potentials with explicit long-range interactions in biomolecular simulations](https://mlsanity.com/paper/2508.10841/) (2025)
- [Twirlator: A Pipeline for Analyzing Subgroup Symmetry Effects in Quantum Machine Learning Ansatzes](https://mlsanity.com/paper/2511.04243/) (2025)
- [Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning](https://mlsanity.com/paper/2605.08885/) (2026)
- [Robust and Interpretable Adaptation of Equivariant Materials Foundation Models via Sparsity-promoting Fine-tuning](https://mlsanity.com/paper/2606.18691/) (2026)
- [A strategic roadmap for an atomistic machine-learning ecosystem](https://mlsanity.com/paper/2609.39090/) (2026)
- [BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials](https://mlsanity.com/paper/2610.02013/) (2026)

## Abstract (arXiv)

Machine learning interatomic potentials, particularly those based on deep equivariant neural networks, have demonstrated state-of-the-art accuracy and computational efficiency in atomistic modeling tasks like molecular dynamics and high-throughput screening. The size of datasets and demands of downstream workflows are growing rapidly, making robust and scalable software essential. This work presents a major overhaul of the NequIP framework focusing on multi-node parallelism, computational performance, and extensibility. The redesigned framework supports distributed training on large datasets and removes barriers preventing full utilization of the PyTorch 2.0 compiler at train time. We demonstrate this acceleration in a case study by training Allegro models on the SPICE 2 dataset of organic molecular systems. For inference, we introduce the first end-to-end infrastructure that uses the PyTorch Ahead-of-Time Inductor compiler for machine learning interatomic potentials. Additionally, we implement a custom kernel for the Allegro model's most expensive operation, the tensor product. Together, these advancements speed up molecular dynamics calculations on system sizes of practical relevance by up to a factor of 18.