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

TL;DR

WGFormer introduces a Wasserstein gradient flow-driven SE(3)-Transformer that efficiently predicts molecular ground-state conformations, outperforming existing methods in accuracy and interpretability by integrating energy minimization within an auto-encoding framework.

Contribution

The paper presents a novel WGFormer model that combines Wasserstein gradient flows with SE(3)-Transformers for improved molecular conformation prediction, bridging energy-based and learning-based approaches.

Findings

Outperforms state-of-the-art methods in conformation prediction accuracy.

Enhances interpretability through energy minimization on latent atom models.

Demonstrates consistent performance improvements across experiments.

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.