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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Raja, Sanjeev, Amin, Ishan, Pedregosa, Fabian, Krishnapriyan, Aditi S.
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866915171822206976
author Raja, Sanjeev
Amin, Ishan
Pedregosa, Fabian
Krishnapriyan, Aditi S.
author_facet Raja, Sanjeev
Amin, Ishan
Pedregosa, Fabian
Krishnapriyan, Aditi S.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2402_13984
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators
Raja, Sanjeev
Amin, Ishan
Pedregosa, Fabian
Krishnapriyan, Aditi S.
Machine Learning
Disordered Systems and Neural Networks
Materials Science
Chemical Physics
Computational Physics
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.
title Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators
topic Machine Learning
Disordered Systems and Neural Networks
Materials Science
Chemical Physics
Computational Physics
url https://arxiv.org/abs/2402.13984