Developments in the NequIP and Allegro Ecosystem of Neural Equivariant Interatomic Potentials

Deep equivariant models have become an indispensable workhorse of machine-learning-based molecular dynamics, achieving unprecedented fidelity to ab initio training data at only a fraction of the computational cost. Notably, the highly scalable Allegro [1] model has been successfully deployed on structures as large as a 44-million-atom HIV capsid [2]. This work presents the architectural advancements in Allegro since then, highlighting key design choices and experimental findings. While the expressivity and learning capabilities of such equivariant models are well recognized, the training infrastructure required to unlock their full potential is equally critical. This aspect becomes more apparent as the field advances toward increasingly ambitious challenges involving larger, more diverse datasets and more intricate learning tasks. To meet these demands and facilitate the exploration of training techniques and heuristics, the NequIP [3] infrastructure (upon which Allegro is based) has undergone a major overhaul, prioritizing configurability and extensibility. We will discuss the core infrastructural changes and the new directions they enable, including techniques for accelerating training and streamlining hyperparameter tuning.

[1] Musaelian, Albert, et al. "Learning local equivariant representations for large-scale atomistic dynamics." Nature Communications 14.1 (2023): 579.

[2] Kozinsky, Boris, et al. "Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size." Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2023.

[3] Batzner, Simon, et al. "E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials." Nature communications 13.1 (2022): 2453.