Knowing when to trust machine-learned interatomic potentials Shams Mehdi1 , Ilkwon Cho1 , Olexandr Isayev1,2,3*
arXiv:2605.00640v1 [cs.LG] 1 May 2026
1
Department of Chemistry, Mellon College of Science, Carnegie Mellon University, Pittsburgh, Pennsylvania, 15213, USA. 2 Department of Materials Science and Engineering, Carnegie Mellon University, Pittsburgh, Pennsylvania, 15213, USA. 3 Ray and Stephanie Lane Computational Biology Department, School of Computer Science, Carnegie Mellon University, Pittsburgh, Pennsylvania, 15213, USA.
*Corresponding author(s). E-mail(s): [email protected]; Contributing authors: [email protected]; [email protected]; Abstract Prevailing machine-learned interatomic potential (MLIP) uncertaintyquantification methods rely on ensembles of independently trained backbones. These methods scale unfavorably with foundation-scale MLIPs, and their member-disagreement signals correlate weakly with per-molecule prediction error. Here we probe the frozen per-atom representations of a pretrained MLIP with a compact discriminative classifier, recasting MLIP uncertainty quantification as selective classification rather than error regression. The resulting method, PROBE (Post-hoc Reliability frOm Backbone Embeddings), produces a perprediction reliability probability that monotonically tracks actual error without modification to the underlying model. Across large held-out evaluation sets and two structurally distinct MLIP architectures, PROBE outperforms ensemble disagreement as a binary reliability signal, which strengthens with the expressiveness of the backbone representation, implying a favorable scaling trajectory toward foundation-scale MLIPs. Multi-head self-attention additionally yields per-atom importance maps, providing chemically interpretable diagnostics at no additional computational cost. PROBE is post-hoc and architecture-agnostic, and is directly deployable on any MLIP that exposes per-atom representations.
1
Keywords: Machine-learned interatomic potentials, Uncertainty quantification, Explainable AI
1 Introduction Machine-learned interatomic potentials (MLIPs) are now widely used in computational chemistry and materials science [1–9]. They map atomic configurations to energies and forces using quantum mechanical reference data, achieving near density functional theory [10, 11] (DFT) accuracy at a fraction of the cost. This enables largescale atomistic simulations [12] and accelerated materials discovery that would be intractable with first-principles methods. The field has advanced through expressive model architectures, from local descriptors toward equivariant and attention-based message passing [4, 13], curation techniques spanning vast chemical spaces [6, 9], foundation model training [9, 14, 15]; broadening MLIP applicability across drug discovery [16], materials design, and automated reaction exploration. Despite these successes, a persistent problem remains: how can a user know when to trust a given MLIP prediction? We address this question directly, by probing the internal representations of a trained MLIP with a lightweight classifier that learns to distinguish reliable from unreliable predictions. Uncertainties associated with MLIP predictions can be categorized into [17]: aleatoric uncertainty, arising from irreducible noise in the data (e.g., DFT numerical errors or basis set incompleteness), and epistemic uncertainty, arising from limited training data, distribution shift or unconverged training of the model. For MLIPs, the practically actionable source is epistemic: failures are systematic and typically tied to underrepresented chemical environments. Even within their training domain, MLIPs do not learn all chemical environments equally well; certain local bonding patterns, rare element combinations, or strained geometries are systematically harder to fit. Long-range intermolecular interactions are fundamentally harder to extrapolate than intramolecular ones. Examples include π −π stacking, charge transfer, and halogen bonding driven by anisotropic σ -holes (e.g., in Br and I). These physics-driven limitations produce large prediction errors that can corrupt high-throughput screening pipelines, active learning loops, and geometry optimization workflows before being noticed [18–20]. Perez et al. [21] recently showed this concretely for the MACE-MPA-0 foundation MLIP on tungsten defects: ensemble disagreement captures epistemic variance but misses model misspecification. Against this backdrop, uncertainty quantification (UQ) has become an explicit research priority for the atomistic-ML community [22, 23]. The dominant approach to MLIP UQ has been the committee or deep ensemble [19, 24]: several independently trained models whose output variance serves as a proxy for uncertainty [18, 19]. Committee potentials have driven much of the progress in uncertainty-guided active learning [19, 25–27], including uncertainty-driven molecular dynamics [28] and hyperactive learning strategies [29] that bias sampling toward high-uncertainty regions. However, Kurniawan et al. [30] recently showed that ensemble uncertainty in out-of-distribution regimes can plateau or even decrease as predictive errors grow. For equivariant neural network potentials specifically, Lu et 2
al. [31] found that ensemble uncertainty estimates are overconfident and poorly predictive of actual errors, suggesting the problem is not merely a matter of method but of a fundamental mismatch between ensemble disagreement and true error. Additionally, there are practical drawbacks of running N models which costs N times as much in both compute and memory. Recent UQ work on foundation-scale MLIPs that scale to tens of millions of parameters, therefore abandons full ensembling in favorof readout-layer or shared-backbone alternatives [32]. The community has pursued a range of approaches involving single model MLIPs to avoid the ensemble overhead. Monte Carlo Dropout [33] applied to MLIPs [34] provides cheap stochastic estimates, though calibration often degrades under distribution shift. Gaussian process [35] force fields [36] such as FLARE [25, 37] offer principled Bayesian posteriors but scale poorly to large training sets. Bayesian neural network potentials [38] and Bayesian learned interatomic potentials (BLIPs) [39] can yield well-calibrated uncertainties but require modifying or re-training the backbone with variational inference. Evidential deep learning for interatomic potentials [40] and loss-trajectory analysis [41] represent newer single-network approaches that avoid additional training but require non-standard training procedures. Conformal prediction has recently been adapted for MLIPs [42], giving distribution-free coverage guarantees without backbone modification, which produces prediction intervals rather than reliability signals. Multi-head committee models for MACE [43] share a backbone to reduce cost, but still require training and maintaining additional output heads. A closely related line of work avoids ensembling by extracting uncertainty signals from a single backbone’s internal representations. Janet et al. [44] and, more recently, Musielewicz et al. [45] used latent-space distance to the training set as a quantitative error metric, the latter on a frozen equivariant GNN potential. Bigi et al. [46] formalised this intuition as a post-hoc prediction rigidity scalar. It measures how sensitive a trained network’s output is to its training data, without additional sampling. Kellner and Ceriotti [47] augmented a shared backbone with lightweight shallow-ensemble heads, recovering much of a full ensemble’s signal at a small fraction of the cost. A previous comparative study [20] concluded that, across this diverse landscape, no single-model UQ method consistently outperforms ensembles for MLIP error estimation. Here, we argue that part of the difficulty with MLIP UQ comes from an overly ambitious problem formulation. Most existing methods try to predict the magnitude of the error, a nonlinear regression problem over a heavy-tailed distribution. We instead ask a simpler question: is this prediction reliable? This binary framing captures most of the practical information a UQ system needs to provide. Downstream decisions are themselves binary: accept a geometry optimization, include a molecule in a training set, or trigger a DFT recalculation. We present PROBE (Post-hoc Reliability frOm Backbone Embeddings), a lightweight framework built on this idea. PROBE attaches a small classifier to the frozen internal representations of a trained MLIP. The classifier is trained to distinguish reliable from unreliable predictions using per-atom embeddings from the backbone, without modifying the backbone itself. We hypothesize that a trained MLIP’s internal representations already encode whether a chemical environment was
3
Fig. 1: PROBE architecture overview. (a) A frozen, pre-trained MLIP processes atomic coordinates (R, Z) and exposes per-atom latent representations {hi } ∈ Rd alongside molecular-level predictions. PROBE attaches to these frozen representations without modifying the underlying MLIP. (b) PROBE architecture. An atom encoder projects per-atom features into a fixed-dimensional space. A molecule encoder applies multi-head self-attention followed by mean and max pooling to produce a fixed-size molecular embedding. A classifier MLP outputs P (reliable) and P (unreliable). (c) Schematic illustration of the multi-head attention mechanism: Different heads can, in principle, attend to distinct structural contexts within a molecule (e.g., strained bonds, polar regions, steric clashes, or conjugated systems), producing complementary contributions to the molecular embedding. Specific head specialization shown is not a claim of the present work. (d) Per-atom importance scores derived from the attention matrix by summing attention weights received by each atom (Eq. 1), enabling identification of the specific atoms driving an unreliable prediction.
well-covered during training, and that a simple classifier can exploit this signal at inference time. This reframing connects MLIP UQ to the literature on selective prediction and confidence-based abstention [48–51], in which a model identifies inputs whose predictions should be trusted versus deferred. Representation-based reliability scoring is established for out-of-distribution detection in computer vision, language modeling[52– 55], and scientific ML [56], but has not been applied to MLIP energetics, where ensemble disagreement remains the default. Methodologically, PROBE is a probing classifier in the sense of Alain and Bengio [57]: a small auxiliary model trained on a frozen backbone’s intermediate features, exploiting the observation that a trained
4
backbone’s internal representations already encode whether a chemical environment was well-covered during training. Figure 1 summarizes PROBE workflow, which has the following key features: 1. Internal MLIP representations for atomic embedding: PROBE operates directly on the per-atom embeddings produced by a single, trained MLIP backbone, capturing rich chemical environment information without any modification to the underlying model. 2. Multi-head self-attention for molecular embedding: Atom-level embeddings are processed through a self-attention module to construct a molecular embedding, enabling PROBE to capture both local and global chemical context while naturally handling variable-sized molecules. Here, the attention mechanism handles the nonadditive nature of prediction error, with different heads learning to specialize on distinct failure modes present in the training data. 3. Data-driven class boundaries: The reliability threshold is defined as a percentile of the training error distribution; we use the 50th percentile (balanced classes) throughout this work. In Section 3 we discuss why moving off the median degrades performance, and why raising the probability cutoff on a balanced model is the better way to target severe outliers. 4. Interpretable atom importance: Attention weights provide chemically meaningful, per-atom reliability maps without any additional computation, identifying which atoms are responsible for anomalously large prediction errors. 5. Architecture agnosticism: PROBE requires only access to per-atom internal representations and can therefore be deployed as a post-hoc module on most modern MLIP architectures, demonstrated here for both AIMNet2 [6] and MACE [5]. We tested PROBE on AIMNet2 and MACE because they use contrasting representation families: chemically informed atom-in-molecule vectors for AIMNet2 [6], and equivariant graph-based node embeddings for MACE [5, 14], which lets us evaluate whether PROBE generalizes across architectures. Details about PROBE architecture, training procedure and complete algorithm is provided in Methods (Section 4).
2 Results 2.1 PROBE achieves high-precision reliability classification for AIMNet2 We trained PROBE on top of AIMNet2 [6] that produces 256-dimensional per-atom AIM embeddings that encode the local chemical environment of each atom through three rounds of message passing. The model outputs per-atom partial charges, total energy both of which are passed to PROBE in addition to the AIM vectors. Results are reported on an additional held-out test set of 3.76M molecules absent from both AIMNet2 and PROBE training set. Additional details on AIMNet2 training and evaluation are provided in Methods (Section 4.3). We evaluate PROBE along three standard UQ axes. Discrimination measures how well the reliability score separates truly reliable from unreliable predictions; we report
5
classification accuracy, Matthews correlation coefficient (MCC), and F1. Ranking measures how well the score orders molecules by actual error magnitude, captured below by the monotonic relationship between P (unreliable) and mean error. A third axis, calibration (agreement between predicted probabilities and empirical frequencies in each probability bin), is a distinct property that we also characterize.
Fig. 2: PROBE performance on AIMNet2 (3.76M held-out molecules, 50thpercentile boundary at 1.64 kcal/mol). (a) Error distributions for the top-100,000 high-confidence reliable (blue) and high-confidence unreliable (orange) predictions, with the class boundary shown as a dashed line. (b) Mean AIMNet2 error as a function of probability cutoff for the reliable (blue) and unreliable (orange) classes. (c) Highconfidence coverage as a function of probability cutoff for both classes. Vertical and horizontal markers indicate cutoffs corresponding to 25%, 50%, and 75% coverage thresholds. (d) Overall classification accuracy (solid navy) and HC subset accuracy (dashed lines, one per coverage threshold) as a function of binned AIMNet2 error magnitude. Shading indicates accurate (blue, overall accuracy ≥ 70%) and buffer (yellow, overall accuracy < 70%) regions.
6
Error separation at high confidence. Fig.2(a) shows error distributions for the top-100,000 HC-reliable and HC-unreliable predictions (∼2.7% of the test set on each side). High-confidence (HC)-reliable predictions have mean error 0.494kcal/mol while HC-unreliable reach 11.319kcal/mol, well separated on either side of the 1.64kcal/mol boundary, confirming that PROBE’s confidence signal isolates genuinely distinct error regimes rather than partitioning noise near the boundary. The two HC subsets occupy disjoint regions of the error distribution on either side of the class boundary, indicating that PROBE’s confidence signal isolates genuinely distinct error regimes rather than merely partitioning noise near the boundary. Table 1 reports analogous means at fixed probability cutoffs covering larger fractions of the test set.
Table 1: PROBE high-confidence performance for AIMNet2 (50thpercentile boundary, 1.64 kcal/mol). Evaluated on 3.76M held-out molecules. Coverage = fraction of test set assigned a high-confidence label at the given cutoff. ērel and ēunrel are mean AIMNet2 errors (kcal/mol) for the HC-reliable and HC-unreliable subsets respectively. Matthews correlation coefficient (MCC), and F1 score are also reported. Cutoff
Coverage
Acc.
MCC
F1
ērel (kcal/mol)
ēunrel (kcal/mol)
All (P ≥ 0.5) P ≥ 0.6 P ≥ 0.7 P ≥ 0.8 P ≥ 0.9 P ≥ 0.95 P ≥ 0.99
100.0% 75.9% 54.0% 36.5% 23.9% 18.0% 10.8%
71.6% 77.2% 83.2% 88.7% 93.2% 95.5% 97.9%
0.402 0.509 0.628 0.746 0.849 0.902 0.957
0.633 0.685 0.749 0.826 0.898 0.936 0.974
1.19 1.06 0.94 0.82 0.71 0.65 0.58
4.21 5.20 6.66 8.29 10.22 11.50 13.14
Confidence monotonically tracks error. Fig.2(b) shows that mean error for reliable predictions decreases monotonically and for unreliable predictions increases monotonically as the probability cutoff increases for creating a subset (molecules with predicted class probability lower than this cutoff are ignored), across the full range P ∈ [0.5, 1.0]. At P ≥ 0.99, HC-reliable mean error reaches 0.58kcal/mol while HCunreliable reaches 13.14kcal/mol (Table1). Calibration analysis (SI Fig. S7a) confirms excellent ECE = 0.039 and Brier score = 0.182 for the test dataset, indicating that P (unreliable) reflects empirical frequencies rather than a mere rank-ordering signal. Coverage-accuracy tradeoff. Fig. 2(c) shows the selective-prediction operating characteristic: monotonic coverage and accuracy trade off with the probability cutoff. Table 1 reports combined accuracy, MCC, F1, and mean errors across a range of operating points. At P ≥ 0.7, PROBE retains 54.0% of the test set at 83.2% accuracy; at P ≥ 0.9, coverage drops to 23.9% while accuracy rises to 93.2%; at P ≥ 0.99, 10.8% of molecules are retained at 97.9% accuracy. All relevant metrics improve monotonically with probability cutoff, indicating that high-confidence predictions are more likely to
7
be correct. Even at strict cutoffs, PROBE still yields high-confidence predictions for a sizeable fraction of the test set. Accuracy as a function of error magnitude. Fig.2(d) shows classification accuracy binned by AIMNet2 error. Near the 1.64kcal/mol boundary (buffer region, shaded yellow) the problem is inherently ambiguous and accuracy drops below 70% overall; away from it, classification is confidently correct on both sides (shaded blue). HC subset accuracy curves lie above the overall curve throughout the buffer and the buffer narrows at stricter cutoffs, confirming that confidence filtering preferentially removes ambiguous near-boundary molecules. Overall, PROBE achieves 71.6% accuracy, F1 = 0.633, and MCC = 0.402 on the full 3.76M test set (Table 1). Confidence filtering converts this into a highly precise classifier: at P ≥ 0.9, accuracy = 93.2%, F1 = 0.898, and MCC = 0.849.
2.2 Comparison with ensemble-based UQ The standard UQ approach for MLIPs uses the standard deviation of an ensemble of independently trained models as a proxy for prediction uncertainty [18, 19, 24]. We benchmark this approach using a 4-model AIMNet2 ensemble trained on the same data and evaluated √ on the same 3.76M molecule test set, scaling ensemble standard deviation by 1/ Natoms to reduce the effect of molecule size. The Spearman correlation between scaled σ and the actual AIMNet2 error is ρ = 0.229, a weak monotonic signal at the individual-molecule level, consistent with prior observations [31] that ensemble members trained on the same data make correlated mistakes in chemically underrepresented regions. To benchmark ensemble σ as a reliability classifier (the quantity PROBE is optimized for), we threshold scaled σ at its median (SI Fig. S2) to define reliable and unreliable classes. This protocol evaluates σ as a classifier under the same decision framing as PROBE. As a baseline, we use majority-class prediction: the class boundary of 1.64 kcal/mol produces a 60/40 reliable/unreliable split on the test set, so always predicting the reliable class yields 60% accuracy.
Table 2: Comparison of UQ methods for AIMNet2 (boundary at 1.64 kcal/mol). Evaluated on 3.76M held-out molecules. The majority-class baseline always predicts reliable, reflecting the 60/40 class split on the test set. HC = high-confidence subset; not applicable (N/A) for methods that do not produce calibrated probabilities. Method
Accuracy (all)
HC Accuracy (P ≥ 0.7)
HC Accuracy (P ≥ 0.9)
Majority-class baseline Ensemble std (4 models) PROBE (ours)
60.0% 57.6% 71.6%
N/A N/A 83.2%
N/A N/A 93.2%
Table 2 compares the majority-class baseline, ensemble thresholding, and PROBE. When ensemble standard deviation is thresholded at its median to produce a binary 8
classifier, it achieves 57.6% accuracy, slightly below the 60% majority-class baseline. In contrast, PROBE reaches 71.6% accuracy across all predictions. At high confidence (P ≥ 0.9), PROBE achieves 93.2% accuracy, a regime the ensemble cannot access without additional calibration because it produces no per-prediction probability. The 4-model ensemble retains a weak continuous-ranker signal (ρ = 0.229) but offers no binary-classification advantage over the majority-class baseline at the median threshold. In practice, PROBE adds less than 1% computational overhead to a single MLIP forward pass. The classifier has 567K parameters, two orders of magnitude smaller than either backbone.
2.3 PROBE extends to MACE representations
Fig. 3: PROBE performance on MACE-OFF23 (50,195 held-out molecules, boundary at 0.36 kcal/mol). (a) Mean MACE error as a function of probability cutoff for the reliable (blue) and unreliable (orange) classes. (b) High-confidence coverage as a function of probability cutoff. Vertical and horizontal markers indicate cutoffs corresponding to 25%, 50%, and 75% coverage thresholds. (c) Overall classification accuracy (solid navy) and HC subset accuracy (dashed lines, one per coverage threshold) as a function of binned MACE error magnitude. Shading indicates accurate (blue, overall accuracy ≥ 70%) and buffer (yellow, overall accuracy < 70%) regions.
We apply PROBE to MACE-OFF23 with no changes to the classifier architecture, hyperparameters, or training protocol used for AIMNet2 in Section 2.1; the only modification is that per-atom embeddings are read from the invariant (L = 0) scalar components of the B-features (Eq. 10 of [5]) from the final tensor product block. This serves as our primary test of architecture transfer. Additional details about the MACE architecture and training data are provided in Methods (Section 4.4). PROBE achieves 80.5% overall accuracy and MCC = 0.612 on the 50,195-molecule test set (Table 3). PROBE’s absolute metrics are higher on MACE-OFF23 than
9
Table 3: PROBE high-confidence performance for MACE-OFF23 (boundary at 0.36 kcal/mol). Evaluated on 50,195 held-out molecules. ērel and ēunrel are mean MACE errors (kcal/mol) for the HC-reliable and HC-unreliable subsets respectively. Cutoff
Coverage
Acc.
MCC
F1
ērel (kcal/mol)
ēunrel (kcal/mol)
All (P ≥ 0.5) P ≥ 0.6 P ≥ 0.7 P ≥ 0.8 P ≥ 0.9 P ≥ 0.95 P ≥ 0.99
100.0% 88.5% 74.6% 57.3% 35.5% 23.5% 10.7%
80.5% 83.8% 87.5% 91.5% 96.0% 97.9% 99.6%
0.612 0.679 0.752 0.832 0.918 0.952 0.985
0.826 0.857 0.889 0.920 0.952 0.967 0.987
0.19 0.17 0.15 0.13 0.12 0.11 0.09
0.89 0.95 1.03 1.19 1.50 1.77 2.20
on AIMNet2. The two evaluations differ simultaneously in test-distribution overlap, training-set size, class-boundary value, richness of MLIP internal representation, and backbone architecture. We therefore do not attribute the numerical gap to any single factor here, and return to scaling-related implication in Section 3. The operationally relevant result is that identical PROBE hyperparameters transfer across both backbones without any tuning. As with AIMNet2, confidence monotonically tracks error in both directions (Fig. 3(a)): mean error for HC-reliable predictions decreases from 0.19 kcal/mol at P ≥ 0.5 to 0.09 kcal/mol at P ≥ 0.99, while HC-unreliable predictions increase from 0.89 to 2.20 kcal/mol over the same range. Coverage filtering shows the expected monotonic tradeoff (Fig. 3(b)). At P ≥ 0.7, PROBE retains 74.6% of the test set at 87.5% accuracy (MCC = 0.752); at P ≥ 0.9, 35.5% of molecules are retained at 96.0% accuracy (MCC = 0.918). Fig. 3(c) shows that the buffer region where overall accuracy falls below 70% is narrow and centered tightly on the 0.36 kcal/mol boundary, and HC subset lines sit well above the overall accuracy curve throughout. Furthermore, MACE-OFF23 shows near-perfect calibration (ECE = 0.011, Brier score = 0.136; SI Fig. S7b), consistent with the rich equivariant representations. These results confirm that PROBE is not contingent on a particular representation family. AIMNet2’s AIM embeddings and MACE’s scalar node features are constructed from different physical priors and trained on datasets differing by two orders of magnitude in size. Both support PROBE with identical classifier hyperparameters. The only backbone-specific change is zeroing the optional partial-charge injection term for MACE, which does not expose atomic charges.
2.4 Attention scores provide an interpretable atom-level view of prediction reliability PROBE’s multi-head self-attention mechanism yields chemically interpretable, peratom importance scores at no additional computational cost. For each molecule i and (m) attention head m, the attention matrix Ai ∈ RNi ×Ni is computed as described in
10
Fig. 4: Atom-level reliability interpretation via PROBE attention scores. (a) Molecular environment containing iodine and bromine colored by per-atom importance score sia (green–yellow–red; red = high importance). (b) Mean-over-heads self-attention matrix for the molecule in (a); query atoms along the y -axis, key atoms along the x-axis. (c) Per-atom importance score sia ; iodine (I) and bromine (Br) dominate. (d) Four further examples of diverse chemical environments flagged as unreliable by PROBE (all molecules, P (unreliable) ≥ 0.95). We emphasize that this is a qualitative illustration of the attention-score output; systematic validation that flagged atoms correspond to training-data underrepresentation would require targeted DFT benchmarks and dataset analysis beyond the scope of this work. (m)
Section 4, where entry Ai,ab is the weight atom a (query, y -axis in Fig. 4b) places on atom b (key, x-axis). To obtain a scalar importance score for each atom b, we sum over all queries, i.e., we ask how much total attention atom b receives from every other atom (including itself) across all heads: M
N
i 1 XX (m) A , sib = Zi m=1 a=1 i,ab
(1)
P where Zi normalizes so that b sib = 1. Atoms with high sib are strongly attended to by the rest of the molecule, meaning they exert disproportionate influence on the molecular embedding that drives the final reliability decision. Fig. 4 illustrates this on high-confidence unreliable predictions from the 3.76M test set for AIMNet2 (Section 2.1). In panel (a), a molecular environment containing iodine and bromine is shown: both heavy halogens accumulate disproportionately high importance scores (panels b-c). This is chemically sensible, since heavy halogens create unusual local electron density distributions that can lie outside the training distribution of AIMNet2, which PROBE correctly identifies. Panel (d) shows four additional
11
Fig. 5: UMAP projections of PROBE molecular embeddings. (a,b) 3.76M held-out AIMNet2 test set. (c,d) 2.4M biologically relevant molecules from Reidenbach et al. (a,c) P (unreliable) projected onto UMAP coordinates. (b) Selected element presence overlays (molecule counts in parentheses). (d) Element presence (Br, I) and RDKit descriptor projections.
examples, selected to span diverse chemistry, in which attention mass concentrates on chemically unusual moieties (heavy halogens, hypervalent sulfur, phosphorus-halogen motifs). These are illustrative rather than a systematic evaluation; the population-level signal is quantified in Section 2.5. In deployment, a practitioner can inspect the importance map of any flagged molecule to identify which structural feature triggered the unreliable prediction. Aggregating importance scores over a large set of unreliable predictions surfaces recurrent structural motifs that the classifier’s representation treats as discriminative of unreliability. This creates a closed feedback loop, PROBE not only flags unreliable predictions but actively guides where to acquire new training data, complementing the active learning workflow described in Section 2.6.
2.5 PROBE molecular embeddings reveal chemical space coverage and reliability boundaries The molecular embedding produced by PROBE is trained solely to separate reliable from unreliable MLIP predictions. If this embedding has learned chemically meaningful structure, molecules should organize into spatially distinct reliable and unreliable regions, for a completely unseen dataset, and requiring no reference energies at inference time. We first projected the molecular embeddings of the 3.76M held-out test
12
set created by PROBE (input to classifier, Section 2.1) onto two dimensions using UMAP [58] and colored each point by P (unreliable) (Fig. 5a). The embedding visibly organizes test molecules into reliable- and unreliable-enriched regions across local clusters, consistent with the claim that the embedding captures chemically distinct success and failure modes. Overlaying element presence onto the same UMAP coordinates (Fig. 5b) is consistent with specific chemistries concentrating in specific regions of the embedding, though we do not claim causal attribution of reliability to element identity alone. Carbon is ubiquitous across the entire manifold, as expected for organic chemistry datasets. Sulfur follows a similar but slightly narrower distribution. By contrast, boron, selenium, and iodine map predominantly onto high-P (unreliable) regions, associating their presence with AIMNet2 failure, while phosphorus occupies a mixed region. The complete element projection for all 14 AIMNet2 elements is provided in SI Fig. S4. We additionally applied PROBE without retraining to a completely separate 2.4Mmolecule set of biologically relevant molecules published by Reidenbach et al. [59]. PROBE labels 82% of these molecules as reliable (P (reliable) ≥ 0.5) and the remaining 18% as unreliable; because no reference energies are available here, this is a classifier self-report rather than a ground-truth measurement.It is consistent in direction with the set’s drug-like, CHNO-dominated composition. The UMAP projection (Fig. 5c) shows a compact, predominantly blue manifold with a distinct tail of highP (unreliable) molecules, again spatially separated from the reliable core. Element presence and RDKit [60] descriptors computed from connectivity (Fig. 5d) show that iodine-containing molecules for example, concentrate in the unreliable tail, while bromine is distributed more broadly. Among molecular properties, higher molecular weight, TPSA, number of H-bond acceptors, and rotatable bond count all correlate positively with P (unreliable). Full element and descriptor projections for this dataset are provided in SI Figs. S5, S6. These results demonstrate that the PROBE molecular embedding encodes actionable chemical information about MLIP reliability. A practitioner can use this analysis on any new compound library to map where molecules fall relative to known failure modes without ground truth labels, and prioritize structural classes for further validation or targeted data acquisition.
2.6 Retrospective active learning As a proof of concept, we show how PROBE can guide targeted data acquisition for MLIP retraining. We compare two selection strategies applied to AIMNet2 retraining: Ensemble (molecules ranked by ensemble standard deviation) and PROBE (molecules ranked by P (unreliable)). Both strategies operate on the same fixed pool of ∼20M molecules and are evaluated on a held-out test set of 13.6M molecules. Cycle 0 trains an ensemble of four independent AIMNet2 models on 1M molecules drawn by stratified random sampling; for PROBE-based selection we use a single AIMNet2 model from this four-member Cycle-0 ensemble. Each subsequent cycle augments the training set with 1M additional molecules selected by the respective strategy, with model parameters loaded from the previous cycle. The ensemble Cycle 0 model
13
starts with lower error than the single PROBE model (RMSE = 3.030 vs. 3.234 kcal/mol), a consequence of ensemble averaging (i.e., better models than the one chosen for PROBE) rather than data selection.
Table 4: Retrospective active learning: RMSE (kcal/mol) across cycles. Each cycle adds 1M molecules selected by the respective strategy. Evaluated on a fixed 13.6M molecule held-out test set. ∆RMSE = cumulative change from Cycle 0. The ensemble Cycle 0 baseline is lower due to ensemble averaging, not data selection. Cycle
Method
RMSE (kcal/mol)
∆ (kcal/mol)
C0 C0
PROBE Ensemble
3.234 3.030
— —
C1 C1
PROBE Ensemble
2.903 2.907
−0.331 −0.123
C2 C2
PROBE Ensemble
2.709 2.817
−0.525 −0.213
Across two retrospective active-learning cycles, PROBE-guided acquisition reduced RMSE by 0.525 kcal/mol (16.2%) versus 0.213 kcal/mol (7.0%) for ensemble-σ selection, more than twice the improvement. By Cycle 2, PROBE not only closes the initial 0.2 kcal/mol Cycle-0 gap but surpasses the ensemble at roughly one-quarter of the per-cycle retraining compute (one model retrained per cycle versus four). We report this as suggestive rather than definitive: the comparison involves a single seed per strategy, two acquisition cycles, no random-selection control, and simultaneously differs in selection strategy and in number of models retrained. A like-for-like benchmark with matched retraining protocols, across multiple seeds and with a random-selection baseline, is needed to attribute the gap to selection quality alone. We leave this to future work.
3 Discussion The question of when to trust an MLIP admits two answers: regress the continuous error magnitude, or predict a binary reliability verdict. We argue the second is more tractable: reducing to classification on the backbone’s own representations and sufficient for most deployment decisions. Grounding the reliability signal in MLIP internal representations rather than committee disagreement is what drives PROBE’s advantage over ensemble-based approaches. The active-learning advantage PROBE demonstrates is the expected consequence of a measured ranking-quality gap. PROBE is one of the few single-model UQ methods demonstrated in this work to outperform ensemble-based selection in an MLIP activelearning scenario; prior comparative studies [20] found that single-model UQ does not consistently match ensembles for MLIP error estimation, and our retrospective experiment suggests the binary-classification framing combined with representation-based 14
scoring can close or reverse this gap in at least one setting, at roughly one-quarter the retraining cost per cycle. The same explanation underlies the static benchmark in Section 2.2: ensemble σ delivers a weak correlation against the AIMNet2 error, whereas PROBE’s P (unreliable) tracks mean error monotonically across the full [0.5, 1.0] probability range on both backbones (Figs. 2(b), 3(a)). A weak ranker produces noisy acquisition decisions; a monotonic one does not. For foundation-scale MLIPs (UMA [9], AllScAIP [15]), maintaining N independent model copies is impractical. PROBE’s post-hoc frozen-backbone design adds negligible inference overhead and requires no backbone retraining, placing it alongside conformal prediction [42], shallow ensembles [47], and prediction-rigidity methods [46] among the few UQ routes practical at foundation scale. Among these, PROBE is the only one that delivers a binary reliability decision and per-atom interpretability (Table 5).
Table 5: Single-backbone UQ methods for MLIPs compared. “Post-hoc” denotes methods requiring no backbone modification or retraining. ✓ = yes; × = no; “partial” = additional trained heads but no full backbone retraining. PROBE is the only method that is simultaneously single-backbone, post-hoc, binary-output, and per-atom interpretable. Method
Ref.
Single backbone
Post-hoc
Output type
Per-atom signal
Shallow ensembles Prediction rigidity Latent-distance UQ Multi-head committee Conformal prediction PROBE (this work)
[47] [46] [45] [43] [42] —
✓ ✓ ✓ ✓ ✓ ✓
partial ✓ ✓ partial ✓ ✓
continuous continuous continuous continuous coverage interval binary probability
× × ✓ × × ✓ (attention)
Beyond tractability, the backbone comparison hints at a favorable scaling trend. PROBE achieves notably higher discrimination on MACE-OFF23 than on AIMNet2 despite MACE being trained on 20× less data. Since the classifier is identical in both cases, the gain must come from the richer representation MACE exposes. This implies a favorable scaling trajectory: as backbone representations improve, PROBE’s reliability signal strengthens, suggesting that next-generation foundation MLIPs (UMA [9], AllScAIP [15]) should yield better, not worse, post-hoc reliability estimates at fixed classifier cost. A systematic scaling study across backbone sizes is a natural next step. The 50th-percentile boundary yields stable training and a balanced class split. Moving the boundary substantially degrades performance as the classifier defaults to the majority class. Users targeting a specific error tolerance are better served by adjusting the probability cutoff on a balanced model rather than redefining the boundary. Importantly, the boundary’s absolute value (1.64 kcal/mol for AIMNet2; 0.36 kcal/mol for MACE-OFF23) is not comparable across backbones; users needing a fixed error tolerance should treat P (reliable) as a relative ranking signal calibrated on a small labeled reference set.
15
The attention mechanism yields per-atom importance scores at no additional computational cost, identifying exactly which structural motifs drive an unreliable prediction and providing actionable guidance on where additional training data or DFT validation is most needed. Two independent signals converge on the same AIMNet2 failure mode: locally, attention maps assign the highest importance to iodine and bromine in high-confidence unreliable predictions (Fig.4); globally, the UMAP projection of PROBE’s molecular embedding concentrates iodine-containing molecules in the unreliable tail on both the 3.76M test set and the Reidenbach 2.4M library (Fig.5b,d). Their agreement across different granularities and datasets supports the interpretation that PROBE has surfaced a genuine training-data gap rather than an artefact of either analysis. Several limitations apply: (i) PROBE is a binary classifier, applications requiring continuous uncertainty estimates would need a multi-class or regression extension. (ii) Like most UQ approaches, performance degrades for chemical environments completely absent from training data. (iii) Training requires reference labels for backbone training molecules, which may be unavailable for large foundation models. (iv) Only energy-prediction reliability is evaluated here; extending PROBE to per-atom quantities such as forces is conceptually direct but not demonstrated in this work. (v) The reliability boundary depends on the backbone’s error distribution and must be recalibrated when the backbone changes. (vi) PROBE’s reliability labels are currently derived from DFT-based reference energies and thus cannot capture errors arising from DFT approximations themselves, but this limitation is not fundamental. Following Kellner et al. [61], PROBE could be calibrated using experimentally-derived reliability labels wherever such data exist, extending its scope beyond DFT-level errors to empirical measurements. Future directions. Three extensions are natural: periodic systems (applying PROBE to materials foundation MLIPs such as MACE-MP [14] and UMA [9]); force uncertainty (classifying per-atom force-error percentiles, most relevant to molecular dynamics); and genuine out-of-distribution evaluation (training on CHNO, testing on heavy halogens or transition metals). Each is enabled by the post-hoc architecture-agnostic design. Conclusion. The backbone’s own internal representations, fed to a lightweight discriminative classifier, yield a reliability verdict sharper than ensemble standard deviation and cheaper to deploy at foundation scale. This verdict is actionable for the binary decisions that dominate practice: whether to include a molecule in the next active-learning cycle, accept a predicted geometry, or trigger a DFT recalculation. Within its scope, PROBE turns “is this prediction trustworthy?” into “with what probability?”, without modifying or retraining the underlying model. As MLIPs scale toward foundation models driving autonomous simulation workflows, knowing when to trust them becomes a prerequisite. PROBE addresses this requirement for the binary-decision regime.
16
4 Methods 4.1 PROBE Architecture and Training PROBE acts as a lightweight binary classifier that takes atom-level representations and other atom-level predictions e.g., partial charges alongside the predicted energy from a single, frozen MLIP as input. We derive training labels from the absolute per(m) molecule energy error, ϵm = |Ê (m) − Eref |. The class √ boundary is itself a percentile of this same training distribution, and the loss is N -normalized. Size-induced error inflation is therefore partially absorbed rather than confounded with unreliability. Per-atom or size-normalized error labels are a straightforward drop-in alternative. Rather than predicting the exact error magnitude, PROBE predicts the probability that the error lies below or above a chosen threshold, set as a percentile of the empirical error distribution of the training data. All the experiments in this work use the 50thpercentile boundary, yielding balanced classes. Atom Encoder. Given a molecular configuration, the backbone MLIP Mθ performs a standard forward pass, producing per-atom internal embeddings hi ∈ Rd alongside the predicted energy Ê and atom-level quantities such as partial charges qi . The atom encoder MLPenc : Rd → R256 is a three-layer network with hidden dimensions (d, 256, 128, 256) with LayerNorm, GELU activations, and dropout (p = 0.1). For AIMNet2, hi is the per-atom AIM vector; for MACE, it is the invariant (L=0) scalar components of the B-features from the final product block. Invariant features are the natural input to a classifier whose target (absolute energy error) is itself a scalar invariant, and using them preserves compatibility without imposing equivariance constraints on the PROBE head. These embeddings are passed through the atom encoder: zi = MLPenc (hi ), producing atom-level features relevant for the classification task. Partial charges are incorporated via a learned linear projection: zi ← zi + Wq qi . Molecule encoder. Atom features are processed through a multi-head selfattention layer with H heads:
z′i = zi + LayerNorm
H X N X
(h)
αij V(h) zj
(2)
h=1 j=1
√ (h) (h)⊤ (h) where αij = softmaxj z⊤ K zj / dk are the attention weights, masked to i Q exclude padding atoms. Each atom aggregates context from all other atoms in the molecule, propagating global chemical environment information into the per-atom representation. Multi-head self-attention operates with H = 32 heads of dimension dk = 8, giving a total attention dimension of 256. The molecular descriptor g ∈ R514 is formed by concatenating masked mean-pool and max-pool of the attended atom features (each R256 ), the scalar predicted energy Ê , and atom count N : "
1 X ′ g= z N i i
# max z′i i
17
Ê
N
(3)
This descriptor is projected to a 256-dimensional molecular embedding e = Wproj g ∈ R256 via a learned linear map Wproj ∈ R256×514 . The embedding e serves as a general-purpose molecular representation and is used directly for chemical space coverage analysis. Classifier. The classifier MLPcls : ∆1 then maps e to binary reliability probabilities. The class boundary is set at the pth percentile of the training error distribution (default p = √50, yielding balanced classes). The loss function is cross-entropy normalized by Nm to prevent large molecules from dominating:
L=
1 X wym · CE(fϕ (em ), ym ) √ , |B| Nm
wc =
m∈B
|D| 2 |Dc |
(4)
where ym ∈ {0, 1} is the binary reliability label and |Dc | is the number of training samples in class c. When classes are balanced, wc = 1 and the loss reduces to plain normalized cross-entropy. The classifier MLPcls : R256 → ∆1 maps e to reliability probabilities through hidden dimensions [128, 32] with LayerNorm, GELU activations, and dropout (p = 0.1). The full model has 5̃67K trainable parameters. Optimization. Parameters ϕ are optimized with AdamW (learning rate η = 5 × 10−5 , weight decay λ = 10−4 , gradient clipping with max norm 1.0). A ReduceLROnPlateau scheduler reduces η by factor 0.9 after 5 epochs without validation loss improvement, with a minimum learning rate of 5 × 10−6 . Early stopping is applied with patience 25 epochs. All models use a 90/10 train/validation split; a fully heldout test set is used exclusively for final evaluation. The MLIP backbone is fully frozen throughout.
4.2 Problem Formulation Let Mθ denote a trained MLIP with frozen parameters θ, mapping a molecular configuration (R, Z) ∈ RN ×3 × ZN to a predicted energy Ê ∈ R and perd atom latent representations {hi }N i=1 ⊂ R . Given a labeled calibration dataset (m) |D| (m) D = {(R(m) , Z(m) , Eref )}m=1 with reference energies Eref , define the per-molecule absolute error (m) ϵm = Ê (m) − Eref . (5) |D| Let bp = Quantilep {ϵm }m=1 denote the pth percentile of the empirical error distribution. We assign binary reliability labels:
( ym =
0 1
if ϵm < bp if ϵm ≥ bp
(reliable) (unreliable).
(6)
1 Goal. Learn a classifier fϕ : {hi }N i=1 , Ê, N → ∆ with parameters ϕ, operating solely on quantities produced by Mθ (no access to Eref at inference time), such that the predicted probability fϕ (·)1 ≈ P (y = 1 | R, Z) generalizes to held-out molecules. The parameters θ of Mθ are never updated.
18
Algorithm 1 PROBE Training Require: Frozen MLIP Mθ ; calibration set D; percentile p ∈ (0, 1); max epochs T Ensure: PROBE classifier parameters ϕ 1: Label generation: (m) 2: Compute ϵm ← |Mθ (R(m) , Z(m) ) − Eref | for all m ∈ D 3: Set boundary bp ← Quantilep ({ϵm }); assign ym ← 1[ϵm ≥ bp ] 4: Compute class weights wc ← |D| / (2 |Dc |), where Dc = {m : ym = c} 5: for epoch t = 1, . . . , T do 6: for each mini-batch B ⊆ D do 7: MLIP representation extraction (no gradient): (m) (m) 8: {hi }, Ê (m) , {qi } ← Mθ (R(m) , Z(m) ) ∀m ∈ B 9: Atom encoding: (m) (m) (m) 10: zi ← MLPenc (hi ) + Wq qi 11: Multi-head self-attention with masking: (m) 12: z′ ← z(m) + LayerNorm MHSA z(m) , mask(m) 13: Molecularh pooling: i P ′(m) ′(m) 1 14: g(m) ← N (m) Ê (m) N (m) maxi zi i zi 15: Molecular embedding: 16: e(m) ← Wproj g(m) 17: Classification and loss: 18: ŷ(m) ← softmax MLPclf (e(m) ) 1 X wym · CE(ŷ(m) , ym ) √ 19: L← |B| N (m) m∈B
Update ϕ ← ϕ − η ∇ϕ L (AdamW) 21: end for 22: Reduce η on validation loss plateau; apply early stopping 23: end for 20:
Training minimizes a size-normalized, class-weighted cross-entropy:
L=
1 X wym · CE(fϕ (cm ), ym ) √ , |B| Nm m∈B
wc =
|D| 2 |Dc |
(7)
where B is a mini-batch, Nm is√the atom count of molecule m, and wc are inversefrequency class weights. The Nm normalization prevents large molecules from dominating the gradient. When classes are balanced, wc = 1 and the loss reduces to plain normalized cross-entropy. This formulation reduces MLIP reliability estimation to supervised binary classification on a learned representation space, avoiding direct regression of the heavy-tailed distribution p(ϵ | R, Z). The full training procedure is summarized in Algorithm 1.
19
4.3 AIMNet2 training and evaluation AIMNet2 is a neural network potential based on the atoms-in-molecules (AIM) framework trained on 20 million hybrid DFT calculations (B97-3c/def2-mTZVP) covering neutral and charged organic molecules across 14 element types. AIMNet2 is initialized from Behler-Parrinello type atomic environment vectors [1]. In AIMNet2, the per-atom learned representation, the atoms-in-molecules (AIM) vector, is explicitly designed as a general-purpose chemical descriptor functionally analogous to the electron density ρ(r) in DFT. It serves as a common substrate refined through iterative self-consistency, specifically message passing with neural charge equilibration, from which multiple physical observables such as energy, forces, partial charges are projected out through distinct learned readout heads. This multi-task design implies that the representation already encodes rich chemical context. We trained an ensemble of four AIMNet2 models using the default architecture and training procedure outlined in [6] that were fully frozen during PROBE training. PROBE was trained using the same dataset used to train AIMNet2 models with a 90/10 train/validation split. The 50th-percentile class boundary on the training set is 1.64 kcal/mol, yielding balanced classes (50%/50%); the full training-set error distribution is shown in the Supplementary Information (SI) Fig. S1. Applying this fixed 1.64 kcal/mol class boundary (the 50th percentile of the training error distribution) to the 3.76M test set yields a 60/40 reliable/unreliable split.
4.4 MACE training and evaluation MACE-OFF23 [5, 14] is a higher-order E (3)-equivariant message-passing neural network that achieves body-order completeness with only two interaction layers. Each interaction layer constructs equivariant features up to angular momentum L. MACEOFF23 is a transferable organic force field trained on the SPICE dataset of drug-like organic molecules. We use the publicly available large model and its training set (855,905 molecules, split 90/10 for PROBE training and validation). Evaluation uses the original held-out test set of 50,195 molecules. The class boundary is set at 0.36 kcal/mol (50th percentile, balanced classes). The resulting class split is 47.3% reliable / 52.7% unreliable on the test set, with mean errors of 0.139 and 1.029 kcal/mol, respectively. Training-set error distribution and train/validation loss curves for training PROBE are provided in the SI Fig. S3. PROBE accesses the scalar node features (node feats) from MACE’s final interaction layer as per-atom embeddings, with the MACE backbone cast to float32 and fully frozen. MACE-OFF23 does not expose atomic partial charges in the same manner as AIMNet2; we therefore set qi = 0 in the partial-charge injection term, reducing the atom encoder to zi = MLPenc (hi ) for this backbone. MACE-level performance remaining competitive with AIMNet2 indicates that partial charges are not essential to PROBE’s signal when the backbone’s per-atom representation is sufficiently expressive.
20
Acknowledgements S.M. would like to thank Filipp Gusev, Bhupalee Kalita, and Roman Zuabatyuk for insightful discussions. This work was made possible by the Office of Naval Research (ONR) through support provided by the Energetic Materials Program (MURI grant no. N00014-21-1-2476). The authors acknowledge computing resources provided through the Advanced Cyberinfrastructure Coordination Ecosystem: Services and Support (ACCESS) program, specifically the NCSA Delta and Delta AI systems. This work was performed, in part, at the Center for Integrated Nanotechnologies, an Office of Science User Facility operated for the U.S. Department of Energy (DOE) Office of Science by Los Alamos National Laboratory (Contract 89233218CNA000001) and Sandia National Laboratories (Contract DE-NA-0003525). Claude Code (Opus 4.5 / Sonnet 4.5) was employed in the implementation of this work.
Data Availability AIMNet2 training data is available through the AIMNet2 repository. MACE training data (SPICE dataset) is publicly available. PROBE model weights and inference code will be made available at publication.
Code Availability PROBE training and inference code, along with analysis notebooks, will be released at https://github.com/isayevlab/PROBE upon publication.
Author Contributions S.M. and O.I. designed research; S.M. developed the framework for PROBE; S.M. and I.C. performed research; S.M., I.C., O.I. wrote the paper.
Competing Interests The authors declare no competing interests.
References [1] Behler, J. & Parrinello, M. Generalized neural-network representation of highdimensional potential-energy surfaces. Physical review letters 98, 146401 (2007). [2] Kulik, H. J. et al. Roadmap on machine learning in electronic structure. Electronic Structure 4, 023004 (2022). [3] Schütt, K. T., Sauceda, H. E., Kindermans, P.-J., Tkatchenko, A. & Müller, K.R. Schnet–a deep learning architecture for molecules and materials. The Journal of chemical physics 148 (2018).
21
[4] Batzner, S. et al. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications 13, 2453 (2022). [5] Batatia, I., Kovacs, D. P., Simm, G., Ortner, C. & Csányi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in neural information processing systems 35, 11423–11436 (2022). [6] Anstine, D. M., Zubatyuk, R. & Isayev, O. Aimnet2: a neural network potential to meet your neutral, charged, organic, and elemental-organic needs. Chemical Science 16, 10228–10244 (2025). [7] Deng, B. et al. Chgnet as a pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence 5, 1031–1041 (2023). [8] Merchant, A. et al. Scaling deep learning for materials discovery. Nature 624, 80–85 (2023). [9] Wood, B. M. et al. Uma: A family of universal models for atoms. arXiv preprint arXiv:2506.23971 (2025). [10] Hohenberg, P. & Kohn, W. Inhomogeneous electron gas. Physical review 136, B864 (1964). [11] Kohn, W. & Sham, L. J. Self-consistent equations including exchange and correlation effects. Physical review 140, A1133 (1965). [12] Mehdi, S., Smith, Z., Herron, L., Zou, Z. & Tiwary, P. Enhanced sampling with machine learning. Annual Review of Physical Chemistry 75, 347–370 (2024). [13] Liao, Y.-L., Wood, B., Das, A. & Smidt, T. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. arXiv preprint arXiv:2306.12059 (2023). [14] Kovács, D. P. et al. Mace-off: Short-range transferable machine learning force fields for organic molecules. Journal of the American Chemical Society 147, 17598–17611 (2025). [15] Qu, E., Wood, B. M., Krishnapriyan, A. S. & Ulissi, Z. W. A recipe for scalable attention-based mlips: unlocking long-range accuracy with all-to-all node attention. arXiv preprint arXiv:2603.06567 (2026). [16] Gómez-Bombarelli, R. et al. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science 4, 268–276 (2018). [17] Hüllermeier, E. & Waegeman, W. Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods. Machine learning 110, 457– 506 (2021). 22
[18] Schran, C., Brezina, K. & Marsalek, O. Committee neural network potentials control generalization errors and enable active learning. The Journal of Chemical Physics 153 (2020). [19] Smith, J. S., Nebgen, B., Lubbers, N., Isayev, O. & Roitberg, A. E. Less is more: Sampling chemical space with active learning. The Journal of chemical physics 148 (2018). [20] Tan, A. R., Urata, S., Goldman, S., Dietschreit, J. C. & Gómez-Bombarelli, R. Single-model uncertainty quantification in neural network potentials does not consistently outperform model ensembles. npj Computational Materials 9, 225 (2023). [21] Perez, D., Subramanyam, A. P., Maliyov, I. & Swinburne, T. D. Uncertainty quantification for misspecified machine learned interatomic potentials. npj Computational Materials 11, 263 (2025). [22] Grasselli, F., Chong, S., Kapil, V., Bonfanti, S. & Rossi, K. Uncertainty in the era of machine learning for atomistic modeling. Digital Discovery 4, 2654–2675 (2025). [23] Heid, E., McGill, C. J., Vermeire, F. H. & Green, W. H. Characterizing uncertainty in machine learning for chemistry. Journal of Chemical Information and Modeling 63, 4012–4029 (2023). [24] Lakshminarayanan, B., Pritzel, A. & Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems 30 (2017). [25] Vandermause, J. et al. On-the-fly active learning of interpretable bayesian force fields for atomistic rare events. npj Computational Materials 6, 20 (2020). [26] Podryabinkin, E. V. & Shapeev, A. V. Active learning of linearly parametrized interatomic potentials. Computational Materials Science 140, 171–180 (2017). [27] Kulichenko, M. et al. Uncertainty-driven dynamics for active learning of interatomic potentials. Nature computational science 3, 230–239 (2023). [28] Zaverkin, V. et al. Uncertainty-biased molecular dynamics for learning uniformly accurate interatomic potentials. npj Computational Materials 10, 83 (2024). [29] van der Oord, C., Sachs, M., Kovács, D. P., Ortner, C. & Csányi, G. Hyperactive learning for data-driven interatomic potentials. npj Computational Materials 9, 168 (2023). [30] Kurniawan, Y., Wen, M., Tadmor, E. B. & Transtrum, M. K. Comparative study of ensemble-based uncertainty quantification methods for neural network
23
interatomic potentials. arXiv preprint arXiv:2508.06456 (2025). [31] Lu, S., Ghiringhelli, L. M., Carbogno, C., Wang, J. & Scheffler, M. On the uncertainty estimates of equivariant-neural-network-ensembles interatomic potentials. arXiv preprint arXiv:2309.00195 (2023). [32] Bilbrey, J. A., Firoz, J. S., Lee, M.-S. & Choudhury, S. Uncertainty quantification for neural network potential foundation models. npj Computational Materials 11, 109 (2025). [33] Gal, Y. & Ghahramani, Z. Dropout as a bayesian approximation: Representing model uncertainty in deep learning 1050–1059 (2016). [34] Wen, M. & Tadmor, E. B. Uncertainty quantification in molecular simulations with dropout neural network potentials. npj computational materials 6, 124 (2020). [35] Williams, C. K. & Rasmussen, C. E. Gaussian processes for machine learning Vol. 2 (MIT press Cambridge, MA, 2006). [36] Bartók, A. P., Kermode, J. et al. Improved uncertainty quantification for gaussian process regression based interatomic potentials. arXiv preprint arXiv:2206.08744 (2022). [37] Vandermause, J., Xie, Y., Lim, J. S., Owen, C. J. & Kozinsky, B. Active learning of reactive bayesian force fields applied to heterogeneous catalysis dynamics of h/pt. Nature Communications 13, 5183 (2022). [38] Farris, R., Telari, E., Artrith, N., Neyman, K. & Bruix, A. Bayesian neural networks versus deep ensembles for uncertainty quantification in machine learning interatomic potentials. arXiv preprint arXiv:2509.19180 (2025). [39] Coscia, D., de Haan, P. & Welling, M. Blips: Bayesian learned interatomic potentials. arXiv preprint arXiv:2508.14022 (2025). [40] Xu, H. et al. Evidential deep learning for interatomic potentials. Communications (2025).
Nature
[41] Vita, J. A., Samanta, A., Zhou, F. & Lordi, V. Ltau-ff: loss trajectory analysis for uncertainty in atomistic force fields. Machine Learning: Science and Technology 6, 015048 (2025). [42] Ho, C. H., Ortner, C. & Wang, Y. Flexible uncertainty calibration for machinelearned interatomic potentials. arXiv preprint arXiv:2510.00721 (2025). [43] Beck, H., Simko, P., Schaaf, L. L., Marsalek, O. & Schran, C. Multi-head committees enable direct uncertainty prediction for atomistic foundation models. The Journal of Chemical Physics 163 (2025). 24
[44] Janet, J. P., Duan, C., Yang, T., Nandy, A. & Kulik, H. J. A quantitative uncertainty metric controls error in neural network-driven chemical discovery. Chemical science 10, 7913–7922 (2019). [45] Musielewicz, J., Lan, J., Uyttendaele, M. & Kitchin, J. R. Improved uncertainty estimation of graph neural network potentials using engineered latent space distances. The Journal of Physical Chemistry C 128, 20799–20810 (2024). [46] Bigi, F., Chong, S., Ceriotti, M. & Grasselli, F. A prediction rigidity formalism for low-cost uncertainties in trained neural networks. Machine Learning: Science and Technology 5, 045018 (2024). [47] Kellner, M. & Ceriotti, M. Uncertainty quantification by direct propagation of shallow ensembles. Machine Learning: Science and Technology 5, 035006 (2024). [48] El-Yaniv, R. et al. On the foundations of noise-free selective classification. Journal of Machine Learning Research 11 (2010). [49] Geifman, Y. & El-Yaniv, R. Selective classification for deep neural networks. Advances in neural information processing systems 30 (2017). [50] Geifman, Y. & El-Yaniv, R. Selectivenet: A deep neural network with an integrated reject option 2151–2159 (2019). [51] DeVries, T. & Taylor, G. W. Learning confidence for out-of-distribution detection in neural networks. arXiv preprint arXiv:1802.04865 (2018). [52] Hendrycks, D. & Gimpel, K. A baseline for detecting misclassified and outof-distribution examples in neural networks. arXiv preprint arXiv:1610.02136 (2016). [53] Lee, K., Lee, K., Lee, H. & Shin, J. A simple unified framework for detecting outof-distribution samples and adversarial attacks. Advances in neural information processing systems 31 (2018). [54] Sun, Y., Ming, Y., Zhu, X. & Li, Y. Out-of-distribution detection with deep nearest neighbors 20827–20840 (2022). [55] Yang, J., Zhou, K., Li, Y. & Liu, Z. Generalized out-of-distribution detection: A survey. International Journal of Computer Vision 132, 5635–5662 (2024). [56] Fannjiang, C., Bates, S., Angelopoulos, A. N., Listgarten, J. & Jordan, M. I. Conformal prediction under feedback covariate shift for biomolecular design. Proceedings of the National Academy of Sciences 119, e2204569119 (2022). [57] Alain, G. & Bengio, Y. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644 (2016).
25
[58] McInnes, L., Healy, J. & Melville, J. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03426 (2018). [59] Reidenbach, D., Nikitin, F., Isayev, O. & Paliwal, S. G. Applications of modular co-design for de novo 3d molecule generation. Digital Discovery 5, 754–768 (2026). [60] Landrum, G. et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum 8, 5281 (2013). [61] Kellner, M., Hansen, T., Bligaard, T., Jacobsen, K. W. & Ceriotti, M. Errors that matter: Uncertainty-aware universal machine-learning potentials calibrated on experiments (2026). arXiv:2604.24607.
26
Supplementary Information Knowing when to trust machine-learned interatomic potentials Shams Mehdi, Ilkwon Cho, Olexandr Isayev∗ Department of Chemistry, Mellon College of Science, Carnegie Mellon University, Pittsburgh, Pennsylvania 15213, USA ∗
1
Supplementary Figures
Fig. S1: AIMNet2 error distribution and training dynamics. (a) Distribution of per-molecule AIMNet2 absolute errors (kcal/mol) with cumulative distribution function (right axis, black curve). Dashed vertical lines indicate the class boundary thresholds used for PROBE training; percentile markers are shown for reference. (b) Training (orange) and validation (blue) loss curves across epochs.
2
Fig. S2: Distribution of scaled ensemble standard deviation for AIMNet2. Histogram of √Nσ , where σ (kcal/mol) is the standard deviation across a 4-model atoms AIMNet2 ensemble and Natoms is the number of atoms in a molecule, evaluated on the 3.76M held-out test set. The dashed vertical line marks the median. The right axis shows the cumulative distribution function (CDF) with percentile annotations.
3
Fig. S3: MACE-OFF23 error distribution and PROBE training curves. (a) Distribution of MACE-OFF23 absolute prediction errors on the 855,905-molecule training set. The dashed vertical line marks the 50/50 class boundary at 0.36 kcal/mol. The right axis shows the cumulative distribution function (CDF) with percentile annotations. (b) PROBE train and validation loss as a function of epoch. The validation loss closely tracks the training loss throughout, indicating no overfitting. Here, we implemented early stopping at epoch 68.
4
Fig. S4: Element presence projected onto PROBE UMAP embedding for the 3.76M AIMNet2 test set. Each panel shows molecules containing at least one atom of the indicated element (colored) overlaid on the full dataset (gray). Molecule counts are given in parentheses. Rare and heavy elements (B, Si, As, Se, Br, I) concentrate in spatially distinct regions of the embedding that correspond to highP (unreliable) clusters, while abundant elements (H, C, N, O) span the full manifold. 5
Fig. S5: Element presence projected onto PROBE UMAP embedding for the 2.4M Reidenbach et al. dataset. Same layout as Fig. SS4. The dataset is dominated by CHNO chemistry. Iodine (11,617) and selenium (751) map onto the unreliable tail while boron clusters at a distant island, consistent with the AIMNet2 test set analysis. Bromine (102,656) is more broadly distributed, occupying both reliable and unreliable regions.
6
Fig. S6: RDKit molecular descriptor projections onto PROBE UMAP embedding for the 2.4M Reidenbach et al. dataset. Descriptors were computed from connectivity information using RDKit. Molecular weight, TPSA, number of Hbond acceptors, and number of rotatable bonds all increase toward the unreliable tail of the embedding, indicating that large, polar, and flexible molecules are more likely to receive unreliable AIMNet2 predictions. Fraction sp3 shows the different trend, highly saturated molecules are concentrated in the reliable core. Number of aromatic rings and H-bond donors show weaker spatial correlation with the reliability score.
7
Fig. S7: PROBE probability calibration. Reliability diagrams for PROBE on (a) AIMNet2 (3.76M held-out molecules) and (b) MACE-OFF23 (50,195 held-out molecules). Each point shows the fraction of molecules truly unreliable within a bin of predicted P (unreliable); the dashed diagonal indicates perfect calibration. Shading shows the gap between the calibration curve and the diagonal. ECE = expected calibration error; BS = Brier score.
8