Source-linked AI summary
On-the-Fly Active Learning of Interpretable Bayesian Force Fields for Atomistic Rare Events
Jonathan Vandermause, Steven B. Torrisi, Simon Batzner, Yu Xie, Lixin Sun, Alexie M. Kolpak, Boris Kozinsky
TL;DR
The paper targets the expense and unpredictability of manually trained ML force fields for rare-event dynamics. It introduces an adaptive Bayesian, Gaussian-process method that selects training structures on the fly and produces interpretable uncertainty estimates. The resulting models require few training environments, agree well with DFT, and can be mapped to faster force-field representations.
Problem
ML force fields commonly require thousands of first-principles calculations and can be unreliable for structures absent from their training data.
Method
FLARE trains interpretable low-dimensional Gaussian-process force fields on the fly, using Bayesian uncertainty to select when new first-principles data are needed.
Results
The models provide accurate force estimates and reliable internal uncertainty estimates while requiring fewer training environments than high-dimensional ML force fields.
Takeaways & Limitations
The approach supports force fields tailored to applications, with good agreement with DFT and mappings to faster regression models approaching classical-force-field speed.
Abstract
from arXiv · showhide
Machine learned force fields typically require manual construction of training sets consisting of thousands of first principles calculations, which can result in low training efficiency and unpredictable errors when applied to structures not represented in the training set of the model. This severely limits the practical application of these models in systems with dynamics governed by important rare events, such as chemical reactions and diffusion. We present an adaptive Bayesian inference method for automating the training of interpretable, low-dimensional, and multi-element interatomic force fields using structures drawn on the fly from molecular dynamics simulations. Within an active learning framework, the internal uncertainty of a Gaussian process regression model is used to decide whether to accept the model prediction or to perform a first principles calculation to augment the training set of the model. The method is applied to a range of single- and multi-element systems and shown to achieve a favorable balance of accuracy and computational efficiency, while requiring a minimal amount of ab initio training data. We provide a fully open-source implementation of our method, as well as a procedure to map trained models to computationally efficient tabulated force fields.
I. INTRODUCTION
The paper addresses unreliable and labor-intensive ML force-field training by introducing an uncertainty-driven Bayesian method that learns interpretable models on the fly. FLARE combines low-dimensional Gaussian-process force fields with active learning to reduce ab initio data needs while supporting complex and rare-event dynamics.
- Motivation: Manual training sets often require thousands of first-principles calculations, while missing uncertainty estimates make reliability and model updating difficult.These limitations hinder applications involving rare events and structures outside the training distribution.
- Method: FLARE uses Gaussian-process regression and internal model-error estimates to automate training during molecular dynamics.The uncertainty determines whether to accept a prediction or perform a new first-principles calculation.
- Results and applicability: Accurate force fields for single- and multi-element systems are obtained with approximately 100 DFT calculations, while the method models rapid crystal melts and rare diffusive events.The resulting models can be mapped to efficient tabulated force fields and are released as open-source software.
- Model design: The force field imposes a physical prior over n-body interactions and learns nonparametric 2- and 3-body functions from ab initio data.This avoids high-dimensional descriptors and sparsification while preserving interpretable energy contributions and kernel hyperparameters.
- Contribution: The force fields combine nonparametric, fully Bayesian, explicitly multi-element modeling with interpretable uncertainty estimates.These features are presented as a combination not previously achieved simultaneously.
A. FLARE: An on-the-fly learning method
FLARE uses low-dimensional, interpretable Gaussian-process force fields whose uncertainty guides on-the-fly DFT data acquisition. The approach supports multi-element systems, rare-event dynamics, and efficient tabulated production models.
- Model design: FLARE uses two- and three-body kernels that compare clusters by interatomic distances and chemical species.Species-mismatched clusters receive zero kernel value, preserving chemical distinction while remaining low dimensional.
- On-the-fly learning: The algorithm calls DFT when any force-component epistemic uncertainty exceeds a threshold based on the GP noise uncertainty.The training set is augmented with the highest-uncertainty local environments, and hyperparameters are reoptimized as data are added.
- Uncertainty characterization: Optimized noise uncertainty closely tracks test RMSE across examined cutoff values, enabling expected error estimation after sufficient training.When training data are insufficient, epistemic uncertainties rise above the noise uncertainty, indicating the need for additional data.
- Uncertainty characterization: For perturbed structures, average epistemic uncertainty exceeds σn = 11.53 meV/˚A when δ > 5%, detecting configurations outside the training set.Higher uncertainties coincide with a wider spread in true force errors, and most errors fall within two standard deviations.
C. Aluminum crystal melt
FLARE adapts an aluminum force field as molecular dynamics moves from an FCC crystal to a liquid. It concentrates DFT calls near initialization and the phase change while retaining low-cost production evaluation.
- Crystal-to-liquid adaptation: Fewer than 50 DFT calls were required during the first 5 ps while the model learned the FCC aluminum phase.DFT calls were frequent initially and became much less frequent after approximately 30 time steps.
- Phase behavior: The final trajectory stabilized around 5000 K and had a liquid-phase radial distribution function after melting.The initial trajectory segment had an FCC-like radial distribution function, whereas the final half was liquid-like.
- Crystal-to-liquid adaptation: After melting, the model made frequent DFT calls because liquid local environments differed substantially from the preceding solid-state environments.The optimized noise uncertainty increased sharply as the system entered the liquid phase.
- Training efficiency: Fewer than 100 DFT calls were needed across 10 ps, with most calls occurring at the simulation start and immediately after melting.The adaptive threshold increased in the liquid phase, yielding roughly similar call counts for learning the solid and liquid phases.
- Computational efficiency: The 2+3-body tabulated FLARE model cost 4.9 × 10^-5 s/atom/timestep, remaining faster than AGNI.The tabulated 2-body model cost 5.6 × 10^-6 s/atom/timestep, versus 2.2 × 10^-6 for EAM.
D. Bulk vacancy and surface adatom diffusion
FLARE learns aluminum force fields for vacancy and adatom diffusion while sampling rare-event trajectories with limited additional DFT. Its predicted migration energetics agree closely with DFT, including a subtle surface-site preference.
- Bulk vacancy diffusion: After approximately 400 ps, no additional DFT calls were required during the 1 ns vacancy-diffusion training run.Vacancy hops appeared every few hundred picoseconds as sharp jumps in the mean squared displacement.
- Bulk vacancy diffusion: FLARE predicted vacancy-migration forces and energies in close agreement with ab initio DFT along the transition path.DFT forces lay within one standard deviation of GP force predictions, and the predicted barrier closely matched DFT.
- Computational efficiency: The vacancy-diffusion training run was over 300 times faster than an equivalent AIMD run.The comparison included DFT calls, GP optimization, force evaluations, and molecular-dynamics updates.
- Surface adatom diffusion: Three-body contributions considerably increased adatom force-field accuracy for the aluminum (111) surface.The selected model used a 7 Å two-body cutoff and 4.5 Å three-body cutoff, with σn = 44.2 meV/˚A after the final DFT call.
- Surface adatom diffusion: FLARE reproduced the slight equilibrium preference for the hcp adatom site, with NEB energies agreeing with DFT within approximately 20 meV.The EAM comparison did not reproduce this hcp preference.
E. Fast-ion diffusion in AgI
FLARE models fast-ion diffusion and the phase transition in silver iodide using a short on-the-fly training simulation followed by large-scale tabulated-force-field dynamics. The predicted transition and diffusion behavior agree with AIMD and experiment within the reported comparisons.
- On-the-fly training: Thirty-nine DFT calls were made during the 15 ps AgI on-the-fly training simulation.The uncertainty threshold was set to σthresh = 2σn, producing a roughly 1 ps delay after a temperature increase before the next DFT call.
- Diffusion and phase transition: FLARE diffusion coefficients showed a sharp increase between 400 K and 450 K, consistent with the experimental fast-ion transition temperature of 420 K.The coefficients were obtained from 1 ns NPT simulations of 10,976 AgI atoms using the tabulated force field.
- Diffusion and phase transition: FLARE and AIMD agreed well with experimental α-phase Ag diffusion coefficients, with activation energies of 0.107, 0.114, and 0.093 eV, respectively.The reported comparison includes a slight vertical offset between the calculated and experimental diffusion data.
- Diffusion and phase transition: Below the transition temperature, FLARE predicted a non-diffusive, non-cubic hcp phase with 12 nearest-neighbor iodine coordination.This accounts for the low-temperature discrepancy from AIMD simulations constrained to the α-phase with a fixed cubic cell.
- Structural validation: The simulated low- and high-temperature structures had c/a = 1.46 and a_lat = 5.30 Å, compared with experimental values of 1.63 and 5.07 Å.These structural values correspond to the β- and α-phase comparisons, respectively.
F. General applicability
FLARE was applied automatically to diverse single- and multi-element systems, with uncertainty-triggered updates during molecular dynamics and low force-prediction errors.
- FLARE was applied to five representative systems spanning carbon, silicon, aluminum oxide, nickel titanium, and two-dimensional boron nitride.
- Temperature increases triggered DFT calls and GP updates when the model detected novel local environments.
- ~20 DFT calls were sufficient after setting the error threshold to twice the noise uncertainty.
- 16 DFT calls trained the aluminum oxide model when 30 local environments were added after each call.
- The training procedure updated training sets and hyperparameters on the fly without human guidance.
- Validation RMSEs were around 0.1 eV/˚A for four systems and 0.42 eV/˚A for carbon at higher temperature.
III. DISCUSSION
The discussion presents FLARE as an interpretable Bayesian force-field framework that uses uncertainty to guide data acquisition and can be accelerated for broader atomistic applications.
- FLARE uses model uncertainties that correlate with out-of-sample error to select force-field training points during molecular dynamics.
- Its nonparametric 2- and 3-body models require fewer training environments than high-dimensional approaches.
- The models combine simple, accurate, physically interpretable energy representations with reliable internal force estimates.
- Trained GP models can be mapped to faster regression models approaching classical-force-field speed.
- This approach targets force fields with good agreement with DFT at several orders of magnitude lower computational cost.
- Extending active learning to complex systems such as polymers and proteins remains an important open challenge.
A. Gaussian process force fields
The force fields use Gaussian processes to learn physically motivated, low-dimensional n-body energy contributions from local atomic environments. The formulation preserves species distinctions, relevant symmetries, energy conservation, and computational efficiency.
- Energy representation: The model represents total energy as a sum of local n-body contributions up to a chosen maximum order N.Each local environment contains the central atom and neighboring clusters within an n-dependent cutoff.
- Energy representation: Small-n cluster descriptors have much lower dimensionality than the full potential energy surface.The implementation uses species vectors and interatomic distance vectors for pair and triplet environments.
- Gaussian-process model: Gaussian process regression learns cluster energy contributions from ab initio force data using kernels that compare rotationally invariant interatomic distance vectors.The local energy kernel is a sum over kernels between corresponding atomic clusters.
- Gaussian-process model: Species-aware kernels and permutations of same-species atoms enforce chemical identity and permutation invariance.The species delta function distinguishes unlike species, while permutation sums preserve invariance under rearrangement of atoms of the same species.
- Force prediction: Differentiating the energy kernel with respect to central-atom coordinates produces an exactly rotationally covariant and energy-conserving force model.The resulting force kernel describes covariance between force components.
- Model choice: N = 3 restricts the model to 2- and 3-body contributions, which the authors found sufficiently accurate across single- and multi-element systems while remaining computationally efficient.Accuracy as a function of maximum kernel order remains an identified area for future study.
- Kernel design: The pair and triplet kernels use squared exponential forms multiplied by smooth quadratic cutoffs to maintain continuity as atoms cross the cutoff boundary.Signal variance controls the maximum uncertainty far from training data, while length scales govern the 2- and 3-body contributions.
- Training: Hyperparameters are optimized by maximizing the training-data log marginal likelihood with BFGS and analytic gradients.This optimization is efficient when the model contains fewer than approximately 1000 training points because it requires covariance-matrix inversion.
B. Mapping to tabulated spline models
The trained Gaussian-process force fields can be mapped to cubic spline models that eliminate the expensive loop over training points. The multi-element mapping uses spline interpolants and can approach the original GP accuracy by increasing control points.
- Spline mapping: Cubic spline mapping eliminates the expensive loop over training points required to calculate the GP kernel vector.The procedure extends an earlier single-element mapping to multi-element kernels.
- Spline mapping: The multi-element mapping constructs cubic spline interpolants for the force-field components.The resulting representation uses 1- and 3-dimensional cubic splines, respectively.
- Accuracy: Increasing the number of spline control points can make the mapped model arbitrarily accurate relative to the original Gaussian process model.Control-point grids for the mapped force fields are reported in the Supplementary Information.
C. Computational details
The implementation couples molecular dynamics, Gaussian-process modeling, and Quantum ESPRESSO for on-the-fly learning. Numerical acceleration and specified simulation ensembles support trajectories spanning hundreds of picoseconds.
- Electronic-structure calculations: Density-functional calculations use Quantum Espresso 6.2.1 with system-specific pseudopotentials, k-point meshes, and energy cutoffs.The calculation settings are listed in the Supplementary Information.
- Software: The on-the-fly learning algorithm is implemented in FLARE, coupling Python-based molecular dynamics and Gaussian-process code with Quantum ESPRESSO.Numba accelerates kernel and distance calculations.
- Molecular dynamics: Training simulations span hundreds of picoseconds, with on-the-fly molecular dynamics trajectories performed in the NVE ensemble using the Verlet algorithm.AgI LAMMPS simulations instead use the NPT ensemble at zero pressure.
- Availability: The trained Gaussian-process models and other simulation data are made available.The work also documents contributions to FLARE, GP-to-spline mapping, and LAMMPS integration.