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 introduces a Wasserstein gradient flow-driven SE(3)-Transformer to achieve both accurate and efficient molecular ground-state conformation prediction.

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.
The code is available at https://github.com/FanmengWang/WGFormer.

Lay Summary:

How can the molecular ground-state conformation, which represents the most stable 3D structure corresponding to the energy-minimized state on the potential energy surface, be accurately and efficiently obtained?
In this work, we propose WGFormer, a Wasserstein gradient flow-driven SE(3)-Transformer, to tackle this task.
Specifically, it can be interpreted as Wasserstein gradient flows, which optimizes molecular conformation by minimizing a physically reasonable energy function defined on the latent mixture models of atoms, thereby significantly improving performance and interpretability. Extensive experiments and analyses comprehensively validate its rationality and superiority, thus providing a new and insightful paradigm for molecular ground-state conformation prediction.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/FanmengWang/WGFormer

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: Molecular ground-state conformation, Molecular modeling, Wasserstein gradient flow

Submission Number: 5077