Calibrated Uncertainty for Molecular Property Prediction

Computational analysis of atomic structures has the potential to speed up the discovery of novel molecules and materials. This process can be accelerated by data-driven methods based on machine learning that are generally less computationally demanding than traditional quantum mechanical methods such as density functional theory (DFT) [1, 2]. Reliable uncertainty estimates are important to assess confidence in predictions and thereby enable decision making and automation in applications such as high-throughput screening and active learning [3]. However, machine learning models can produce badly calibrated uncertainty estimates and it is therefore crucial to detect and handle uncertainty carefully.<br/><br/>We have recently developed a complete framework for training and evaluating ensembles of message passing neural network models that can produce accurate predictions of properties of molecules with well calibrated uncertainty estimates in and out of the training data distribution. The proposed method differs from previous work by considering both aleatoric and epistemic uncertainty in a unified framework, and by recalibrating the predictive distribution on unseen data. We demonstrate the approach through computer experiments on two publicly available benchmark datasets for molecular property prediction, QM9 [4] and PC9 [5].<br/><br/>The model achieves good predictive performance across all experiments and extending the model to include predictive uncertainty did not reduce accuracy compared to the original model without uncertainty. The uncertainty estimates are generally well calibrated such that high uncertainty is assigned to high error instances and low uncertainty is assigned to low error instances on average. We test out of distribution performance by training on QM9 and testing on examples exclusive to the more diverse PC9 which leads to relative high error but still with good uncertainty calibration.<br/><br/>The proposed method provides a general framework for training and evaluating neural network ensemble models that are able to produce accurate predictions of properties of molecules with well calibrated uncertainty estimates in and out of the training data distribution.<br/><br/>References:<br/>[1] Pavlo O Dral. Quantum chemistry in the age of machine learning.J. Phys.Chem. Lett., 11(6):2336–2347, March 2020.<br/>[2] O. Anatole von Lilienfeld and Kieron Burke. Retrospective on a decadeof machine learning for chemical discovery.Nature Communications,11(1):4895, Sep 2020.<br/>[3] Kevin Tran, Willie Neiswanger, Junwoong Yoon, Qingyang Zhang, EricXing, and Zachary W Ulissi. Methods for comparing uncertainty quan-tifications for material property predictions.Mach. Learn.: Sci. Technol.,1(2):025006, May 2020.<br/>[4] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O AnatoleVon Lilienfeld. Quantum chemistry structures and properties of 134 kilomolecules.Scientific data, 1(1):1–7, 2014.<br/>[5] Marta Glavatskikh, Jules Leguy, Gilles Hunault, Thomas Cauchy, andBenoit Da Mota. Dataset’s chemical diversity limits the generalizabilityof machine learning predictions.Journal of Cheminformatics, 11, 11 2019.