Scalable Bayesian Uncertainty Quantification for Neural Network Potentials: Promise and Pitfalls

TL;DR

This paper demonstrates scalable Bayesian uncertainty quantification for neural network potentials in molecular dynamics, highlighting the effectiveness of stochastic gradient MCMC and deep ensembles in capturing uncertainties reliably.

Contribution

It introduces scalable Bayesian UQ methods for NN potentials using SG-MCMC, compares with deep ensembles, and discusses their reliability and limitations for MD simulations.

Findings

SG-MCMC provides reliable uncertainty estimates for MD observables.

Deep ensembles achieve comparable results to SG-MCMC with less training.

Both methods effectively capture aleatoric and epistemic uncertainties.

Abstract

Neural network (NN) potentials promise highly accurate molecular dynamics (MD) simulations within the computational complexity of classical MD force fields. However, when applied outside their training domain, NN potential predictions can be inaccurate, increasing the need for Uncertainty Quantification (UQ). Bayesian modeling provides the mathematical framework for UQ, but classical Bayesian methods based on Markov chain Monte Carlo (MCMC) are computationally intractable for NN potentials. By training graph NN potentials for coarse-grained systems of liquid water and alanine dipeptide, we demonstrate here that scalable Bayesian UQ via stochastic gradient MCMC (SG-MCMC) yields reliable uncertainty estimates for MD observables. We show that cold posteriors can reduce the required training data size and that for reliable UQ, multiple Markov chains are needed. Additionally, we find that SG-MCMC and the Deep Ensemble method achieve comparable results, despite shorter training and less hyperparameter tuning of the latter. We show that both methods can capture aleatoric and epistemic uncertainty reliably, but not systematic uncertainty, which needs to be minimized by adequate modeling to obtain accurate credible intervals for MD observables. Our results represent a step towards accurate UQ that is of vital importance for trustworthy NN potential-based MD simulations required for decision-making in practice.