Bayesian NequIP: Uncertainty in Atomistic Modeling

• The paper extends NequIP by incorporating a Bayesian framework with Gaussian priors and adaptive SGHMC, enabling uncertainty-aware predictions of energies and forces.
• It distinguishes between two constructions: a weight-space posterior sampling approach and a Gaussian-process surrogate integrated into Bayesian optimization for structure determination.
• Benchmark results indicate state-of-the-art accuracy and significant efficiency gains in atomistic simulations and adsorption search compared to deterministic methods.

Searching arXiv for the cited Bayesian NequIP papers and directly related NequIP context. Bayesian NequIP denotes Bayesian augmentations of the E(3)-equivariant interatomic potential NequIP that are used for uncertainty-aware atomistic modeling and for accelerated structure search. In one line of work, NequIP is turned into a stochastic neural network with a Gaussian prior over parameters, Gaussian likelihoods for energies and forces, and posterior sampling by an adaptive stochastic-gradient Hamiltonian Monte Carlo procedure, yielding uncertainty-aware interatomic force models with state-of-the-art accuracy and a good measure of uncertainty (Rensmeyer et al., 2023). In another line of work, NequIP is coupled to Bayesian optimization in the BOSS framework for adsorption-structure determination, where a Gaussian-process surrogate reduces the number of structure evaluations and NequIP reduces the cost of each evaluation (Jestilä et al., 2023).

1. NequIP as the common substrate

NequIP is an E(3)-equivariant graph neural network for interatomic potentials. It takes as input a molecular configuration $x = \{(r_i,z_i)\}_{i=1}^n$, where $r_i \in \mathbb{R}^3$ are nuclear coordinates and $z_i$ their atomic numbers, and outputs a scalar energy $E(x)$ and forces $F_i = -\nabla_{r_i}E(x)$. In the uncertainty-aware formulation, each atom carries features $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$, one vector per irrep up to $\ell_{\max}=2$, with $d_\ell=64$ for all $\ell$; radial functions use a learnable Bessel basis of size $8$ plus two hidden layers with $r_i \in \mathbb{R}^3$0 units and SiLU activation and a cutoff of $r_i \in \mathbb{R}^3$1 Å; five interaction blocks are stacked; and the final invariant features are mapped by an MLP to virtual atomic energies $r_i \in \mathbb{R}^3$2 with total energy $r_i \in \mathbb{R}^3$3 (Rensmeyer et al., 2023). In the adsorption-search formulation, the NequIP configuration uses $r_i \in \mathbb{R}^3$4 message-passing layers, feature multiplicity $r_i \in \mathbb{R}^3$5, maximum rotation order $r_i \in \mathbb{R}^3$6, and cutoff radius $r_i \in \mathbb{R}^3$7 Å, again relying on spherical harmonics and tensor-valued features to ensure E(3) equivariance (Jestilä et al., 2023).

The common architectural core should not obscure the fact that the literature uses the label “Bayesian NequIP” for two different Bayesian constructions. One places a posterior distribution over network parameters and treats energy and force prediction probabilistically. The other places a Gaussian-process prior over an unknown potential-energy surface in a Bayesian optimization loop and uses NequIP as the fast evaluator inside that loop. The shared element is NequIP’s equivariant energy–force model; the Bayesian layer differs substantially between the two usages.

2. Bayesian neural-network formulation of NequIP

In the uncertainty-aware force-field formulation, all network parameters, including the weights and biases of the base and projection layers and the two new standard-deviation MLPs, are collected into $r_i \in \mathbb{R}^3$8 and assigned the isotropic Gaussian prior

$r_i \in \mathbb{R}^3$9

The dataset is $z_i$0, assumed i.i.d., where each $z_i$1 contains the true DFT energy $z_i$2 and forces $z_i$3. Under $z_i$4, the network predicts means $z_i$5 and $z_i$6, and it also predicts standard deviations $z_i$7 and a single scalar $z_i$8 for all three components of every atomic force. These standard deviations are produced by small MLP heads acting on the invariant features from the NequIP base. The likelihood is

