Uncertainty-aware predictions of molecular X-ray absorption spectra using neural network ensembles

TL;DR

This paper demonstrates how neural network ensembles can predict molecular X-ray absorption spectra and quantify uncertainty, enabling better out-of-sample predictions and potential integration with active learning for structure refinement.

Contribution

It introduces an uncertainty-aware neural network ensemble approach for predicting molecular spectra and their uncertainties, applicable to larger molecules and out-of-sample data.

Findings

Model provides accurate spectra predictions with reliable uncertainty estimates.

Uncertainty bounds correlate with actual prediction errors.

Method generalizes to larger molecules and distorted structures.

Abstract

As machine learning (ML) methods continue to be applied to a broad scope of problems in the physical sciences, uncertainty quantification is becoming correspondingly more important for their robust application. Uncertainty aware machine learning methods have been used in select applications, but largely for scalar properties. In this work, we showcase an exemplary study in which neural network ensembles are used to predict the X-ray absorption spectra of small molecules, as well as their point-wise uncertainty, from local atomic environments. The performance of the resulting surrogate clearly demonstrates quantitative correlation between errors relative to ground truth and the predicted uncertainty estimates. Significantly, the model provides an upper bound on the expected error. Specifically, an important quality of this uncertainty-aware model is that it can indicate when the model is predicting on out-of-sample data. This allows for its integration with large scale sampling of structures together with active learning or other techniques for structure refinement. Additionally, our models can be generalized to larger molecules than those used for training, and also successfully track uncertainty due to random distortions in test molecules. While we demonstrate this workflow on a specific example, ensemble learning is completely general. We believe it could have significant impact on ML-enabled forward modeling of a broad array of molecular and materials properties.