Flexible Uncertainty Calibration for Machine-Learned Interatomic Potentials

TL;DR

This work addresses the challenge of reliable uncertainty quantification for machine-learned interatomic potentials by reimagining conformal prediction as a learnable, environment-dependent calibration. By replacing the global CP quantile with a neural-parameterized quantile function $q_{\theta}(\mathbf{X})$ trained via a pinball-loss objective (and optionally a weighted variant), the method delivers adaptive, site-resolved predictive intervals with strong calibration and sharpness. Across ionic, catalytic, and molecular benchmarks using the MACE-MP-0b3 foundation model, the approach achieves substantial improvements in uncertainty–error alignment (e.g., $\rho$ rising from $0.386$ to $0.589$ on LiCl) while incurring negligible computational overhead and enabling robust transfer across exchange–correlation functionals (e.g., HEA25 with $\rho=0.625$). The framework supports active learning in MD, generalizes to unseen atomic environments, and reduces labeling costs by maintaining reliable calibrated uncertainties under cross-domain shifts, paving the way for more trustworthy atomistic simulations.

Abstract

Reliable uncertainty quantification (UQ) is essential for developing machine-learned interatomic potentials (MLIPs) in predictive atomistic simulations. Conformal prediction (CP) is a statistical framework that constructs prediction intervals with guaranteed coverage under minimal assumptions, making it an attractive tool for UQ. However, existing CP techniques, while offering formal coverage guarantees, often lack accuracy, scalability, and adaptability to the complexity of atomic environments. In this work, we present a flexible uncertainty calibration framework for MLIPs, inspired by CP but reformulated as a parameterized optimization problem. This formulation enables the direct learning of environment-dependent quantile functions, producing sharper and more adaptive predictive intervals at negligible computational cost. Using the foundation model MACE-MP-0 as a representative case, we demonstrate the framework across diverse benchmarks, including ionic crystals, catalytic surfaces, and molecular systems. Our results show order-of-magnitude improvements in uncertainty-error correlation, enhanced data efficiency in active learning, and strong generalization performance, together with reliable transfer of calibrated uncertainties across distinct exchange-correlation functionals. This work establishes a principled and data-efficient approach to uncertainty calibration in MLIPs, providing a practical route toward more trustworthy and transferable atomistic simulations.

Figures (14)

• Figure 1: Conceptual illustration of the uncertainty calibration framework. (A) Workflow and calibration schemes. (B) Representative comparison of calibration outcomes. Before calibration and regular CP show systematic misalignment between predicted and true errors. Class CP improves alignment by grouping atomic environments into discrete classes, shown in different colors. Flexible UC further refines this by learning a smooth, environment-dependent mapping.
• Figure 2: Comparison of uncertainties from LLPR, regular CP, class-based CP, and flexible UC on LiCl dataset. $\rho$ denotes the spearman rank correlation coefficient.
• Figure 3: Uncertainty estimates from LLPR, regular CP, class CP, and flexible UC on catalytic surface data. Left: undoped configurations (in-calibration). Right: Pt-doped configurations (out-of-calibration). Regular and class-based CP show similar behavior in both cases, whereas flexible UC substantially improves generalization to unseen Pt-doped environments.
• Figure 4: Species-resolved calibration of force uncertainties on Pt-doped surfaces. Predicted force uncertainties versus actual errors are shown for different elements. Top: regular conformal prediction. Bottom: flexible calibration. Although Pt atoms are excluded from the calibration set, the proposed method yields improved correlation ($\rho$) across all species.
• Figure 5: Uncertainty estimates from LLPR and flexible UC, together with true force errors, for a representative Pt-doped surface configuration (atomic structure shown on the left). Errors and uncertainties are assigned on a per-atom basis, with darker colors indicating atoms exhibiting larger values.