Improved Uncertainty Estimation of Graph Neural Network Potentials Using Engineered Latent Space Distances
Joseph Musielewicz, Janice Lan, Matt Uyttendaele and John R. Kitchin, “Improved Uncertainty Estimation of Graph Neural Network Potentials Using Engineered Latent Space Distances”, The Journal of Physical Chemistry C 128(49), 20799-20810 (2024). https://doi.org/10.1021/acs.jpcc.4c04972
This summary was written by a language model from the paper and has not been checked line by line. Rely on the published paper for details.
Note: the extracted text available for this summary is the arXiv preprint (arXiv:2407.10844v2, August 2024) with its supporting information, so details may differ slightly from the journal article.
Question
Graph neural network potentials such as EquiformerV2 and GemNet-OC are now used in workflows like AdsorbML, and in the public OCP demo, to find minimum adsorption energies with little or no DFT. Without DFT checks, users need to know when a relaxed energy predicted by the network can be trusted. The paper asks how to build and fairly evaluate uncertainty estimates for this relaxed structure to relaxed energy (RS2RE) task, and which family of methods gives the best calibrated uncertainties (Musielewicz et al. 2024).
Approach
The authors relaxed the roughly 25,000 structures of the OC20 in-domain validation set with the public 31M-parameter EquiformerV2 checkpoint and computed VASP single points on the final frames; this served as the calibration set, and the same procedure on the out-of-domain-both validation set gave the test set. They benchmarked four classes of uncertainty estimator for EquiformerV2’s relaxed energies. Ensembles were trained on the structure to energy and forces task: an architecture ensemble (11 members mixing GemNet-OC, eSCN and EquiformerV2), a bootstrap ensemble (10 EquiformerV2 models each trained on 66% of OC20) and a parameter ensemble (6 EquiformerV2 models of different sizes), with variance taken at the first frame, the last frame, or the mean or max over the relaxation trajectory. Latent distance methods indexed latent vectors of the roughly 460,000 relaxed training systems with FAISS and used the nearest-neighbor L2 distance, either between system-averaged vectors or per atom (then combined by mean, max or sum). Latent spaces came from the full equivariant EquiformerV2 embedding, its rotationally invariant l=0 node channel, edge-aligned m=0 channels, or other networks (GemNet-OC, PaiNN, eSCN). The remaining classes were mean variance estimation heads fine-tuned on residuals and a transformer sequence regressor reading latent vectors along the trajectory.
Because the RS2RE error distribution is sharply peaked at zero and far from Gaussian (much more so than for the structure to energy and forces task), the authors argue for distribution-free validation. Every method was recalibrated with a linear fit of binned RMSE against root mean variance (20 bins) on the calibration set, then judged on the test set by the R2 of the binned points against the parity line and by Pernot’s bootstrap confidence interval test on Var(Z), whose ideal value is 1.
Key findings
The common metrics (miscalibration area before and after recalibration, negative log likelihood, Spearman rank correlation, AUROC) disagreed with one another about which method was best, while the two distribution-free calibration measures agreed. The best latent distance method, GemNet-OC node embeddings with the sum of per-atom distances, reached a parity R2 of 0.952 with a test-set CI(Var(Z)) of [1.13, 1.79], ahead of the best ensemble (architecture ensemble with trajectory-mean variance, 0.905 and [1.42, 2.57]), an MVE head (0.849 and [1.66, 2.46]) and the residual sequence model (0.331 and [3.38, 4.47]). All methods passed the global calibration test on the calibration set but became miscalibrated on the out-of-domain test set. The ensemble showed local overconfidence between about 0.3 and 1 eV and a recalibration-induced floor on its uncertainties, while the MVE head became overconfident above about 0.2 eV, and both residual-fitting approaches overfit.
For ensembles, member diversity mattered most (architecture beat bootstrap and parameter ensembles), and the trajectory-mean variance was generally best while the first frame was always worst. For distances, the choice of latent space dominated: the full equivariant EquiformerV2 embedding was noisy (UMAP plots showed overlapping Cu, Pd and Pt clusters that separated in the invariant space), and a system-mean distance on it gave a parity R2 of -14.7. Per-atom distances generally beat system-averaged ones, and edge-aligned invariant embeddings beat the l=0 node embedding. Strikingly, the invariant latent space of the less accurate GemNet-OC predicted EquiformerV2’s error better than EquiformerV2’s own. In interpretable tests, the GemNet-OC distance gave very large uncertainties for Cu bulk and larger still for TiO2 bulk (neither bulk materials nor oxides are in OC20), smaller uncertainties for a strained Al slab that grew away from the minimum, and a smooth, monotonic rise for increasing O coverage on Pt that tracked the measured errors.
Significance
The work gives the community a concrete RS2RE uncertainty benchmark with a recalibration protocol, and shows that a cheap post hoc nearest-neighbor distance in a well-chosen, rotationally invariant latent space can beat expensive ensembles for screening workflows such as AdsorbML. It also identifies rotational equivariance as an obstacle to meaningful latent distances, pointing to future architectures or distance measures that keep equivariance while yielding useful similarity.