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

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.

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.