Active Learning of Graph Neural Network Potentials from Quantum and Experimental Data
Project Description
This project develops data‑efficient methods for training neural network (NN) potentials to model molecular interactions in materials science, such as additive manufacturing. Because experimental and density functional theory (DFT) data are costly, active learning is used to identify high‑uncertainty configurations that require labeling. Reliable uncertainty quantification (UQ) is therefore essential. Conventional ensemble‑based UQ often yields overconfident estimates, while Bayesian Markov chain Monte Carlo (MCMC) methods are accurate but computationally prohibitive for NN potentials with often hundreds of different parameter sets (see Figure below). To address this, the project investigates efficient single‑pass UQ approaches. These include deep‑kernel learning with Gaussian processes, enabling uncertainty estimates for energies and forces, and evidential deep learning, which predicts a Student‑t distribution to separate epistemic from aleatoric uncertainty. Benchmarking on Lennard‑Jones molecular dynamics (MD) datasets evaluates each method’s ability to detect out‑of‑distribution states.
Results
- Demonstrated that standard neural‑network ensembles frequently produce overconfident and unreliable uncertainty estimates, making them unsuitable for guiding active learning in molecular simulations.
- Verified that stochastic‑gradient guided Monte Carlo (SGGMC) provides more principled Bayesian UQ, effectively reducing overconfidence across diverse molecular configurations.
- Identified the computational limitations of stochastic‑gradient MCMC methods, as they require evaluating hundreds of parameter samples, making them impractical for screening long MD trajectories.
- Developed a deep‑kernel learning framework using Gaussian processes in the embedding space of a NN to quantify uncertainties for both energies and forces.
- Implemented evidential deep learning, enabling prediction of a Student‑t distribution and separation of aleatoric vs. epistemic uncertainty, which is crucial for active learning decisions.
- Generated a comprehensive Lennard‑Jones MD benchmark dataset across temperatures and densities to evaluate UQ methods under out‑of‑distribution conditions.
- Confirmed that the new UQ approaches better identify high‑error states and correlate more closely with true prediction errors.
Follow-up
Future efforts should focus on improving the robustness of the UQ methods either through regularization to prevent collapse or by developing alternative approaches to model force uncertainties. These methods will be used to train uncertainty-aware NN potentials from experimental data.
Another approach will be to combine the DiffTRe method with our efficient UQ schemes to create NN potentials that are consistent with experimental data while providing reliable uncertainty estimates.
First, this enables active learning guided by experiments at specific state-points where the model signals high uncertainty. From this, an exhaustive database of NN potentials consistent with experiments can be built. Second, practitioners running MD simulations via our NN potentials may obtain an uncertainty estimate at negligible additional cost, thereby overcoming the high computational burden of current MCMC approaches.
Thaler, Stephan; Fuchs, Paul; Cukarska, Ana; Zavadlav, Julija: JaxSGMC: Modular stochastic gradient MCMC in JAX. SoftwareX 26, 2024, 101722, https://doi.org/10.1016/j.softx.2024.101722
Thaler, Stephan; Mayr, Felix; Thomas, Siby; Gagliardi, Alessio; Zavadlav, Julija: Active learning graph neural networks for partial charge prediction of metal-organic frameworks via dropout Monte Carlo. npj Computational Materials 10 (1), 2024, https://doi.org/10.1038/s41524-024-01277-8
Thaler, Stephan; Doehner, Gregor; Zavadlav, Julija: Scalable Bayesian Uncertainty Quantification for Neural Network Potentials: Promise and Pitfalls. Journal of Chemical Theory and Computation 19 (14), 2023, 4520–4532, https://doi.org/10.1021/acs.jctc.2c01267