Bayesian E(3)-Equivariant Interatomic Potential with Iterative Restratification of Many-body Message Passing

TL;DR

This work tackles uncertainty quantification in interatomic potential models by developing Bayesian E(3)-equivariant MLPs that use iterative restratification of many-body interactions. Central to the approach is the joint energy–force negative log-likelihood loss, NLL_JEF, which models uncertainties in both energies and forces, enabling accurate predictions and calibrated uncertainty estimates. The authors implement the RACE architecture with eight-headed mean-variance estimators and evaluate multiple approximate Bayesian methods (DE, SWAG, IVON, LA) across QM9, PSB3, rMD17, and 3BPA benchmarks, demonstrating improved calibration, OOD detection, and data-efficient active learning via BALD. The results show competitive accuracy with state-of-the-art models while providing uncertainty-guided capabilities for active learning, OOD detection, and energy/force calibration, highlighting the practicality of Bayesian equivariant networks for large-scale atomistic simulations.

Abstract

Machine learning potentials (MLPs) have become essential for large-scale atomistic simulations, enabling ab initio-level accuracy with computational efficiency. However, current MLPs struggle with uncertainty quantification, limiting their reliability for active learning, calibration, and out-of-distribution (OOD) detection. We address these challenges by developing Bayesian E(3) equivariant MLPs with iterative restratification of many-body message passing. Our approach introduces the joint energy-force negative log-likelihood (NLL$_\text{JEF}$) loss function, which explicitly models uncertainty in both energies and interatomic forces, yielding superior accuracy compared to conventional NLL losses. We systematically benchmark multiple Bayesian approaches, including deep ensembles with mean-variance estimation, stochastic weight averaging Gaussian, improved variational online Newton, and laplace approximation by evaluating their performance on uncertainty prediction, OOD detection, calibration, and active learning tasks. We further demonstrate that NLL$_\text{JEF}$ facilitates efficient active learning by quantifying energy and force uncertainties. Using Bayesian active learning by disagreement (BALD), our framework outperforms random sampling and energy-uncertainty-based sampling. Our results demonstrate that Bayesian MLPs achieve competitive accuracy with state-of-the-art models while enabling uncertainty-guided active learning, OOD detection, and energy/forces calibration. This work establishes Bayesian equivariant neural networks as a powerful framework for developing uncertainty-aware MLPs for atomistic simulations at scale.

Figures (4)

• Figure 1: Overview of the proposed Bayesian E(3)-equivariant machine learning potential framework.(a) Two model variants: the Base model predicts only energies and forces, while the Mean-variance Estimator model additionally outputs predictive uncertainties. (b) The RACE model architecture, consisting of an embedding layer that initializes node features $\bm{A}_i^{(0)}$ from atomic numbers $\{Z_i\}$ and encodes local environments via pair-wise vectors $\bm{r}_{ij}$, angular edge features $e_{ij}$, and radial basis functions $\tilde{e}_\text{RBF}(r_{ij})$. (c) Each interaction layer updates node features in a ResNet-like scheme and predicts per-atom energies $E_i$ through a readout block; in Bayesian variants, this block also outputs uncertainties for energies and forces. (d) Bayesian neural network approaches, including Deep Ensemble, Stochastic Weight Averaging Gaussian, Improved Variational Online Newton, and Laplace Approximation, were used to obtain predictive distributions. (e) Downstream applications of uncertainty, including out-of-distribution detection, active learning with uncertainty-based sample selection, and model recalibration using reliability plots and confidence intervals.
• Figure 2: Predicted–empirical error scatter plots for the oBN25 dataset using a MVE, b DE, c SWAG, and d IVON. Blue dots correspond to liquid BN (ID) and orange dots to solid BN (OOD). Gray lines represent reference quantile bands from the ideal Gaussian distribution, with dashed lines marking the mode. Points aligned with the bands indicate well-calibrated uncertainty–error correlation, while points above or below indicate overconfidence or underconfidence, respectively.
• Figure 3: Calibration plots for energy predictions in a liquid phase and b solid phase, and for force predictions in c liquid phase and d solid phase of the oBN25 dataset. Methods compared: MVE, DE, SWAG, and IVON. Curves above the diagonal indicate underconfidence, whereas curves below indicate overconfidence.
• Figure 4: CE of Bayesian models before and after post-hoc recalibration on the oBN25 test dataset. Results are reported separately for energy and force in liquid-phase (ID) and solid-phase (OOD). Panels: a MVE, b DE, c SWAG, d IVON. Panels e and f illustrate 90% confidence intervals for RACE-DE on liquid BN before and after recalibration, respectively, showing how post-hoc adjustment corrects the raw Bayesian intervals.