Billions of Atoms with Machine Learning Interatomic Potentials: Performance Portability of FLARE

Machine learning interatomic potentials (MLIPs) have become a prevalent approach to bridging the gap between slow-but-accurate ab-initio calculations and fast-but-inaccurate classical potentials for molecular dynamics. There is now a wide spectrum of state-of-the-art MLIPs that can quantitatively predict the behavior and properties of materials while being orders of magnitude faster than ab-initio calculations. Yet many materials properties require time scales or system sizes that are not easily accessible with existing MLIPs. The issue is compounded by most MLIPs not taking advantage of modern hardware accelerators such as GPUs.

In this work, we present a performance portable implementation of the FLARE interatomic potential. FLARE uses descriptors from the Atomic Cluster Expansion together with a sparse Gaussian process. After training, the model can be rewritten as a simple quadratic form for faster inference. We implement FLARE in LAMMPS with the Kokkos performance portability library, enabling efficient molecular dynamics simulations on GPUs across a wide range of system sizes. We demonstrate the speed and accuracy by running long trajectories for NaCl and obtaining well-converged Green-Kubo estimates of the thermal conductivity.