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

arXiv: https://arxiv.org/abs/2410.09795 · HTML: https://mlsanity.com/paper/2410.09795/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 51 references that are in the map.
- First posted on arXiv 13 Oct 2024.

## Builds on (most cited first, 25 of 25)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [How Attentive are Graph Attention Networks?](https://mlsanity.com/paper/2105.14491/) (2021)
- [Recipe for a General, Powerful, Scalable Graph Transformer](https://mlsanity.com/paper/2205.12454/) (2022)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [Learning Gradient Fields for Molecular Conformation Generation](https://mlsanity.com/paper/2105.03902/) (2021)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Molecular geometry prediction using a deep generative graph neural network](https://mlsanity.com/paper/1904.00314/) (2019)
- [Learning Neural Generative Dynamics for Molecular Conformation Generation](https://mlsanity.com/paper/2102.10240/) (2021)
- [A Mean-Field Optimal Control Formulation of Deep Learning](https://mlsanity.com/paper/1807.01083/) (2018)
- [An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming](https://mlsanity.com/paper/2105.07246/) (2021)
- [Redesigning the Transformer Architecture with Insights from Multi-particle Dynamical Systems](https://mlsanity.com/paper/2109.15142/) (2021)
- [Molecule3D: A Benchmark for Predicting 3D Geometries from Molecular Graphs](https://mlsanity.com/paper/2110.01717/) (2021)
- [A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications](https://mlsanity.com/paper/2403.00485/) (2024)
- [UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning](https://mlsanity.com/paper/2405.10343/) (2024)
- [MMPolymer: A Multimodal Multitask Pretraining Framework for Polymer Property Prediction](https://mlsanity.com/paper/2406.04727/) (2024)
- [3D Graph Contrastive Learning for Molecular Property Prediction](https://mlsanity.com/paper/2208.06360/) (2022)
- [Gradual Optimization Learning for Conformational Energy Minimization](https://mlsanity.com/paper/2311.06295/) (2023)

## Built on by (most cited first, 4 of 4)

- [Universally Invariant Learning in Equivariant GNNs](https://mlsanity.com/paper/2510.13169/) (2025)
- [Geometric Mixture Models for Electrolyte Conductivity Prediction](https://mlsanity.com/paper/2510.15403/) (2025)
- [Energy-Guided Generative Modeling for Low-Energy Molecular Structure Discovery](https://mlsanity.com/paper/2512.22597/) (2025)
- [MuCO: Generative Peptide Cyclization Empowered by Multi-stage Conformation Optimization](https://mlsanity.com/paper/2602.11189/) (2026)

## Abstract (arXiv)

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.