Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization

This repository is the official implementation of Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization (NeuIPS 2024). [PDF]

Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization
Siyi Gu, Minkai Xu, Alexander Powers, Weili Nie, Tomas Geffner, Karsten Kreis, Jure Leskovec, Arash Vahdat, Stefano Ermon
Stanford University, NVIDIA

Environment

conda create -n target python=3.8
conda activate target
conda install pytorch==2.0.1 pytorch-cuda=11.7 -c pytorch -c nvidia
conda install pyg -c pyg
pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
conda install rdkit openbabel tensorboard pyyaml easydict python-lmdb -c conda-forge

# For Vina Docking
pip install meeko==0.1.dev3 scipy pdb2pqr vina==1.2.2 
python -m pip install git+https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3

Training from scratch

Data Generation

SFT Data

The data preparation follows (https://github.com/guanjq/targetdiff).

Preference Data

python gen_data.py

Train

python scripts/train_ipo.py configs/training_ipo.yml

Sampling

python scripts/sample_diffusion.py configs/sampling.yml --data_id {i} # Replace {i} with the index of the data. i should be between 0 and 99 for the testset.

Sampling Results

https://drive.google.com/drive/folders/1Auvigp6FLgNKY0i8eVLQf5loFwrIdW0G?usp=sharing

Evaluation

Evaluation from sampling results

python scripts/evaluate_diffusion.py {OUTPUT_DIR} --docking_mode vina_score --protein_root data/test_set

The docking mode can be chosen from {qvina, vina_score, vina_dock, none}

Evaluation results

https://drive.google.com/drive/folders/1eRCcALnBpuVgjUqqRucpZSTtF6oT9pX3?usp=sharing

Citation

@inproceedings{gu2024aligning,
  title={Aligning Target-Aware Molecule Diffusion Models with Exact Energy Optimization},
  author={Gu, Siyi and Xu, Minkai and Powers, Alexander and Nie, Weili and Geffner, Tomas and Kreis, Karsten and Leskovec, Jure and Vahdat, Arash and Ermon, Stefano},
  booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
  year={2024},
  url={https://openreview.net/forum?id=EWcvxXtzNu}
}

MinkaiXu / AliDiff

readme