Source-linked AI summary
Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules
Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, Stephan Günnemann
TL;DR
Machine-learning models largely focus on equilibrium molecules, limiting evidence for reactive transition-state configurations far from equilibrium. The paper proposes DimeNet++, develops the COLL collision dataset, and evaluates uncertainty methods; DimeNet++ is 8x faster than DimeNet while improving predictions by 10% on average, whereas the uncertainty approaches remain insufficient.
Problem
Most molecular machine-learning models focus on equilibrium or near-equilibrium structures, while reactions require predictions for far-from-equilibrium transition states spanning broader chemical and configurational space.
Method
The paper redesigns DimeNet as DimeNet++, constructs COLL from reactive molecular-collision configurations with DFT references, and evaluates ensembling and mean-variance estimation for uncertainty quantification.
Results
DimeNet++ achieves an 8x runtime improvement and 10% average prediction improvement over DimeNet, while reducing COLL error by 76% on average at only 17% greater runtime than SchNet.
Takeaways & Limitations
DimeNet++ extends molecular prediction toward reactive non-equilibrium configurations, but reliable and fast uncertainty estimates remain an important direction for future work.
Takeaways & Limitations
Ensembling has considerable computational overhead, and no general method is available for estimating force uncertainty for models that provide only energy mean and standard deviation.
Abstract
from arXiv · showhide
Many important tasks in chemistry revolve around molecules during reactions. This requires predictions far from the equilibrium, while most recent work in machine learning for molecules has been focused on equilibrium or near-equilibrium states. In this paper we aim to extend this scope in three ways. First, we propose the DimeNet++ model, which is 8x faster and 10% more accurate than the original DimeNet on the QM9 benchmark of equilibrium molecules. Second, we validate DimeNet++ on highly reactive molecules by developing the challenging COLL dataset, which contains distorted configurations of small molecules during collisions. Finally, we investigate ensembling and mean-variance estimation for uncertainty quantification with the goal of accelerating the exploration of the vast space of non-equilibrium structures. Our DimeNet++ implementation as well as the COLL dataset are available online.
1 Introduction
The paper targets molecular prediction far from equilibrium, where reactions involve transition states and broader chemical and configurational space. It proposes DimeNet++, introduces the COLL dataset, and evaluates uncertainty-estimation approaches for non-equilibrium exploration.
- Motivation: Equilibrium-focused molecular models do not cover reaction transition states with open-shell structures, stretched bonds, and distorted angles.Reliable prediction therefore requires broader chemical and configurational coverage.
- Contributions: 8x runtime improvement and 10% average prediction improvement are achieved by DimeNet++ over DimeNet, with energy predictions improving by 20%.The model predicts molecular energies and atomic forces and generalizes across configurations and molecules.
- Contributions: 140 000 collision configurations comprise the COLL dataset, covering wider energy and force ranges than previous benchmarks.The dataset contains pairs of small molecules reacting at high kinetic energies.
- Contributions: Ensembling and mean-variance estimation are investigated to detect when non-equilibrium structures leave the training-data region.The paper concludes that both approaches are insufficient because of computational overhead and unreliable energy and force uncertainties.
2 DimeNet++
DimeNet++ retains DimeNet’s directional message-passing basis while redesigning its costly architecture. The changes target computationally expensive interaction representations and embedding hierarchies without sacrificing accuracy.
- DimeNet foundation: DimeNet updates messages between atoms, allowing it to model bond angles and interatomic distances jointly.It uses a spherical 2D Fourier-Bessel basis to embed distances and angles.
- Architecture bottleneck: DimeNet uses around 15x as many message embeddings and triplet representations as atoms on QM9, making directional message passing especially expensive.The interaction block is approximately 15x more expensive than other model components.
- Fast interactions: DimeNet++ replaces the expensive bilinear interaction with a Hadamard product and adds MLPs for basis representations.This recovers original accuracy at a fraction of the computational cost.
- Embedding hierarchy: Down- and upprojection layers reduce embedding sizes in high-volume parts of the model while removing information bottlenecks.The hierarchy leverages unequal numbers of embeddings across model components.
- Other improvements: Four layers perform on par with the original six for U0, while larger batches slow convergence and mixed precision causes complete precision breakdown.The reported precision issue is consistent with DimeNet’s relative error being below float16 machine precision, 5 · 10−4.
3 COLL Dataset
COLL is built from molecular-collision simulations to benchmark highly distorted, reactive configurations. Its reference energies and forces are recomputed with density functional theory under procedures designed to address difficult electronic-structure behavior.
- Dataset construction: 140 000 random collision snapshots form COLL, with energies and forces recomputed using density functional theory.Collision trajectories were generated with GFN2-xTB, then evaluated using revPBE, def2-TZVP, and D3 dispersion corrections.
- Configurational coverage: COLL contains highly distorted structures with stretched bonds, distorted angles, and prevalent open-shell electronic structures.All calculations use broken spin-symmetry, unlike equilibrium or near-equilibrium benchmark sets such as QM9.
- Electronic-structure treatment: Multiple randomized initial wavefunctions and 5000 K Fermi-smearing address multiple self-consistent-field solutions and convergence challenges.The lowest-energy solution is selected to reduce discontinuities in the reference potential energy surfaces.
4 Uncertainty Quantification
The paper examines uncertainty estimates for both energies and forces in non-equilibrium molecular dynamics. Ensembling and mean-variance estimation have important limitations: ensembles are costly, while energy uncertainty does not generally provide force uncertainty.
- Motivation: Non-equilibrium dynamics require uncertainty estimates for both energy E and force F because systems can leave the training distribution.The vast state space reachable in high-energy simulations motivates detecting degraded predictive performance.
- Ensembling: Ensembling is differentiable and can estimate force error, but requires predictions from multiple separate models, creating substantial computational overhead.
- Mean-variance estimation: Mean-variance estimation provides mean and standard deviation for energy, but differentiating the energy mean does not yield an analogous force uncertainty.
- Limitations: There is no general way to estimate σF for these models, so σE must serve as a proxy whose correlation with force error is uncertain.
5 Experiments
Experiments evaluate DimeNet++ on QM9 and the reactive COLL benchmark, testing accuracy, runtime, and uncertainty estimation. DimeNet++ improves the accuracy–speed trade-off and outperforms SchNet on COLL, while reliable fast uncertainty estimation remains unresolved.
- DimeNet++ improvements: Replacing the bilinear layer with a Hadamard product decreases runtime by a factor of 5, while smaller embeddings both accelerate the model and improve accuracy.Each proposed change either reduces runtime or improves error on the QM9 U0 validation set.
- QM9: DimeNet++ performs best overall on QM9 and for most properties individually, while being 8x faster than DimeNet.The comparison includes SchNet, MGCN, and DeepMoleNet.
- COLL: 76 % lower average error and only 17 % slower than SchNet are reported for DimeNet++ on the COLL dataset.The table reports MAE in eV and eV/Å.
- COLL: COLL is significantly more challenging than QM9, with both SchNet and DimeNet++ showing MAE around 10x higher than on QM9.
- Uncertainty quantification: Ensembling estimates force error well but struggles with energy uncertainty, while MVE is especially poor for force uncertainty and still leaves substantial overhead unresolved.Reliable and fast uncertainty estimates remain an important direction for future work.