Latent Ewald summation for machine learning of long-range interactions (2408.15165v2)

Abstract: Machine learning interatomic potentials (MLIPs) often neglect long-range interactions, such as electrostatic and dispersion forces. In this work, we introduce a straightforward and efficient method to account for long-range interactions by learning a latent variable from local atomic descriptors and applying an Ewald summation to this variable. We demonstrate that in systems including charged and polar molecular dimers, bulk water, and water-vapor interface, standard short-ranged MLIPs can lead to unphysical predictions even when employing message passing. The long-range models effectively eliminate these artifacts, with only about twice the computational cost of short-range MLIPs.

• The paper presents the Latent Ewald Summation method, a novel approach that learns latent variables from local atomic descriptors to incorporate long-range interactions into MLIPs.
• The method shows significant improvements with an order of magnitude reduction in RMSE for molecular dimers and enhanced prediction accuracy in bulk water and interfacial systems.
• The LES technique integrates seamlessly into existing MLIP architectures, offering higher data efficiency and realistic handling of long-range electrostatics without complex empirical corrections.

Latent Ewald Summation for Machine Learning of Long-Range Interactions

Introduction

Machine Learning Interatomic Potentials (MLIPs) serve as a valuable tool for quick and accurate predictions of energies and forces for atomic configurations, achieving these predictions by learning from reference quantum mechanical calculations. Despite their utility, state-of-the-art MLIPs conventionally adopt a short-range approximation. This implies the potential energy perceived by an atom is influenced by its immediate atomic neighborhood, leading to the total energy being expressed as a sum of individual atomic contributions, offering linear scalability with respect to the system size. However, this short-range approximation often fails to account for long-range interactions such as Coulomb and dispersion forces, resulting in unphysical predictions in specific systems, including liquid-vapor interfaces and dilute ionic solutions.

Various approaches have been proposed to incorporate long-range interactions in MLIPs, such as empirical corrections for electrostatics and dispersion, predicting partial charges for each atom, and integrating message-passing neural networks (MPNNs). Although these methods have had varying degrees of success, they come with limitations in terms of implementation complexity or accuracy. This paper presents the Latent Ewald Summation (LES) methodology, a straightforward mechanism to account for long-range interactions in atomistic systems by learning a latent variable from local atomic descriptors and applying Ewald summation to this variable. The LES technique can be integrated into most existing MLIP architectures, including those based on local atomic environments and MPNNs.

Methodology

The LES method focuses on periodic atomic systems and employs a dual approach where the total potential energy comprises both short-range and long-range components.

For the short-range part, the energy is summed over the atomic contributions: $E^{sr} = \sum_{i} E_{\theta}(B_i)$

Here, $E_{\theta}$ is a multilayer perceptron (MLP) with parameters $\theta$ mapping invariant atomic features (denoted as $B$) to the short-range atomic energy.

For the long-range component, another MLP with parameters $\phi$ maps the invariant features of each atom $i$ to a hidden variable: $q_i = Q_{\phi}(B_i)$

The structure factor $S(\mathbf{k})$ of the hidden variable is computed as: $$S(\mathbf{k}) = \sum_i q_i e^{i\mathbf{k}\mathbf{r}_i}$$

The long-range energy is then obtained via the Ewald summation: $$E^{lr} = \dfrac{1}{V} \sum_{0<k<k_{c}} \dfrac{e^{-\sigma^2 k^2/2}}{k^2} |S(\mathbf{k})|^2$$

The LES method provides two interpretations:

1. The hidden variable $q$ can be considered environment-dependent partial charges.
2. As a communication mechanism for atoms far apart in the simulation, akin to recently proposed Ewald-based long-range message-passing methods.

Results and Discussion

Molecular Dimers:

LES was benchmarked on various molecular dimer classes (charged, polar, and apolar). Short-range (SR) models with a cutoff $r_{\text{cut}}$ of 5 Å showed significant prediction errors for binding energies and forces of dimers separated by distances beyond the cutoff. In contrast, long-range (LR) models using LES demonstrated substantial improvements with better accuracy in energy and force predictions, especially for charged and polar dimers, achieving an order of magnitude reduction in root mean square errors (RMSEs).

Bulk Water:

For bulk water, LES dramatically enhanced both energy and force prediction capabilities when compared to SR models. The presence of message-passing layers (T=1) improved the accuracy further, but the inclusion of the LES component yielded models with higher data efficiency and significantly smaller errors in prediction. The LR components effectively mitigated unphysical artifacts in long-range dipole correlation functions observed in SR models.

Water Interfaces:

For the water liquid-vapor interface dataset, LES overcame the shortcomings of SR models that fail to capture interfacial properties correctly. SR models, even with message passing, exhibited incorrect extended dipole ordering into the bulk, while LR models accurately captured the dielectric screening and polarization behavior at liquid-vapor interfaces.

Conclusion

The LES method offers a simple and general approach to integrate long-range interactions in MLIP architectures without requiring user-defined empirical corrections, partial charges, or complex charge equilibration schemes. The method is straightforward to integrate into various existing MLIP frameworks and cost-efficient, incurring only about twice the computational cost of short-range MLIPs. LES has shown promise in accurately predicting the binding curves of molecular dimers, dipole correlations in bulk fluids, and interfacial properties, making it highly applicable for systems where long-range interactions are significant, such as aqueous solutions and biological processes. Scaling implementations and further optimizations can potentially enhance the efficiency and applicability of LES across broader chemistry and materials science domains.