How to Train a Shallow Ensemble

TL;DR

This paper investigates training strategies for shallow ensembles in machine learning interatomic potentials, emphasizing calibration, computational efficiency, and the importance of modeling force uncertainties for reliable uncertainty quantification.

Contribution

It introduces an efficient fine-tuning protocol for shallow ensembles that maintains calibration quality while significantly reducing training time.

Findings

Explicit NLL optimization improves calibration.

Modeling force uncertainties is crucial for reliable calibration.

Fine-tuning reduces training time by up to 96% without losing calibration quality.

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.