[Paper Review] Single-model uncertainty quantification in neural network potentials does not consistently outperform model ensembles

Aik Rui Tan, Shingo Urata+3 more|arXiv (Cornell University)|May 2, 2023

Machine Learning in Materials Science73 references4 citations

TL;DR

This study evaluates single-model uncertainty quantification (UQ) methods—mean-variance estimation, deep evidential regression, and Gaussian mixture models—against neural network ensembles for neural network interatomic potentials (NNIPs). Despite lower computational cost, none of the single-model approaches consistently outperformed ensembles across datasets and metrics, with ensembling remaining superior for generalization and robustness in active learning loops.

ABSTRACT

Neural networks (NNs) often assign high confidence to their predictions, even for points far out-of-distribution, making uncertainty quantification (UQ) a challenge. When they are employed to model interatomic potentials in materials systems, this problem leads to unphysical structures that disrupt simulations, or to biased statistics and dynamics that do not reflect the true physics. Differentiable UQ techniques can find new informative data and drive active learning loops for robust potentials. However, a variety of UQ techniques, including newly developed ones, exist for atomistic simulations and there are no clear guidelines for which are most effective or suitable for a given case. In this work, we examine multiple UQ schemes for improving the robustness of NN interatomic potentials (NNIPs) through active learning. In particular, we compare incumbent ensemble-based methods against strategies that use single, deterministic NNs: mean-variance estimation, deep evidential regression, and Gaussian mixture models. We explore three datasets ranging from in-domain interpolative learning to more extrapolative out-of-domain generalization challenges: rMD17, ammonia inversion, and bulk silica glass. Performance is measured across multiple metrics relating model error to uncertainty. Our experiments show that none of the methods consistently outperformed each other across the various metrics. Ensembling remained better at generalization and for NNIP robustness; MVE only proved effective for in-domain interpolation, while GMM was better out-of-domain; and evidential regression, despite its promise, was not the preferable alternative in any of the cases. More broadly, cost-effective, single deterministic models cannot yet consistently match or outperform ensembling for uncertainty quantification in NNIPs.

Motivation & Objective

• To evaluate whether single-model uncertainty quantification (UQ) methods can outperform neural network ensembles in training robust neural network interatomic potentials (NNIPs).
• To assess the effectiveness of various UQ techniques—mean-variance estimation, deep evidential regression, and Gaussian mixture models—across different data regimes: in-domain interpolation and out-of-domain generalization.
• To determine the suitability of single deterministic NNs for active learning in atomistic simulations, where uncertainty quantification is critical for avoiding unphysical structures and biased dynamics.
• To identify the limitations of promising single-model UQ approaches in comparison to the computationally heavier ensemble-based method.

Proposed method

• The study employs three benchmark datasets: rMD17 (in-domain interpolation), ammonia inversion (moderate extrapolation), and bulk silica glass (highly extrapolative generalization).
• For each dataset, multiple UQ methods are applied: mean-variance estimation (MVE), deep evidential regression, Gaussian mixture models (GMM), and neural network ensembles.
• Uncertainty is evaluated using metrics that correlate model error with predicted uncertainty across diverse test sets, including adversarial sampling to probe out-of-distribution performance.
• Active learning loops are implemented using UQ scores to guide data acquisition, measuring how well each method identifies informative, high-error regions.
• The GMM approach uses latent space representations from trained NNs to fit mixture models post-training, enabling uncertainty estimation without architectural changes.
• Deep evidential regression models uncertainty via a Dirichlet distribution over predictive outputs, with parameters optimized end-to-end using a proper scoring rule.

Figure 1: Illustration of the uncertainty quantification (UQ) methods used. $x$ denotes the input to the neural networks (NNs), while $y$ is the predicted property. In the case of NNIPs, $x$ generally represents the positions and atomic numbers of the input structure, whereas $y$ is the energy and/o

Experimental results

Research questions

• RQ1 Can single-model UQ methods such as mean-variance estimation, deep evidential regression, or Gaussian mixture models consistently outperform neural network ensembles in uncertainty quantification for NNIPs?
• RQ2 How do different UQ methods perform across varying degrees of distribution shift, from in-domain interpolation to highly extrapolative generalization?
• RQ3 To what extent can single deterministic NNs with UQ capabilities achieve robustness comparable to ensembles in active learning loops for atomistic simulations?
• RQ4 Does the numerical stability and generalization capability of deep evidential regression make it a viable alternative to ensembles in NNIP training?
• RQ5 How well do UQ methods rank uncertainty for out-of-distribution frames, particularly in challenging systems like silica glass?

Key findings

• Neural network ensembles consistently outperformed all single-model UQ methods across all datasets and evaluation metrics, particularly in generalization and robustness during active learning.
• Mean-variance estimation (MVE) performed well only in in-domain interpolation tasks, failing to systematically identify high-uncertainty regions in out-of-distribution regimes.
• Gaussian mixture models (GMM) showed superior performance for out-of-domain generalization, especially in the silica glass dataset, but were less effective in interpolative regimes.
• Deep evidential regression, despite theoretical promise, exhibited numerical instability during optimization and did not outperform ensembles in any metric or dataset.
• No single single-model UQ method matched the generalization capability of ensembles, indicating that cost-effective deterministic models cannot yet consistently replace ensembles for robust NNIP training.
• The study confirms that ensemble-based UQ remains the most reliable method for identifying out-of-distribution frames via adversarial sampling, even though it is computationally expensive.

Figure 2: (a) Hexbin plots showing (predicted) uncertainties versus squared errors of atomic forces in all molecules of all 5-fold test sets from the rMD17 data set for each considered UQ method. Note that since the uncertainties (-NLL) for the GMM method contain negative values, all uncertainties a