Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

WGFormer: An SE(3)-Transformer Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction

Fanmeng Wang, Minjie Cheng, Hongteng Xu

TL;DR

This work tackles the challenge of predicting molecular ground-state conformations with high accuracy and efficiency. It introduces WGFormer, a Wasserstein gradient flow-driven SE(3)-Transformer that operates in an auto-encoding framework, encoding low-quality conformations and decoding ground-state coordinates via an MLP, while enforcing latent energy minimization on atom mixtures. The approach connects transformer-based attention to Wasserstein gradient flows, leveraging energy functionals E^0 and E^∞ and entropic OT to guide conformation optimization. Empirically, WGFormer achieves state-of-the-art results on Molecule3D and QM9 with notable speedups, demonstrating the value of a physically grounded, interpretable framework for conformation prediction and downstream applications like docking and property prediction. The work suggests a promising direction for energy-aware, geometry-preserving molecular modeling and invites further theoretical development and broader task extensions.

Abstract

Predicting molecular ground-state conformation (i.e., energy-minimized conformation) is crucial for many chemical applications such as molecular docking and property prediction. Classic energy-based simulation is time-consuming when solving this problem, while existing learning-based methods have advantages in computational efficiency but sacrifice accuracy and interpretability. In this work, we propose a novel and effective method to bridge the energy-based simulation and the learning-based strategy, which designs and learns a Wasserstein gradient flow-driven SE(3)-Transformer, called WGFormer, for ground-state conformation prediction. Specifically, our method tackles this task within an auto-encoding framework, which encodes low-quality conformations by the proposed WGFormer and decodes corresponding ground-state conformations by an MLP. The architecture of WGFormer corresponds to Wasserstein gradient flows -- it optimizes conformations by minimizing an energy function defined on the latent mixture models of atoms, thereby significantly improving performance and interpretability. Extensive experiments demonstrate that our method consistently outperforms state-of-the-art competitors, providing a new and insightful paradigm to predict ground-state conformation.

WGFormer: An SE(3)-Transformer Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction

TL;DR

This work tackles the challenge of predicting molecular ground-state conformations with high accuracy and efficiency. It introduces WGFormer, a Wasserstein gradient flow-driven SE(3)-Transformer that operates in an auto-encoding framework, encoding low-quality conformations and decoding ground-state coordinates via an MLP, while enforcing latent energy minimization on atom mixtures. The approach connects transformer-based attention to Wasserstein gradient flows, leveraging energy functionals E^0 and E^∞ and entropic OT to guide conformation optimization. Empirically, WGFormer achieves state-of-the-art results on Molecule3D and QM9 with notable speedups, demonstrating the value of a physically grounded, interpretable framework for conformation prediction and downstream applications like docking and property prediction. The work suggests a promising direction for energy-aware, geometry-preserving molecular modeling and invites further theoretical development and broader task extensions.

Abstract

Predicting molecular ground-state conformation (i.e., energy-minimized conformation) is crucial for many chemical applications such as molecular docking and property prediction. Classic energy-based simulation is time-consuming when solving this problem, while existing learning-based methods have advantages in computational efficiency but sacrifice accuracy and interpretability. In this work, we propose a novel and effective method to bridge the energy-based simulation and the learning-based strategy, which designs and learns a Wasserstein gradient flow-driven SE(3)-Transformer, called WGFormer, for ground-state conformation prediction. Specifically, our method tackles this task within an auto-encoding framework, which encodes low-quality conformations by the proposed WGFormer and decodes corresponding ground-state conformations by an MLP. The architecture of WGFormer corresponds to Wasserstein gradient flows -- it optimizes conformations by minimizing an energy function defined on the latent mixture models of atoms, thereby significantly improving performance and interpretability. Extensive experiments demonstrate that our method consistently outperforms state-of-the-art competitors, providing a new and insightful paradigm to predict ground-state conformation.

Paper Structure

This paper contains 28 sections, 2 theorems, 44 equations, 9 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Molecular Conformation Generation
  4. Applications of Ground-State Conformation
  5. Proposed Model
  6. An Auto-Encoding Framework for Ground-State Conformation Prediction
  7. The Architecture of WGFormer
  8. Revisiting WGFormer through the Lens of Wasserstein Gradient Flows
  9. Continuous Counterpart of Attention Head
  10. Wasserstein Gradient Flows When $M=0$ or $\rightarrow\infty$
  11. Learning Algorithm
  12. Experiments
  13. Experimental Setup
  14. Comparisons
  15. Ablation Studies and Analyses
  16. ...and 13 more sections

Key Result

Proposition 4.1

For $\mu=\sum_{i=1}^{N}\delta_{\bm{x}_i}$, the $T_{\mu}^M$ in Eq. eq:operator is the Wasserstein gradient flow, i.e., $T_{\mu}^{M}=-\nabla_W E^M(\mu)$, when $M=0$ or $M\rightarrow\infty$. The corresponding energy functions are where $\Pi_{\mu}=\{\pi\in\mathcal{M}(\mathbb{R}^{D}\times\mathbb{R}^D)|\int_{\bm{x}}\mathrm{d}\pi(\bm{x},\bm{x}')=\mu(\bm{x}'),\int_{\bm{x}'}\mathrm{d}\pi(\bm{x},\bm{x}')=\

Figures (9)

  • Figure 1: An illustration of the proposed model architecture, in which WGFormer corresponds to Wasserstein gradient flows and minimizes a physically reasonable energy function defined on the latent mixture models of atoms.
  • Figure 2: An illustration of the $h$-th head in the $l$-th WGFormer.
  • Figure 3: The visual comparison for the ground-state conformations predicted by our WGFormer and the top-3 baselines.
  • Figure 4: The efficiency comparison for our WGFormer and the top-3 baselines on each dataset.
  • Figure 5: Some cases about the change of latent and potential energy as the number of layers increases, where latent energy is obtained by solving Eq. \ref{['eq:eot']} and potential energy is obtained by xTB bannwarth2019gfn2. More cases are provided in Figure \ref{['fig: energy_more']}.
  • ...and 4 more figures

Theorems & Definitions (4)

  • Proposition 4.1
  • Proposition 4.2
  • proof
  • proof