Improved Uncertainty Estimation of Graph Neural Network\nPotentials Using Engineered Latent Space Distances

Graph neural networks\n(GNNs) have been shown to be astonishingly\ncapable models for molecular property prediction, particularly as\nsurrogates for expensive density functional theory calculations of\nrelaxed energy for novel material discovery. However, one limitation\nof GNNs in this context is the lack of useful uncertainty prediction\nmethods, as this is critical to the material discovery pipeline. In\nthis work, we show that uncertainty quantification for relaxed energy\ncalculations is more complex than uncertainty quantification for other\nkinds of molecular property prediction due to the effect that structure\noptimizations have on the error distribution. We propose that distribution-free\ntechniques are more useful tools for assessing calibration, recalibrating,\nand developing uncertainty prediction methods for GNNs performing\nrelaxed energy calculations. We also develop a relaxed energy task\nfor evaluating uncertainty methods for equivariant GNNs, based on\ndistribution-free recalibration and using the Open Catalyst Project\ndata set. We benchmark a set of popular uncertainty prediction methods\non this task and show that conformal prediction methods, with our\nnovel latent space distance improvements, are the most well-calibrated\nand efficient approach for relaxed energy calculations. Finally, we\ndemonstrate that our latent space distance method produces results\nwhich align with our expectations on an atom clustering example, and\non specific equation of state and adsorbate coverage examples from\noutside the training data set.