Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules

Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, Stephan Günnemann

TL;DR

The paper tackles molecular property prediction in non-equilibrium regimes by introducing DimeNet++ to accelerate directional message passing without sacrificing accuracy, and by creating COLL, a large non-equilibrium dataset to stress-test models. It systematically analyzes uncertainty quantification methods, finding ensembles provide better force uncertainty at a cost and that mean-variance estimation struggles for forces. Empirically, DimeNet++ delivers substantial speedups and accuracy gains over DimeNet and outperforms SchNet on COLL, with code and data released for community use. Together, these contributions push ML for molecules beyond equilibrium dynamics toward reliable, uncertainty-aware predictions in reactive systems.

Abstract

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.

Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules

TL;DR

The paper tackles molecular property prediction in non-equilibrium regimes by introducing DimeNet++ to accelerate directional message passing without sacrificing accuracy, and by creating COLL, a large non-equilibrium dataset to stress-test models. It systematically analyzes uncertainty quantification methods, finding ensembles provide better force uncertainty at a cost and that mean-variance estimation struggles for forces. Empirically, DimeNet++ delivers substantial speedups and accuracy gains over DimeNet and outperforms SchNet on COLL, with code and data released for community use. Together, these contributions push ML for molecules beyond equilibrium dynamics toward reliable, uncertainty-aware predictions in reactive systems.

Abstract

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.

Paper Structure

This paper contains 5 sections, 3 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. DimeNet$^{++}$
  3. COLL Dataset
  4. Uncertainty Quantification
  5. Experiments

Figures (4)

  • Figure 1: DimeNet$^{++}$ architecture. $\square$ denotes the layer's input and $\|$ denotes concatenation. Changes to regular DimeNet are highlighted in red.
  • Figure 2: DimeNet's original "directional message passing" block.
  • Figure 3: Example configurations from the COLL dataset. COLL covers a much broader range of configurational space, including stretched bonds and distorted angles.
  • Figure 4: Distribution of atomization energy per atom. COLL covers a much wider range.