Methodological and Performance Advances in the NequIP and Allegro Suite of Deep Equivariant Interatomic Potentials

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 revamp of the NequIP framework, emphasizing multi-node parallelism, computational performance, and modular extensibility [1]. The redesigned framework supports distributed training on large datasets, removes barriers preventing full utilization of the PyTorch 2.0 compiler for training and inference, and enables seamless integration of custom Triton kernels as well as optimized GPU kernels from OpenEquivariance and CuEquivariance. Training performance improvements have enabled the creation of the first NequIP framework foundation potentials, now available at https://www.nequip.net/. Inference-side advancements, including compiler integration and optimized tensor-product kernels, accelerate molecular dynamics simulations on systems of practical relevance by factors of 5 to 18, making the NequIP and Allegro foundation potentials among the fastest and most accurate models on the Matbench Discovery benchmark.

[1] C. W. Tan, M. L. Descoteaux et al. High-performance training and inference for deep equivariant interatomic potentials. arXiv preprint arXiv:2504.16068 (2025).