Robust and scalable uncertainty estimation with conformal prediction for machine-learned interatomic potentials

Abstract

Uncertainty quantification (UQ) is important to machine learning (ML) force fields to assess the level of confidence during prediction, as ML models are not inherently physical and can therefore yield catastrophically incorrect predictions. Established a-posteriori UQ methods, including ensemble methods, the dropout method, the delta method, and various heuristic distance metrics, have limitations such as being computationally challenging for large models due to model re-training. In addition, the uncertainty estimates are often not rigorously calibrated. In this work, we propose combining the distribution-free UQ method, known as conformal prediction (CP), with the distances in the neural network's latent space to estimate the uncertainty of energies predicted by neural network force fields. We evaluate this method (CP+latent) along with other UQ methods on two essential aspects, calibration, and sharpness, and find this method to be both calibrated and sharp under the assumption of independent and identically-distributed (i.i.d.) data. We show that the method is relatively insensitive to hyperparameters selected, and test the limitations of the method when the i.i.d. assumption is violated. Finally, we demonstrate that this method can be readily applied to trained neural network force fields with traditional and graph neural network architectures to obtain estimates of uncertainty with low computational costs on a training dataset of 1 million images to showcase its scalability and portability. Incorporating the CP method with latent distances offers a calibrated, sharp and efficient strategy to estimate the uncertainty of neural network force fields. In addition, the CP approach can also function as a promising strategy for calibrating uncertainty estimated by other approaches.