Fast and flexible long-range models for atomistic machine learning.

Area of Science:

• Computational Chemistry
• Materials Science
• Machine Learning

Background:

• Atomistic machine learning (ML) models typically use local approximations, limiting their ability to model long-range interactions like electrostatics.
• Existing methods to address long-range effects in ML are often inefficient and require ad-hoc implementations.

Purpose of the Study:

• To develop a unified framework for integrating established long-range interaction algorithms into atomistic ML.
• To provide efficient and modular implementations for evaluating non-bonded interactions in ML models.
• To introduce novel descriptors suitable for ML applications dominated by long-range physics.

Main Methods:

• Incorporation of Ewald summation, classical particle-mesh Ewald (PME), and particle-particle/particle-mesh (PPPM) into atomistic ML frameworks.
• Development of reference implementations in PyTorch (torch-pme) and an experimental one in JAX (jax-pme).
• Introduction of purified descriptors that focus on non-local atomic environments.

Main Results:

• Fast, feature-rich, and modular implementations for accurate evaluation of physical long-range forces.
• Seamless combination of long-range models with local ML schemes using automatic differentiation.
• Demonstrated utility in molecular dynamics simulations, ML potential training, and evaluation of long-range equivariant descriptors.

Conclusions:

• The developed framework and libraries effectively address the limitations of local approximations in atomistic ML.
• Enables the construction of accurate (semi)empirical baseline potentials and complex ML architectures incorporating physical interactions.
• Facilitates advanced molecular simulations and the development of novel ML models for systems with significant long-range effects.