$z_i$9

and the posterior density is

$E(x)$0

Equivalently, the negative-log potential is

$E(x)$1

and posterior sampling targets $E(x)$2 (Rensmeyer et al., 2023).

This construction is explicitly motivated by the use of Bayesian neural networks for modeling uncertainty, active learning, and incorporating prior physical knowledge. The NequIP backbone remains the mechanism that enforces equivariance and produces energies and forces, while the Bayesian formulation adds a probabilistic layer over both the parameters and the predictive outputs. In this sense, Bayesian NequIP is not merely an ensemble heuristic but a stochastic neural network model with a defined prior, likelihood, posterior, and predictive density.

3. Adaptive SGHMC posterior sampling and training protocol

A central obstacle in applying Bayesian neural networks to interatomic force modeling is posterior sampling for state-of-the-art architectures. Standard SGHMC introduces an auxiliary momentum variable $E(x)$3 and simulates

$E(x)$4

$E(x)$5

with diagonal mass $E(x)$6 and friction $E(x)$7. In NequIP, however, different parameter groups exhibit gradient scales differing by orders of magnitude, so a single fixed $E(x)$8 or stepsize $E(x)$9 causes some parameters to move too slowly while others explode. The proposed remedy is an AMSGrad-based diagonal mass adaptation:

$F_i = -\nabla_{r_i}E(x)$0

$F_i = -\nabla_{r_i}E(x)$1

$F_i = -\nabla_{r_i}E(x)$2

where the maximum is elementwise. During the first $F_i = -\nabla_{r_i}E(x)$3 steps, regarded as burn-in, $F_i = -\nabla_{r_i}E(x)$4 is updated; afterward it is frozen as $F_i = -\nabla_{r_i}E(x)$5, yielding a true SGHMC chain. The algorithm uses minibatches $F_i = -\nabla_{r_i}E(x)$6 of size $F_i = -\nabla_{r_i}E(x)$7 and the stochastic-gradient estimate

$F_i = -\nabla_{r_i}E(x)$8

The update rule given in the pseudocode is

$F_i = -\nabla_{r_i}E(x)$9

$v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$0

with $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$1, $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$2, $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$3, and $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$4 (Rensmeyer et al., 2023).

The reported training protocol fixes $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$5, $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$6, and $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$7. For the RMD17 datasets, the batch size is $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$8 and the learning rate $v_i^{\ell,p} \in \mathbb{R}^{d_\ell}$9 is decayed exponentially from $\ell_{\max}=2$0 to $\ell_{\max}=2$1 over the first $\ell_{\max}=2$2 steps of burn-in; after step $\ell_{\max}=2$3, the mass is frozen and a cosine cyclic schedule is used,

$\ell_{\max}=2$4

with $\ell_{\max}=2$5 and $\ell_{\max}=2$6, sampling one $\ell_{\max}=2$7 every cycle and obtaining $\ell_{\max}=2$8 samples in about $\ell_{\max}=2$9 extra steps. For the PEDOT dataset, the scheme is identical except that $d_\ell=64$0, burn-in lasts $d_\ell=64$1 steps, and the mass is frozen after $d_\ell=64$2 steps. The paper reports that $d_\ell=64$3 samples are a good trade-off and that a single chain suffices to produce approximately independent samples.

4. Predictive uncertainty, calibration, and benchmark behavior

Given posterior samples $d_\ell=64$4, the predictive density for a new configuration is approximated by

$d_\ell=64$5

with point estimates

$d_\ell=64$6

For a force component $d_\ell=64$7, the total predictive variance is

$d_\ell=64$8

The first term is identified as aleatoric uncertainty, and the second as epistemic uncertainty; analogous formulas hold for the energy (Rensmeyer et al., 2023).

