Structure-Aware E(3)-Invariant Molecular Conformer Aggregation Networks

TL;DR

The paper introduces ConAN, an $E(3)$-invariant framework that fuses 2D molecular graphs with ensembles of 3D conformers via a differentiable Fused Gromov-Wasserstein (FGW) barycenter. It combines a 2D MPNN, a 3D SchNet-based conformer encoder, and a FGW-barycenter aggregator to produce a unified, permutation- and rotation-invariant representation for downstream property prediction. The authors prove invariance properties, establish a fast $\mathcal{O}(1/K)$ convergence rate for the empirical FGW barycenter, and develop an entropic Sinkhorn-based solver for scalable training on GPUs. Empirically, ConAN-FGW achieves state-of-the-art or competitive results across MoleculeNet regression tasks, SARS-CoV-2 classification, and conformer-ensemble benchmarks, with analysis showing that a small number of conformers ($K\approx 5$) can suffice. The work also demonstrates substantial efficiency gains over prior FGW approaches and highlights the practical impact of geometry-aware conformer aggregation for molecular modeling.

Abstract

A molecule's 2D representation consists of its atoms, their attributes, and the molecule's covalent bonds. A 3D (geometric) representation of a molecule is called a conformer and consists of its atom types and Cartesian coordinates. Every conformer has a potential energy, and the lower this energy, the more likely it occurs in nature. Most existing machine learning methods for molecular property prediction consider either 2D molecular graphs or 3D conformer structure representations in isolation. Inspired by recent work on using ensembles of conformers in conjunction with 2D graph representations, we propose $\mathrm{E}$(3)-invariant molecular conformer aggregation networks. The method integrates a molecule's 2D representation with that of multiple of its conformers. Contrary to prior work, we propose a novel 2D-3D aggregation mechanism based on a differentiable solver for the Fused Gromov-Wasserstein Barycenter problem and the use of an efficient conformer generation method based on distance geometry. We show that the proposed aggregation mechanism is $\mathrm{E}$(3) invariant and propose an efficient GPU implementation. Moreover, we demonstrate that the aggregation mechanism helps to significantly outperform state-of-the-art molecule property prediction methods on established datasets.

Figures (7)

• Figure 1: Overview of the proposed conformer aggregation network with alanine dipeptide as example input.
• Figure 2: Illustration of the feature-based and structural distances of conformers (here: alanine dipeptide) we use for the computation of the Fused Gromov-Wasserstein barycenter.
• Figure 3: Ablation study on the effect of number conformers on the FGW barycenter component on valid sets.
• Figure 4: (left) box-plot distribution of mean, variance, maximum, and minimum distances among conformers; (right) distribution of the same values where sample top-$k$ closest conformers.
• Figure 5: Comparing runtimes of FGW-Mixup, ConAN-FGW (single and multi-GPU).

Theorems & Definitions (6)

• Theorem 3.1
• Theorem 4.1
• Lemma 2.1
• proof : Proof of \ref{['lemma_eq_upper_bound']}
• Definition 2.2: Strongly convex and smooth functions
• Lemma 2.3: Corollary 4.4 from le_gouic_fast_2022