A Simple and Scalable Kernel Density Approach for Reliable Uncertainty Quantification in Atomistic Machine Learning

Abstract

Machine learning models are increasingly used to predict material properties and accelerate atomistic simulations, but the reliability of their predictions depends on the representativeness of the training data. We present a scalable, GPU-accelerated uncertainty quantification framework based on $k$-nearest-neighbor kernel density estimation (KDE) in a PCA-reduced descriptor space. This method efficiently detects sparsely sampled regions in large, high-dimensional datasets and provides a transferable, model-agnostic uncertainty metric without requiring retraining costly model ensembles. The framework is validated across diverse case studies varying in: i) chemistry, ii) prediction models (including foundational neural network), iii) descriptors used for KDE estimation, and iv) properties whose uncertainty is sought. In all cases, the KDE-based score reliably flags extrapolative configurations, correlates well with conventional ensemble-based uncertainties, and highlights regions of reduced prediction trustworthiness. The approach offers a practical route for improving the interpretability, robustness, and deployment readiness of ML models in materials science.

References in corpus (14)

• On representing chemical environments
• On-the-fly machine learning force field generation: Application to melting points
• Phase transitions of hybrid perovskites simulated by machine-learning force fields trained on-the-fly with Bayesian inference
• Sample Size Planning for Classification Models
• SchNetPack: A Deep Learning Toolbox For Atomistic Systems
• MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
• A foundation model for atomistic materials chemistry
• Neural Network Potentials for Chemistry: Concepts, Applications and Prospects
• A First Principles Theory of Nuclear Magnetic Resonance J-Coupling in solid-state systems
• Single-model uncertainty quantification in neural network potentials does not consistently outperform model ensembles
• Differentiable sampling of molecular geometries with uncertainty-based adversarial attacks
• Fast Uncertainty Estimates in Deep Learning Interatomic Potentials
• A reactive neural network framework for water-loaded acidic zeolites
• Optimal Training of Mean Variance Estimation Neural Networks