Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Equivariant Evidential Deep Learning for Interatomic Potentials

Zhongyao Wang, Taoyong Cui, Jiawen Zou, Shufei Zhang, Bo Yan, Wanli Ouyang, Weimin Tan, Mao Su

TL;DR

This work tackles uncertainty quantification for vector-valued interatomic forces by developing e^2IP, an $SE(3)$-equivariant evidential framework that models the force uncertainty as a full $3\times3$ SPD covariance. The covariance is constructed via a Lie-algebra parameterization and matrix exponential to guarantee positive definiteness and rotation consistency, enabling a single forward pass with Hamiltonian-like uncertainty decomposition into aleatoric and epistemic components. The method couples with NIW priors to yield a multivariate Student-$t$ predictive distribution and introduces spectral stabilization to maintain numerical robustness. Empirical results across liquid water, DWCT, and silica-OOD benchmarks show superior uncertainty calibration, better error–uncertainty ranking, and favorable data efficiency compared to non-equivariant baselines and ensembles, with backbone-agnostic applicability and significant inference-time gains. This positions geometry-aware tensor uncertainty as a practical alternative to ensembles for uncertainty-aware molecular simulation and active-learning workflows.

Abstract

Uncertainty quantification (UQ) is critical for assessing the reliability of machine learning interatomic potentials (MLIPs) in molecular dynamics (MD) simulations, identifying extrapolation regimes and enabling uncertainty-aware workflows such as active learning for training dataset construction. Existing UQ approaches for MLIPs are often limited by high computational cost or suboptimal performance. Evidential deep learning (EDL) provides a theoretically grounded single-model alternative that determines both aleatoric and epistemic uncertainty in a single forward pass. However, extending evidential formulations from scalar targets to vector-valued quantities such as atomic forces introduces substantial challenges, particularly in maintaining statistical self-consistency under rotational transformations. To address this, we propose \textit{Equivariant Evidential Deep Learning for Interatomic Potentials} ($\text{e}^2$IP), a backbone-agnostic framework that models atomic forces and their uncertainty jointly by representing uncertainty as a full $3\times3$ symmetric positive definite covariance tensor that transforms equivariantly under rotations. Experiments on diverse molecular benchmarks show that $\text{e}^2$IP provides a stronger accuracy-efficiency-reliability balance than the non-equivariant evidential baseline and the widely used ensemble method. It also achieves better data efficiency through the fully equivariant architecture while retaining single-model inference efficiency.

Equivariant Evidential Deep Learning for Interatomic Potentials

TL;DR

This work tackles uncertainty quantification for vector-valued interatomic forces by developing e^2IP, an -equivariant evidential framework that models the force uncertainty as a full SPD covariance. The covariance is constructed via a Lie-algebra parameterization and matrix exponential to guarantee positive definiteness and rotation consistency, enabling a single forward pass with Hamiltonian-like uncertainty decomposition into aleatoric and epistemic components. The method couples with NIW priors to yield a multivariate Student- predictive distribution and introduces spectral stabilization to maintain numerical robustness. Empirical results across liquid water, DWCT, and silica-OOD benchmarks show superior uncertainty calibration, better error–uncertainty ranking, and favorable data efficiency compared to non-equivariant baselines and ensembles, with backbone-agnostic applicability and significant inference-time gains. This positions geometry-aware tensor uncertainty as a practical alternative to ensembles for uncertainty-aware molecular simulation and active-learning workflows.

Abstract

Uncertainty quantification (UQ) is critical for assessing the reliability of machine learning interatomic potentials (MLIPs) in molecular dynamics (MD) simulations, identifying extrapolation regimes and enabling uncertainty-aware workflows such as active learning for training dataset construction. Existing UQ approaches for MLIPs are often limited by high computational cost or suboptimal performance. Evidential deep learning (EDL) provides a theoretically grounded single-model alternative that determines both aleatoric and epistemic uncertainty in a single forward pass. However, extending evidential formulations from scalar targets to vector-valued quantities such as atomic forces introduces substantial challenges, particularly in maintaining statistical self-consistency under rotational transformations. To address this, we propose \textit{Equivariant Evidential Deep Learning for Interatomic Potentials} (IP), a backbone-agnostic framework that models atomic forces and their uncertainty jointly by representing uncertainty as a full symmetric positive definite covariance tensor that transforms equivariantly under rotations. Experiments on diverse molecular benchmarks show that IP provides a stronger accuracy-efficiency-reliability balance than the non-equivariant evidential baseline and the widely used ensemble method. It also achieves better data efficiency through the fully equivariant architecture while retaining single-model inference efficiency.
Paper Structure (62 sections, 49 equations, 6 figures, 4 tables, 1 algorithm)

This paper contains 62 sections, 49 equations, 6 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Preliminaries
  4. Problem setup.
  5. Equivariant backbone features.
  6. Equivariance requirements for mean and covariance
  7. Irrep structure of symmetric rank-2 tensors
  8. Methodology
  9. Overview.
  10. Multivariate Evidential Learning Framework
  11. Physically interpretable scale parameterization.
  12. Posterior predictive.
  13. Aleatoric and epistemic uncertainty decomposition.
  14. Constraints and stable scalar parameterization.
  15. Equivariant SPD covariance via Lie algebra
  16. ...and 47 more sections

Figures (6)

  • Figure 1: e$^{2}$IP framework. An SE(3)-equivariant GNN predicts the force mean and NIW evidential parameters, and constructs an equivariant SPD covariance $\Sigma_0$ (via a symmetric-matrix Lie-algebra head and matrix exponential) to produce tensor-valued uncertainty in a single forward pass.
  • Figure 2: Uncertainty diagnostics on liquid water.Top-left: reliability diagram (observed vs. nominal coverage; dashed line is ideal). Top-right: PIT CDF deviation $\hat{F}_U(p)-p$ (zero is ideal). Ensemble (raw) is strongly overconfident, while post-hoc calibration over-corrects and becomes under-confident; $\text{e}^2$IP stays closest to the ideal in both views. Bottom: alignment between force error and predicted epistemic scale, plotting Force RMSE against $\sqrt{\mathrm{tr}(\mathbf{U}_{\text{epi}})/3}$ (both in log10), shown as 2D density (counts in log scale) for $\text{e}^2$IP (left) and Ensemble (right). $\text{e}^2$IP exhibits a clearer monotonic relationship and higher Spearman correlation (0.78 vs. 0.61), indicating better error-aware ranking.
  • Figure 3: Data efficiency on MD22 subsets. Force MAE (top; lower is better) and Spearman rank correlation (bottom; higher is better) on Buckyball-Catcher (left) and DWCT (right) as a function of training-set usage (30/50/70/100%). We compare $\text{e}^2$IP with eIP xu2025evidential and Ours w/o equiv.. $\text{e}^2$IP achieves consistently lower MAE across data regimes, with the largest gains on DWCT, while maintaining stable rank correlation.
  • Figure 4: Spectral stabilization is necessary for numerical robustness. We track the batch-wise mean and min--max band of the eigenvalue ratio $\lambda_{\max}(\Sigma_0)/\lambda_{\min}(\Sigma_0)$ during early training. left (w/ damper): the ratio stays bounded and training remains stable. right (w/o damper): the ratio rapidly explodes and becomes erratic; negative values indicate that $\lambda_{\min}(\Sigma_0)$ turns non-positive due to finite-precision effects, after which Cholesky factorization fails (dashed line).
  • Figure 5: Linear--Tanh damping function with threshold $\tau=4$ and ceiling $c=5$. The mapping is linear near the origin and smoothly saturates for large inputs.
  • ...and 1 more figures