The empirical evaluation uses PEDOT conducting-polymer molecular dynamics and Revised MD17. The PEDOT dataset contains $d_\ell=64$9 training and $\ell$0 validation structures for chains of length $\ell$1, with testing also on length $\ell$2 and zero training on that length. The Revised MD17 setup uses $\ell$3 training structures, $\ell$4 validation structures, and approximately $\ell$5 test structures per molecule. Metrics include MAE-F and RMSE-F in kcal/(mol·Å), MAE-E and RMSE-E in kcal/mol, mean log-likelihood under the predictive distribution, ROC AUC for outlier detection with threshold $\ell$6 kcal/(mol·Å), and calibration via Expected Calibration Error normalized by bin width, denoted NECE. The main baseline is Dropout-based NequIP with a single dropout layer of rate $\ell$7 trained with AMSGrad on identical data splits. On PEDOT, MAE-F is approximately $\ell$8–$\ell$9 kcal/(mol·Å), improving slightly from $8$0 to $8$1, while energy MAE is about $8$2–$8$3 kcal/mol. On length-$8$4 PEDOT, NECE improves from $8$5 for $8$6 to $8$7 for $8$8, and $8$9 improves from $r_i \in \mathbb{R}^3$00 to $r_i \in \mathbb{R}^3$01. On RMD17, force MAE is comparable to Dropout on easy molecules but $r_i \in \mathbb{R}^3$02–$r_i \in \mathbb{R}^3$03 better on aspirin and malonaldehyde; $r_i \in \mathbb{R}^3$04 is consistently better by $r_i \in \mathbb{R}^3$05–$r_i \in \mathbb{R}^3$06 nats; ROC AUC for outlier detection at $r_i \in \mathbb{R}^3$07 is $r_i \in \mathbb{R}^3$08–$r_i \in \mathbb{R}^3$09 versus $r_i \in \mathbb{R}^3$10–$r_i \in \mathbb{R}^3$11 for Dropout; and NECE is $r_i \in \mathbb{R}^3$12–$r_i \in \mathbb{R}^3$13, indicating some overconfidence but still lower than Dropout’s $r_i \in \mathbb{R}^3$14–$r_i \in \mathbb{R}^3$15. The paper also reports that the Bayesian procedure requires about twice the GPU hours of deterministic NequIP to reach convergence, while inference with $r_i \in \mathbb{R}^3$16 Monte Carlo predictions costs approximately $r_i \in \mathbb{R}^3$17 times a single NequIP forward pass plus small MLP overhead for the standard deviations, and in most workflows DFT data generation dominates runtime.

5. Bayesian optimization with NequIP for adsorption structure determination

In the adsorption-search usage, Bayesian NequIP refers to a protocol that combines Bayesian optimization structural inference with NequIP. The decision variables are split into conformer-search internal coordinates $r_i \in \mathbb{R}^3$18 and adsorption-search variables $r_i \in \mathbb{R}^3$19, consisting of the Cartesian translation of the molecule above the surface and three Euler roll-pitch-yaw angles. A Gaussian-process prior is placed on the unknown potential-energy surface $r_i \in \mathbb{R}^3$20. After $r_i \in \mathbb{R}^3$21 evaluations at points $r_i \in \mathbb{R}^3$22, the posterior predictive mean and variance are

$r_i \in \mathbb{R}^3$23

$r_i \in \mathbb{R}^3$24

where $r_i \in \mathbb{R}^3$25, $r_i \in \mathbb{R}^3$26, $r_i \in \mathbb{R}^3$27 is the noise or jitter hyperparameter, and $r_i \in \mathbb{R}^3$28. The acquisition function may be Expected Improvement,

$r_i \in \mathbb{R}^3$29

with $r_i \in \mathbb{R}^3$30, or Upper-Confidence Bound,

$r_i \in \mathbb{R}^3$31

The BO loop is initialized with a small number, for example $r_i \in \mathbb{R}^3$32, of random or space-filling samples, then iterates by fitting $r_i \in \mathbb{R}^3$33, choosing $r_i \in \mathbb{R}^3$34, evaluating the expensive objective, augmenting the dataset, and stopping when a budget such as $r_i \in \mathbb{R}^3$35 points is reached or a convergence criterion is met (Jestilä et al., 2023).

