Bayesian Neural Networks versus deep ensembles for uncertainty quantification in machine learning interatomic potentials

Technical Summary: Bayesian Neural Networks Versus Deep Ensembles for Uncertainty Quantification in Machine Learning Interatomic Potentials

Problem Statement
Machine Learning Interatomic Potentials (MLIPs) have become essential for atomistic modeling, offering a computationally efficient alternative to ab initio methods like Density Functional Theory (DFT). However, standard deep learning architectures are deterministic and lack inherent mechanisms to quantify predictive uncertainty. This limitation is critical in materials science, where models are frequently applied to out-of-distribution data. Without reliable uncertainty estimates, simulations risk producing misleading results, and active learning strategies for model refinement lack guidance. While Deep Ensembles (DE) and Bayesian Neural Networks (BNNs) are established strategies for addressing this, their comparative performance in the specific context of MLIPs—particularly regarding joint energy and force predictions across varying data regimes—requires systematic evaluation.

Methodology
The authors implemented and compared two primary uncertainty quantification approaches within the ænet-PyTorch framework:

1. Deep Ensembles (DE): Following the formulation by Lakshminarayanan et al., five independently trained neural networks were initialized with different random seeds. Uncertainty was quantified as the sample standard deviation of the ensemble's predictions.
2. Variational Bayesian Neural Networks (VBNN): The authors implemented BNNs using Variational Inference (VI) via the Pyro probabilistic programming framework. Two specific variational guides were tested:
   • Local Reparametrization Trick (LRT): Utilizing a mean-field approximation with a fully factorized Gaussian posterior, optimized using the Bayes by Backprop algorithm with LRT to reduce gradient variance.
   • Radial Guide (RAD): A non-mean-field approach introducing radial transformations to capture global correlations among weights, optimized via standard Monte Carlo sampling of the Evidence Lower Bound (ELBO).

Experimental Setup
The study evaluated these methods on two distinct datasets:

• TiO2 Dataset: 7,815 structures representing various phases of titanium dioxide. Models were trained on both the full dataset (High-Data regime) and a 20% subset (Low-Data regime) to assess data efficiency. Both energy-only and joint energy-force training were performed.
• QM7 Dataset: 7,165 small organic molecules used to test generalizability across chemically diverse systems. Models were trained in the high-data regime on atomization energies only.

Evaluation Metrics
Performance was assessed using a comprehensive suite of metrics:

• Accuracy: Mean Absolute Error (MAE), Root Mean Squared Error (RMSE), and Negative Log-Likelihood (NLL).
• Uncertainty Quality: Root Mean Squared Calibration Error (RMSCE), Sharpness (average predictive uncertainty), Spearman rank correlation ($r_s$) between uncertainty and error, and an "Overlap Score" measuring the alignment between high-uncertainty predictions and high-error regions.

Key Results

1. Predictive Accuracy:

   • In the High-Data regime (TiO2), DE consistently outperformed VBNNs, achieving the lowest MAE (6 meV/atom) and RMSE (9 meV/atom), as well as the most negative NLL (-4.2), indicating superior accuracy and calibration.
   • In the Low-Data regime, DE maintained its superiority in accuracy (MAE 35 meV/atom). However, LRT showed competitive accuracy (MAE 54 meV/atom) and a lower RMSE (259 meV/atom) than DE, suggesting better robustness to outliers. RAD performed poorly in low-data settings.
   • On the QM7 dataset, a reversal occurred: BNNs (LRT and RAD) achieved lower energy RMSE (13.0–13.3 meV/atom) than DE (15.2 meV/atom), suggesting the variational prior provides beneficial regularization for chemically diverse spaces.

2. Uncertainty Quantification:

   • Calibration: DE generally provided the best-calibrated uncertainties (lowest RMSCE) in the TiO2 experiments. In the low-data regime, RAD surprisingly maintained excellent calibration (RMSCE 0.05), comparable to its high-data performance, while LRT and DE saw degraded calibration.
   • Sharpness: DE produced the sharpest (most confident) predictions. LRT tended to produce overly broad uncertainty intervals (high sharpness values), particularly in low-data regimes.
   • Correlation and Overlap: DE demonstrated the strongest monotonic correlation ($r_s$) between predicted uncertainty and actual error in the low-data regime (0.45), outperforming both LRT (0.28) and RAD (0.19). DE also achieved the highest Overlap Score (55%), indicating it most effectively identified high-error regions.

3. Force Predictions:

   • In the low-data regime, VBNNs (LRT and RAD) achieved significantly lower force RMSE than DE, likely due to the regularization provided by the prior. However, force calibration remained poor across all models, with only DE showing reasonable calibration in the low-data regime.

4. Computational Efficiency:

   • Training: DE has a lower per-epoch cost but requires training five independent models. Total training time for DE was comparable to LRT in the high-data regime but slightly higher in the low-data regime.
   • Inference: DE is significantly faster at inference (0.03 s) compared to VBNNs (0.64–0.71 s), as BNNs require 20 Monte Carlo samples per prediction, including automatic differentiation for forces.

Significance and Conclusions
The paper concludes that while Deep Ensembles are theoretically less grounded than BNNs, they empirically outperform variational BNNs in most metrics regarding accuracy, calibration, and inference efficiency, particularly when the chemical space is relatively narrow (TiO2). DE is recommended as the preferred choice when reliable uncertainty quantification is the priority and implementation simplicity is desired.

However, the study highlights that VBNNs, specifically the LRT variant, offer a competitive trade-off and may be superior in data-scarce or chemically diverse regimes (as seen in QM7) where the variational prior acts as an effective regularizer. The authors note that BNNs are more sensitive to random initialization and require multiple training runs to ensure robustness. The work provides practical guidance for selecting uncertainty-aware MLIPs, suggesting that the choice between DE and VBNNs should depend on the specific data regime and the chemical diversity of the target system.

Contributions

• Introduction of a GPU-accelerated VBNN implementation for MLIPs within the ænet-PyTorch framework, supporting joint energy-force training.
• A systematic comparison of DE, LRT, and RAD across high- and low-data regimes on both inorganic (TiO2) and organic (QM7) datasets.
• Demonstration that while DE is generally superior, the relative performance of BNNs can reverse in chemically diverse datasets, challenging the universality of ensemble superiority in all contexts.
• Provision of a new "Overlap Score" metric to evaluate the utility of uncertainty estimates for active learning scenarios.