Uncertainty quantification in graph neural networks with shallow ensembles
Tirtha Vinchurkar, Kareem Abdelmaqsoud and John R. Kitchin, “Uncertainty quantification in graph neural networks with shallow ensembles”, Machine Learning: Science and Technology 6(4), 045007 (2025). https://doi.org/10.1088/2632-2153/ae0bf0
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:2504.12627v1, April 2025), not the final published version, so details may differ slightly from the journal article.
Question
Graph neural network potentials are fast and accurate inside their training distribution, but they can fail silently on structures unlike those they were trained on, which is exactly where materials discovery tends to look. Deep ensembles are the usual benchmark for flagging such cases, yet they require training many models. The paper asks whether a much cheaper scheme, Direct Propagation of Shallow Ensembles (DPOSE, introduced by Kellner and Ceriotti), can give a graph neural network uncertainty estimates that are low for in-domain inputs and high for out-of-domain inputs, across molecules, catalytic surfaces and metal molecular dynamics frames (Vinchurkar et al. 2025).
Approach
DPOSE keeps a shared network body and replaces only the last layer with many output heads, so the ensemble costs about as much as one extra hidden layer. The mean of the heads is the prediction and their variance is the uncertainty, and training uses a Gaussian negative log-likelihood loss that balances squared error against the predicted variance. The authors split the final layer of SchNet into 64 heads (the “SchNet Ensemble”, SE) and built three models. QM9-SE was trained on all of QM9 (130,831 equilibrium organic molecules with C, H, O, N, F and up to nine heavy atoms) for 1000 epochs with Adam at a learning rate of 5e-5, batch size 100 and an 80:20 split. OC20-SE was made by fine-tuning only the new final layer of a SchNet checkpoint pretrained on all of OC20, using intermetallic slabs, with nonmetal slabs held out. AuMD-SE fine-tuned an OC20-pretrained SchNet on 797 bulk gold frames (FCC, BCC, SC, HCP and diamond cubic, equilibrium and non-equilibrium) from an earlier gold dataset, holding out 6,278 amorphous cluster frames and, later, silver structures. Out-of-domain tests included stretching or compressing a C-C bond, introducing elements absent from training, increasing molecular size, and compressing or expanding the volume per atom.
Key findings
For QM9, the variance was smallest at the equilibrium C-C bond length and grew as the bond was stretched or compressed in ethane, cyanogen, acetonitrile and glyoxal. Unseen elements produced large jumps: CF4 had a variance of about 0.09 while CCl4 reached 62, CH4 at 0.09 rose to 2.7 for SiH4, HF at 0.21 compared with 3.2 for HCl, and SiCl4 reached 56. Uncertainty also rose with chain length beyond the training range, from about 0.09 for methane to 0.66 for octadecane and from 0.08 for methanol to 0.70 for octadecanol, with values staying flat up to about C8 and increasing (not monotonically) from C9 onward.
For OC20, the fine-tuned model was accurate on intermetallic slabs, with parity R2 near 1, and less accurate on nonmetal slabs. The median variance for nonmetal slabs was about 0.03, while the upper whisker for intermetallic slabs stayed below 0.01. Under volume compression and expansion, intermetallic slabs showed growing energy error away from equilibrium but essentially constant uncertainty (a spread of about 0.002 eV/atom), so fine-tuning made the model overconfident for strained in-domain structures. For nonmetal slabs the predicted energy minimum was misplaced relative to the true equilibrium volume, yet the variance did behave as expected, being lowest near the predicted minimum and larger away from it.
For gold, bulk frames were predicted more accurately than amorphous frames, but the mean variance was slightly higher for bulk, with maximum variances differing by only about 0.001 eV/atom, so DPOSE did not separate the two morphologies. Sorting frames by energy showed that both variance and error increased at higher energies for FCC, BCC and amorphous structures, which the authors attribute to fewer high-energy frames in the data; they suggest short-range order in amorphous gold resembles the bulk lattices. Changing the element did separate domains clearly: mean variance was 0.00022 eV/atom for gold and 0.61304 eV/atom for silver. The overall conclusion is that DPOSE flags out-of-domain data when it is far from the training set, especially new atom types, but struggles to distinguish structurally different systems of the same element.
Significance
The work shows that a plug-in, nearly free modification of an existing GNN potential can supply useful out-of-domain warnings, and it is candid about where the signal fails (overconfident strained intermetallics, indistinguishable bulk and amorphous gold). These limitations matter for anyone hoping to use ensemble variance as an active learning acquisition signal, which the authors name as the next step.