[Paper Review] Coherent energy and force uncertainty in deep learning force fields

Peter Bjørn Jørgensen, Jonas Busk+2 more|arXiv (Cornell University)|Dec 7, 2023

Machine Learning in Materials ScienceMaterials Science16 references8 citations

TL;DR

The paper links energy and force aleatoric uncertainty in deep learning force fields by modeling a correlated noise process, derives force variance expressions, and validates coherent uncertainty on molecular datasets using an equivariant message passing neural network ensemble.

ABSTRACT

In machine learning energy potentials for atomic systems, forces are commonly obtained as the negative derivative of the energy function with respect to atomic positions. To quantify aleatoric uncertainty in the predicted energies, a widely used modeling approach involves predicting both a mean and variance for each energy value. However, this model is not differentiable under the usual white noise assumption, so energy uncertainty does not naturally translate to force uncertainty. In this work we propose a machine learning potential energy model in which energy and force aleatoric uncertainty are linked through a spatially correlated noise process. We demonstrate our approach on an equivariant messages passing neural network potential trained on energies and forces on two out-of-equilibrium molecular datasets. Furthermore, we also show how to obtain epistemic uncertainties in this setting based on a Bayesian interpretation of deep ensemble models.

Motivation & Objective

• Motivate the need for coherent uncertainty estimates for energies and forces in ML force fields.
• Propose a noise model that links energy aleatoric uncertainty to force uncertainty through a spatially correlated process.
• Derive closed-form expressions for energy and force variances under colored noise.
• Demonstrate epistemic uncertainty via deep ensembles and Bayesian interpretation.
• Evaluate the approach on out-of-equilibrium molecular datasets (ANI-1x and Transition1x) using an equivariant message passing neural network.

Proposed method

• Model energy observations as E_obs = E_theta + rho_theta * eta with eta having differentiable autocorrelation.
• Derive force variance: Var(-∂E_obs/∂r) = gamma_hat * rho_theta^2 + (∂rho_theta/∂r)^2 (Eq. 8).
• Interpret gamma_hat as squared inverse length scale of the noise kernel and learn or set it as a hyperparameter.
• Use wide-sense stationary noise to simplify variance derivations (Appendix A).
• Adopt a Bayesian deep ensemble framework to obtain epistemic uncertainty via posterior sampling (Appendix B).
• Train equivariant PaiNN-based ensembles with calibrated aleatoric/epistemic components and evaluate on energy and force MAE/RMSE, NLL, and calibration metrics (Table 1).

Figure 1: Ensemble model (with 10 instances) fitted to energy (left) and force (right) observations with correlated noise.

Experimental results

Research questions

• RQ1 How can energy and force aleatoric uncertainties be coherently connected in ML force fields?
• RQ2 What is the impact of modeling correlated (colored) noise on the uncertainty estimates for energies and forces?
• RQ3 Can epistemic uncertainty be captured via deep ensembles in this coherent framework?
• RQ4 How does the proposed method perform on out-of-equilibrium molecular data (ANI-1x and Transition1x) compared to vanilla and white-noise baselines?
• RQ5 What are the practical calibration and predictive performance benefits of this coherent uncertainty approach?

Key findings

• A colored-noise energy model yields a closed-form force variance that couples with the energy uncertainty through a hyperparameter gamma_hat (Eq. 8).
• The force variance includes a term with the squared gradient of the variance output and a term proportional to rho_theta^2, enabling coherent energy/force uncertainty estimates.
• Epistemic uncertainty is obtainable from Bayesian deep ensembles, separating aleatoric from epistemic contributions (Appendix B).
• On ANI-1x and Transition1x, colored/noise-based ensembles achieve competitive MAE/RMSE with improved or calibrated uncertainty metrics (NLL, ENCE, RZV, CV) compared to vanilla and white-noise baselines (Table 1).
• Calibrated uncertainty can be improved by validation-set calibration without sacrificing prediction accuracy.

Figure 2: Reliability diagrams for models trained on the ANI-1x dataset.