Uncertainty Estimation for Molecules: Desiderata and Methods
by Tom Wollschläger, Nicholas Gao, Bertrand Charpentier and Stephan Günnemann
Published at the 40th International Conference on Machine Learning (ICML), 2023
Abstract
Graph Neural Networks (GNNs) are promising surrogates for quantum mechanical calculations as they establish unprecedented low errors on collections of molecular dynamics (MD) trajectories. Thanks to their fast inference times they promise to accelerate computational chemistry applications. Unfortunately, despite low in-distribution (ID) errors, such GNNs might be horribly wrong for out-of-distribution (OOD) samples. Uncertainty estimation (UE) may aid in such situations by communicating the model's certainty about its prediction. Here, we take a closer look at the problem and identify six key desiderata for UE in molecular force fields, three 'physics-informed' and three 'application-focused' ones. To overview the field, we survey existing methods from the field of UE and analyze how they fit to the set desiderata. By our analysis, we conclude that none of the previous works satisfies all criteria. To fill this gap, we propose Localized Neural Kernel (LNK) a Gaussian Process (GP)-based extension to existing GNNs satisfying the desiderata. In our extensive experimental evaluation, we test four different UE with three different backbones and two datasets. In out-of-equilibrium detection, we find LNK yielding up to 2.5 and 2.1 times lower errors in terms of AUC-ROC score than dropout or evidential regression-based methods while maintaining high predictive performance.
Cite
Please cite our paper if you use the model, experimental results, or our code in your own work:
@inproceedings{wollschlaeger_uq4mol_2023,
title = {Uncertainty Estimation for Molecules: Desiderata and Methods},
author = {Wollschl{\"a}ger, Tom and Gao, Nicholas and Charpentier, Bertrand and Ketata, Mohamed Amine and G{\"u}nnemann, Stephan},
booktitle={International Conference on Machine Learning (ICML)},
year = {2023}
}