For adsorption searches, the objective is the binding energy

$r_i \in \mathbb{R}^3$36

which is minimized. NequIP serves as the machine-learning interatomic potential that replaces DFT as the expensive oracle at active-learning steps. In this protocol, the NequIP model uses $r_i \in \mathbb{R}^3$37 message-passing layers, feature multiplicity $r_i \in \mathbb{R}^3$38, $r_i \in \mathbb{R}^3$39, and $r_i \in \mathbb{R}^3$40 Å. Training minimizes

$r_i \in \mathbb{R}^3$41

with equal weighting of the energy and force terms; Adam is used with learning rate $r_i \in \mathbb{R}^3$42 and EMA decay $r_i \in \mathbb{R}^3$43. The network outputs atomic contributions $r_i \in \mathbb{R}^3$44 such that $r_i \in \mathbb{R}^3$45, and forces are obtained by analytic differentiation, $r_i \in \mathbb{R}^3$46. Between BO iterations, local relaxations of candidate minima are carried out with BFGS in ASE using the NequIP calculator until $r_i \in \mathbb{R}^3$47 eV/Å. To address conformational flexibility, the protocol adopts a building-block strategy: gas-phase conformers are located first, and where full conformational mapping is infeasible, external conformer search tools such as CREST metadynamics are used; the resulting conformers are then frozen as rigid building blocks and only the $r_i \in \mathbb{R}^3$48D adsorption pose is optimized with BOSS.

6. Performance claims, limitations, and recurrent points of confusion

For adsorption-structure determination, the reported performance gains are substantial in terms of search cost. DFT-only global adsorption searches for flexible xylotetraose are described as intractable, requiring more than $r_i \in \mathbb{R}^3$49 years of CPU time for $r_i \in \mathbb{R}^3$50 structures. BOSS with DFT reduces the number of required DFT evaluations by two to three orders of magnitude, for example from approximately $r_i \in \mathbb{R}^3$51 to approximately $r_i \in \mathbb{R}^3$52–$r_i \in \mathbb{R}^3$53. NequIP accelerates each energy and force evaluation by approximately $r_i \in \mathbb{R}^3$54–$r_i \in \mathbb{R}^3$55 over DFT. Combined, the protocol identifies known DFT global minima at approximately $r_i \in \mathbb{R}^3$56–$r_i \in \mathbb{R}^3$57 lower wall time, enabling $r_i \in \mathbb{R}^3$58 pose samples instead of $r_i \in \mathbb{R}^3$59 by DFT at comparable accuracy (Jestilä et al., 2023).

A recurrent source of confusion is that the Bayesian ingredient is different in the two main usages. In the uncertainty-aware force-modeling work, the Bayesian object is the posterior over NequIP parameters and the predictive density is a Monte Carlo average over sampled networks. In the adsorption-search work, the Bayesian object is the Gaussian-process surrogate over adsorption configurations, while NequIP functions as the fast machine-learning interatomic potential inside the optimization loop. Another common misunderstanding is to equate uncertainty estimation with perfect calibration. The uncertainty-aware force-modeling results explicitly report slight overconfidence tails, even though calibration as measured by NECE is better than the Dropout baseline. A further limitation concerns conformational flexibility: the adsorption-search paper states that the applicability of Bayesian optimization for the conformational analysis of the more flexible xylotetraose molecule is restricted by the sample-complexity bottleneck, and that this bottleneck can be bypassed with external conformer search tools before lower-dimensional adsorption optimization (Rensmeyer et al., 2023).

These two usages therefore occupy different positions in atomistic-learning workflows. One addresses posterior uncertainty within the interatomic force model itself; the other addresses sample efficiency in global structure determination. Taken together, they show how Bayesian methodology can enter NequIP-based atomistic modeling either through weight-space posterior inference or through Gaussian-process-guided search over molecular and adsorption configurations.