Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators

2024-02-21

AI Summary

Machine learning force fields (MLFFs) often suffer from numerical instability in molecular dynamics (MD) simulations, limiting long-timescale modeling and introducing bias in observable estimation. To address this, we propose StABlE—a novel training framework featuring the first differentiable Boltzmann estimator, enabling end-to-end automatic differentiation of MD simulations. StABlE identifies and rectifies unstable regions without requiring additional *ab initio* calculations. By jointly incorporating quantum-mechanical energy/force data and system-level observables, it establishes the first semi-empirical, general-purpose training paradigm that simultaneously ensures first-principles accuracy and experimental observability. Extensive validation across organic molecules, tetrapeptides, and condensed-phase systems demonstrates significantly improved simulation stability, larger feasible time steps, enhanced data efficiency, and up to 40% reduction in observable prediction error—while maintaining full compatibility with mainstream MLFF architectures.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Intelligent Robots: State Estimation
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Machine learning force fields (MLFFs) are an attractive alternative to ab-initio methods for molecular dynamics (MD) simulations. However, they can produce unstable simulations, limiting their ability to model phenomena occurring over longer timescales and compromising the quality of estimated observables. To address these challenges, we present Stability-Aware Boltzmann Estimator (StABlE) Training, a multi-modal training procedure which leverages joint supervision from reference quantum-mechanical calculations and system observables. StABlE Training iteratively runs many MD simulations in parallel to seek out unstable regions, and corrects the instabilities via supervision with a reference observable. We achieve efficient end-to-end automatic differentiation through MD simulations using our Boltzmann Estimator, a generalization of implicit differentiation techniques to a broader class of stochastic algorithms. Unlike existing techniques based on active learning, our approach requires no additional ab-initio energy and forces calculations to correct instabilities. We demonstrate our methodology across organic molecules, tetrapeptides, and condensed phase systems, using three modern MLFF architectures. StABlE-trained models achieve significant improvements in simulation stability, data efficiency, and agreement with reference observables. The stability improvements cannot be matched by reducing the simulation timestep; thus, StABlE Training effectively allows for larger timesteps. By incorporating observables into the training process alongside first-principles calculations, StABlE Training can be viewed as a general semi-empirical framework applicable across MLFF architectures and systems. This makes it a powerful tool for training stable and accurate MLFFs, particularly in the absence of large reference datasets. Our code is available at https://github.com/ASK-Berkeley/StABlE-Training.

Problem

• Enhance MLFF simulation stability
• Avoid additional ab-initio calculations
• Improve agreement with reference observables

Innovation

• StABlE Training enhances MLFF stability.
• MD simulations guided by quantum calculations.
• Automatic differentiation in stochastic algorithms.