[Paper Review] How to Train a Shallow Ensemble

Moritz Schäfer, Matthias Kellner+2 more|arXiv (Cornell University)|Feb 17, 2026

Machine Learning in Materials Science0 citations

TL;DR

The paper analyzes training strategies for shallow ensembles in ML interatomic potentials, showing that explicit energy and force uncertainty losses improve calibration, and that full-model fine-tuning can achieve comparable calibration with much lower cost than training from scratch.

ABSTRACT

Shallow ensembles provide a convenient strategy for uncertainty quantification in machine learning interatomic potentials, that is computationally efficient because the different ensemble members share a large part of the model weights. In this work, we systematically investigate training strategies for shallow ensembles to balance calibration performance with computational cost. We first demonstrate that explicit optimization of a negative log-likelihood (NLL) loss improves calibration with respect to approaches based on ensembles of randomly initialized models, or on a last-layer Laplace approximation. However, models trained solely on energy objectives yield miscalibrated force estimates. We show that explicitly modeling force uncertainties via an NLL objective is essential for reliable calibration, though it typically incurs a significant computational overhead. To address this, we validate an efficient protocol: full-model fine-tuning of a shallow ensemble originally trained with a probabilistic energy loss, or one sampled from the Laplace posterior. This approach results in negligible reduction in calibration quality compared to training from scratch, while reducing training time by up to 96%. We evaluate this protocol across a diverse range of materials, including amorphous carbon, ionic liquids (BMIM), liquid water (H$_2$O), barium titanate (BaTiO$_3$), and a model tetrapeptide (Ac-Ala3-NHMe), establishing practical guidelines for reliable uncertainty quantification in atomistic machine learning.

Motivation & Objective

• Motivate calibrated uncertainty quantification (UQ) for ML interatomic potentials to support reliable simulations.
• Compare different shallow ensemble and last-layer approaches for energy and force uncertainty estimation.
• Identify training strategies that balance calibration quality with computational cost.
• Provide practical guidelines and benchmarks across diverse materials for uncertainty-aware potentials.

Proposed method

• Use Gaussian negative log-likelihood (NLL) loss to train a shallow ensemble by sharing the backbone and training the last layer jointly (DPOSE).
• Compare energy-focused NLL training (SE E) with last-layer Laplace-based LLPR methods (LLPR E).
• Incorporate force uncertainties by training with a force loss (NLL for forces) and/or including force Hessian contributions (E,F variants).
• Apply post-hoc calibration to predict uncertainties via a scalar alpha scaling of uncertainties (variance) on held-out data.
• Evaluate calibration with Relative Log-Likelihood (RLL) across multiple datasets and architectures; assess training costs and the impact of full-model fine-tuning vs last-layer updates.

Experimental results

Research questions

• RQ1 Do energy-uncertainty calibrated shallow ensembles (SE E) and LLPR-based methods (LLPR E) provide reliable energy uncertainty estimates across diverse materials?
• RQ2 Can incorporating force uncertainty (E,F) or a force NLL improve force calibration without prohibitive cost?
• RQ3 Is full-model fine-tuning necessary to recover good calibration when LLPR or force-informed training underperforms with a frozen backbone?
• RQ4 How do these approaches perform across diverse datasets (mol liquids, solids, and biomolecules) and architectures?

Key findings

• Shallow ensembles trained with energy NLL (SE E) yield consistently positive energy RLL across all tested datasets, indicating informative energy uncertainty.
• LLPR E (energy-only) often yields negative RLLs on several datasets, showing weaker energy-calibration performance and failure modes on some outlier structures.
• Full-model fine-tuning of LLPR E initializations can recover calibration quality, achieving RLL comparable to scratch-trained shallow ensembles, while reducing training time substantially.
• Incorporating force uncertainties requires a force NLL loss (E,F) or including force Hessian contributions (LLPR E,F) to avoid miscalibration of forces; energy-only calibration fails to provide reliable force estimates.
• Independent energy/force post-hoc calibration (separate alpha_E and alpha_F) helps some datasets but can fail for others, and element-wise miscalibration (e.g., Boron/Fluorine in BMIM) can persist.
• Element-specific miscalibration persists for LLPR E and SE E in BMIM, which motivates force-informed training to correct local uncertainties.