cs.LGOct 26, 2026
SaveMulti-Fidelity Training Reshapes Uncertainty Decomposition in Atomistic Models
Organizations: Department of Computer Science, University of Pisa · Atomistic Simulations, Italian Institute of Technology
Abstract
Multi-fidelity training, which combines abundant low-fidelity data with scarce high-fidelity labels, is increasingly common in atomistic modeling, but its effect on predictive uncertainty remains largely unexplored. We investigate this question using deep ensembles of mean-variance MACE models trained on ANI-1x, contrasting high-fidelity training from scratch with low-fidelity pretraining followed by fine-tuning. We characterize aleatoric, epistemic, and total uncertainty throughout training and under two distribution shifts: unseen molecular systems and higher-energy configurations. We find that pretraining reduces prediction error and improves uncertainty-based ranking across all test sets, but does not yield well-calibrated predictive variances; calibration deteriorates more severely under energy extrapolation. Training-dynamics analysis localizes most of this calibration degradation to high-fidelity optimization, during which the learned aleatoric variance becomes overconfident on held-out data. Moreover, multi-fidelity fine-tuning produces a more effective uncertainty-based acquisition function than high-fidelity training from scratch. These findings show that uncertainty ranking and calibration capture distinct properties and must be assessed separately in multi-fidelity atomistic models, especially under extrapolation.
Figures & tables
Figure 1: Overview of the uncertainty estimation pipeline. A mean–variance MACE model maps a structure to an energy and a variance. In an ensemble of such models, the average predicted variance gives the aleatoric uncertainty and the disagreement between members the epistemic uncertainty.
Figure 2: Accuracy and uncertainty for the converged models . (a) RMSE of the ensemble mean, (b) AUSE and (c) ENCE. Colour denotes the protocol (teal HF-only, red LF HF) and marker the signal (circle AU, square EU, diamond TU); in (c) TU, the predictive variance, is highlighted and its components AU and EU are faded. Bars are 95% intervals over five splits, and RMSE and ENCE use log axes with geometric means. The wide HF-only System-ID RMSE interval comes from a single configuration in one split, which TU ranks as the most uncertain of 5,000.
Figure 3: Ranking and calibration behavior under different distribution shifts . ENCE (log ) against AUSE ( ) for the converged models under (a) system shift and (b) energy shift; lower left is best. Colour: protocol; marker: signal, with TU highlighted and AU and EU faded. For every signal a thin line joins the ID (hollow) to the OOD test set (filled).
Figure 4: Uncertainty evolution during training for the multi-fidelity case . LF HF through DFT pretraining (epochs 0–299, scored against DFT labels) and CC fine-tuning (300–399, shaded, CC labels) on its own training data (a, b), the ID (c, d) and the OOD test set (e, f). Note that at epoch 300 the target labels change and, except for Energy-OOD, the evaluated configurations also change; values across the transition are therefore not directly comparable. Colour: split (purple system, brown energy); line style: signal (dashed AU, dotted EU, solid TU), with TU highlighted in the calibration panels; bands: 95% intervals over five splits.
Figure 5: Comparison of uncertainty evolution for the high-fidelity target. AUSE (top) and ENCE (bottom, log scale) on System-OOD (a, c) and Energy-OOD (b, d) against the HF-stage epoch, for HF-only from scratch (teal, 300 epochs) and CC fine-tuning of LF HF (red, 100 epochs). Line style: signal (dashed AU, dotted EU, solid TU); in the calibration panels TU is highlighted. Lines are exponential moving averages of the five-split mean (geometric for ENCE); bands are 95% intervals. Last-epoch values differ from Figure 2 , whose models are best-validation checkpoints.
Figure 6: Uncertainty-guided acquisition. One round acquires 500 of 2,500 Energy-OOD pool configurations; models are evaluated on the disjoint held-out half. (a) RMSE reduction relative to random selection (positive favours uncertainty); markers are five-split means with 95% intervals, grey dots individual splits. (b) Within-system energy quantile of the acquired configurations; grey is random selection, which reproduces the pool.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| Split | DFT | CC |
|---|---|---|
| Train | 50,000 | 5,000 |
| Validation | 10,000 | 1,000 |
| Test-ID | 50,000 | 5,000 |
| Test-OOD | 50,000 | 5,000 |
Table 1: Number of configurations per split for the DFT and CC data sets, identical for both shifts and all five split seeds.
Figure 7: Training dynamics of HF-only ensembles under system shift. Columns: training set, ID and OOD test set. Rows: Spearman ( ), AUSE ( ), ENCE ( , TU highlighted, AU and EU faded), mean per-atom AU and EU, RMSE of the ensemble mean and Gaussian NLL of the predictive distribution, against epoch. Lines are exponential moving averages of the five-split means (geometric on log axes); bands are 95% intervals.
Figure 8: Training dynamics of LF HF ensembles under system shift. As Figure 7 . Epochs 0–299 are DFT pretraining (the LF-only ensemble, DFT labels), epochs 300–399 (shaded) CC fine-tuning; RMSE and NLL jump at epoch 300 partly because the target changes.
Figure 9: Training dynamics of HF-only ensembles under energy shift. As Figure 7 ; the Energy-OOD column is support-filtered.
Figure 10: Training dynamics of LF HF ensembles under energy shift. As Figure 8 ; the Energy-OOD column is support-filtered and, during pretraining, scored against DFT labels.
Figure 11: Training and validation curves. Training loss (Gaussian NLL, solid) and validation loss (dashed), and the corresponding RMSE, for LF-only (DFT), HF-only (CC) and the fine-tuning stage of LF HF (CC), under system (top two rows) and energy shift (bottom two rows). Loss: five-split arithmetic means; RMSE: geometric means; bands are 95% intervals. For HF-only and LF HF the validation loss flattens or rises while the training loss keeps falling, which is why validation-based selection keeps early fine-tuning checkpoints.
Figure 12: EU is smaller than AU at every energy depth and for nearly all of training. (a) Mean EU over mean AU on the OOD test sets during LF HF training, DFT pretraining (epochs 0–299) then CC fine-tuning (300–399, shaded); bands are 95% intervals. (b) Converged energy-split models: median per-configuration EU/AU on Energy-ID and in within-system energy-quantile bins of Energy-OOD, for LF-only (grey, DFT labels), HF-only and LF HF.
Figure 13: Reliability of the converged ensembles under system shift, before recalibration. Rows: LF-only (DFT labels), HF-only and LF HF (CC labels); columns: ID and OOD test set. In each split, the support-filtered test configurations are sorted by each signal and cut into 15 equal-count bins; each point is the five-split geometric mean of the root mean predicted variance (RMV, ) and of the RMSE ( ) of one bin, per atom, on common log-log axes. Bands are 95% log-space intervals of the RMSE over five splits (RMV intervals not drawn); the dashed diagonal marks perfect calibration, and points above it are overconfident. The corner text gives the five-split geometric-mean ENCE of each signal.
Figure 14: Reliability of the converged ensembles under energy shift, before recalibration. As Figure 13 .
| HF-only | LF HF | |||||
|---|---|---|---|---|---|---|
| Test set | AU | EU | TU | AU | EU | TU |
| System-ID | ||||||
| System-OOD | ||||||
| Energy-ID | ||||||
| Energy-OOD | ||||||
Table 2: Spearman rank correlation ( ) between uncertainty and squared error of the converged models, five-split mean 95% half-width.
Figure 15: Isotonic recalibration fitted in domain transfers poorly to energy extrapolation. ENCE of the converged models on (a) System-OOD and (b) Energy-OOD before (hollow) and after (filled) isotonic recalibration fitted on the in-domain validation set, per protocol and signal.
Figure 16: Reliability of the converged ensembles under system shift, after isotonic recalibration. Rows: LF-only (DFT labels), HF-only and LF HF (CC labels); columns: ID and OOD test set. In each split, the support-filtered test configurations are sorted by each signal and cut into 15 equal-count bins; each point is the five-split geometric mean of the root mean predicted variance (RMV, ) and of the RMSE ( ) of one bin, per atom, on common log-log axes. Bands are 95% log-space intervals of the RMSE over five splits (RMV intervals not drawn); the dashed diagonal marks perfect calibration, and points above it are overconfident. The corner text gives the five-split geometric-mean ENCE of each signal. The isotonic map of each signal is fitted on the in-domain validation set of each split.
Figure 17: Reliability of the converged ensembles under energy shift, after isotonic recalibration. As Figure 16 .
Figure 18: Uncertainty-guided acquisition from the Energy-OOD pool costs little in domain. Energy-ID RMSE reduction relative to random selection of the same budget, for the acquisition experiment of Section 4.4 ; negative values favour random selection. Markers are five-split means with 95% intervals, grey dots the individual splits.
More in this topic
Multifidelity Modeling
Multifidelity modeling combines inexpensive approximate models with costly accurate simulations or measurements to improve prediction and decision-making.
See all 36 papers