🚀 Introduction

Our WGFormer achieves state-of-the-art performance in molecular ground-state conformation predictioin. In particular, our WGFormer 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 the model performance and interpretability.

⚒️ Environment

We have provided an env.yml file that contains the necessary environment dependencies. To set up your environment, please execute:

conda env create -f env.yml
conda activate WGFormer

📦 Datasets

We use the Molecule3D and QM9 provided in GTMGC as our datasets. Please download them and organize the ./data directory as follows:

data
├── Molecule3D
│   ├── Random
│   │   ├── dict.txt
│   │   ├── random_train.csv
│   │   ├── random_train_0.sdf
│   │   ├── random_train_1.sdf
│   │   ├── random_train_2.sdf
│   │   ├── random_train_3.sdf
│   │   ├── random_valid.csv
│   │   ├── random_valid.sdf
│   │   ├── random_test.csv
│   │   └── random_test.sdf
│   ├── Scaffold
│   │   ├── dict.txt
│   │   ├── scaffold_train.csv
│   │   ├── scaffold_train_0.sdf
│   │   ├── scaffold_train_1.sdf
│   │   ├── scaffold_train_2.sdf
│   │   ├── scaffold_train_3.sdf
│   │   ├── scaffold_valid.csv
│   │   ├── scaffold_valid.sdf
│   │   ├── scaffold_test.csv
│   │   └── scaffold_test.sdf
├── QM9
│   ├── dict.txt
│   ├── gdb9.sdf
│   ├── train_indices.csv
│   ├── valid_indices.csv
│   └── test_indices.csv

Then, execute the command below to obtain the processed datasets:

cd data
python process_QM9.py
python process_Molecule3D.py

Here, the processed QM9 dataset has been provided in this link.

💪 Experiments

Training

Our model weights have been provided in this link. You can also train from scratch by running the following script.

bash train.sh

Inference

bash inference.sh valid
bash inference.sh test

Evaluation

python eval.py valid
python eval.py test

👍 Acknowledgments

This code is built upon Uni-Mol and Uni-Core. Thanks for their contribution.

📌 Citation

If you find this work useful for your research, please consider citing it. 😊

@inproceedings{wang2025wgformer,
  title={WGFormer: An SE (3)-Transformer Driven by Wasserstein Gradient Flows for Molecular Ground-State Conformation Prediction},
  author={Wang, Fanmeng and Cheng, Minjie and Xu, Hongteng},
  booktitle={International Conference on Machine Learning},
  pages={62652--62671},
  year={2025},
  organization={PMLR}
}