Triplet Envelope Functions for increasing machine learning interatomic potential efficiency and stability

TL;DR

This work introduces triplet envelope functions as a higher-order, energy-conserving alternative to KNN sparsification for MLIPs. Building on MEAM-inspired screening, it couples radial envelopes with triplet-based attenuation and a logistic low-pass filter to prune neighborhoods while preserving a smooth potential energy surface. Experiments on silicon and water show that triplet envelopes double training and inference speed, triple memory efficiency, and enhance stability with minimal or no loss in accuracy, even at larger cutoffs like $8$ Å. The approach enables larger radii to model open structures more efficiently, marking a promising direction for scalable, stable MLIP simulations.

Abstract

Central to interatomic potential efficiency is the radial envelope function that enables linear scaling with computational cost by defining a local neighborhood of atoms. This has enabled MLIPs to revolutionize materials science over the past decade by providing DFT accuracy with linear scaling computational cost in molecular dynamics workflows. However, MLIPs still have a relatively high computational cost compared to empirical interatomic potentials, preventing them from transforming molecular dynamics workflows. A central issue is that MLIPs use relatively large cutoff radii, converging to 6A over the last few years. The large cutoffs prioritize accuracy of any material over efficiency in any particular region of phase space, capturing dispersion effects and low density materials at the expense of increased computational cost in higher density materials. Past work has aimed to address this with KNN graph sparsification, which, while significantly reducing cost, has the drawback of breaking energy conservation. In this work, we propose higher-order envelope functions that prune local atomic neighborhoods through physically inspired geometric functions to provide the memory and efficiency benefits of KNN graph sparsification while eliminating non-conservative energy dynamics. Through numerical experiments on solids and liquids with 5-8A cutoffs, we show that triplet envelope functions complement radial envelope functions by doubling training and inference speed, tripling memory efficiency, and increasing simulation stability while not impacting accuracy or data efficiency for the most common 6A cutoff. Moreover, experiments with 8A radial cutoffs show triplet envelope functions create a pathway to larger cutoff radii for efficiently and accurately modeling open structures with large interatomic distances, showing a promising new direction for engineering MLIP efficiency.

Figures (5)

• Figure 1: Geometric representation of triplet screening. (a) Schematic of the ellipse-based screening geometry, where the position of atom $k$ relative to atoms $i$ and $j$ determines its contribution to the envelope function $f_{ijk}$. (b) Progression of the screening function $S_{ijk}$ across representative three-body configurations, illustrating the transition from fully screened ($S_{ijk}=0$) to unscreened ($S_{ijk}=1$) states.
• Figure 2: Comparison of neighborhoods defined by radial and triplet envelope functions for a radial cutoff of 6Å. The top two rows depict the local neighborhood of a central atom (blue), where the top row is colored by chemical identity and the middle row by the triplet screening parameter. The bottom row depicts the local environments after pruning with the triplet envelope function.
• Figure 3: (left) The edges per atom for a radial (solid) envelope function compared to a triplet (dashed) envelope function for crystalline silicon, melted silicon, and liquid water. (right) The ratio of these two are used to find the maximum speed-up for each system.
• Figure 4: The change in training accuracy, training speed, inference speed, and inference memory pressure when including triplet envelope functions as compared to radial envelope functions. The training speed-up, inference speed-up, and the capacity ratio are the ratio of triplet to radial envelope function showing the improvement of the triplet envelope function in throughput and capacity.
• Figure 5: Comparison of the conserved energy of radial, KNN, triplet, and triplet low-pass edge sparsification for models trained on water and silicon.