Apax: A Flexible and Performant Framework For The Development of Machine-Learned Interatomic Potentials

Abstract

We introduce Atomistic learned potentials in JAX (apax), a flexible and efficient open source software package for training and inference of machine-learned interatomic potentials. Built on the JAX framework, apax supports GPU acceleration and implements flexible model abstractions for fast development. With features such as kernel-based data selection, well-calibrated uncertainty estimation, and enhanced sampling, it is tailored to active learning applications and ease of use. The features and design decisions made in apax are discussed before demonstrating some of its capabilities. First, a data set for the room-temperature ionic liquid EMIM+BF4- is created using active learning. It is highlighted how continuously learning models between iterations can reduce training times up to 85 % with only a minor reduction of the models' accuracy. Second, we show good scalability in a data-parallel training setting. We report that a Gaussian Moment Neural Network model, as implemented in apax, achieves higher accuracy and up to 10 times faster inference times than a performance-optimized Allegro model. A recently published Li3PO4 dataset, reported with comparable accuracy and inference performance metrics, is used as a point of comparison. Moreover, the inference speeds of the available simulation engines are compared. Finally, to highlight the modularity of apax, an equivariant message-passing model is trained as a shallow ensemble and used to perform uncertainty-driven dynamics.

References in corpus (19)

• E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials
• PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments and Partial Charges
• DeePMD-kit v2: A software package for Deep Potential models
• Committee neural network potentials control generalization errors and enable active learning
• Gaussian Moments as Physically Inspired Molecular Descriptors for Accurate and Scalable Machine Learning Potentials
• Uncertainty estimation for molecular dynamics and sampling
• Fast Uncertainty Estimates in Deep Learning Interatomic Potentials
• Deep Ensembles vs. Committees for Uncertainty Estimation in Neural-Network Force Fields: Comparison and Application to Active Learning
• A Novel Approach to Describe Chemical Environments in High Dimensional Neural Network Potentials
• Scalable Hybrid Deep Neural Networks/Polarizable Potentials Biomolecular Simulations including long-range effects
• Transfer learning for chemically accurate interatomic neural network potentials
• Force-Field-Enhanced Neural Network Interactions: from Local Equivariant Embedding to Atom-in-Molecule properties and long-range effects
• Lifelong Machine Learning Potentials
• Fast and Sample-Efficient Interatomic Neural Network Potentials for Molecules and Materials Based on Gaussian Moments
• Neural-Network Assisted Study of Nitrogen Atom Dynamics on Amorphous Solid Water. I. Adsorption & Desorption
• Thermally Averaged Magnetic Anisotropy Tensors via Machine Learning Based on Gaussian Moments
• Metadynamics sampling in atomic environment space for collecting training data for machine learning potentials
• Enabling robust offline active learning for machine learning potentials using simple physics-based priors
• FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials

Cited by in corpus (2)

• Breaking scaling relations with inverse catalysts: a machine learning exploration of trends in $\mathrm{C}{\mathrm{O}}_2$ hydrogenation energy barriers
• Active learning potentials for first-principles phase diagrams using replica-exchange nested sampling