Flexible Cutoff Learning: Optimizing Machine Learning Potentials After Training

TL;DR

Results show that FCL enables training of a single general-purpose MLIP that can be adapted to diverse applications through post-training cutoff optimization, eliminating the need for retraining.

Abstract

We introduce Flexible Cutoff Learning (FCL), a method for training machine learning interatomic potentials (MLIPs) whose cutoff radii can be adjusted after training. Unlike conventional MLIPs that fix the cutoff radius during training, FCL models are trained by randomly sampling cutoff radii independently for each atom. The resulting model can then be deployed with different per-atom cutoff radii depending on the application, enabling application-specific optimization of the accuracy-cost tradeoff. Using a differentiable cost model, these per-atom cutoffs can be optimized for specific target systems after training. We demonstrate FCL with a modified MACE architecture trained on the MAD dataset. For a subset featuring molecular crystals, optimized per-atom cutoffs reduce computational cost by more than 60% while increasing force errors by less than 1%. These results show that FCL enables training of a single general-purpose MLIP that can be adapted to diverse applications through post-training cutoff optimization, eliminating the need for retraining.

Figures (10)

• Figure 1: Simplified representation of a model that depends explicitly on atom-wise cutoff radii. For a fixed geometry, the model can be evaluated for different realizations of $r_\text{cut}^{(1)},...,r_\text{cut}^{(n)}$. In this schematic figure, the model is assumed to be decomposable into on-site contributions that are computed from atom-centered atomic environments. In particular, each contribution is obtained by evaluating a function (denoted as $V$) on atomic environments of different sizes that depend on the choice of the atom-wise cutoff radii.
• Figure 2: Schematic overview of the Flexible Cutoff Learning and the optimization of cutoff radii after training. (a) Per-atom cutoff radii $r_\text{cut}^{(i)}$ are sampled during the training process. (b) Optimization of error and cost with respect to cutoff radii is performed after training. Smaller receptive fields lead to faster evaluation but may reduce the model accuracy.
• Figure 3: Test errors as a function of a global cutoff radius. Results for "Flexible Cutoff Learning" are obtained from evaluating a single model with different global cutoff radii. "Static Cutoff Learning" correspond to a separate model for each cutoff radius. All results are averaged over three training runs with different seeds. Error bars correspond to the empirical standard deviation.
• Figure 4: Energy as a function of distance in diatomic systems for different inference cutoff radii. All curves are obtained by evaluating a single FCL model.
• Figure 5: Force RMSE versus average number of pairs per atom for different MAD test subsets. Each curve shows results for varying $\lambda$ (from left to right: $10^{-2}, 10^{-3}, 10^{-4}, 10^{-5}, 10^{-6}$), demonstrating the accuracy-cost tradeoff achievable through cutoff optimization.