EQGAT-diff: Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation

Tuan Le^1* ,  Julian Cremer^1* ,  Djork-Arné Clevert¹ ,  Kristof T. Schütt¹ , 

^*Equal contribution    ¹Pfizer Research & Development   

ICLR, 2024

[Paper]      

This is the official repository for EQGAT-diff - a model for 3D molecule generation via equivariant (continuous and discrete) denoising diffusion. The code is inside eqgat_diff and is work in progress. If you have any questions, feel free to reach out to us: julian.cremer@pfizer.com, tuan.le@pfizer.com.

Find the inference subdirectory: ./eqgat-diff/tree/main/inference Find the link to download the model weights for the QM9 and Geom-Drugs models here: https://drive.google.com/drive/folders/1_qqhKRU4GI43uSMFB0cMgreLxubZGYWZ?usp=share_link

For QM9, the provided checkpoint is the model after 300 epochs of training of EQGAT^{x0}{disc} as in Table 2, while for GEOM-Drugs, I have attached the checkpoint to the model EQGAT^{x0}{disc} that was further trained as described in Table 3.

Please place the model weights into the subfolders ./eqgat-diff/tree/main/weights, like so:

weights: geom: best_mol_stab.ckpt qm9: best_mol_stab.ckpt

We provide two notebooks to play around with the model (e.g. ./inference/sampling_geom.ipynb) Make sure to extract the two data folders here: ./data

Below, we show some sampling trajectories from our models. The corresponding .gif files can be found in the `assets/` directory.

Unconditional Sampling Trajectory

Conditional Sampling Trajectory on maximized polarizability

Conditional Sampling Trajectory on minimized polarizability

Conditional Sampling Trajectory on a fixed protein pocket

Acknowledgement

This study was partially funded by the European Union’s Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie Actions grant agreement “Advanced machine learning for Innovative Drug Discovery (AIDD)” No. 956832.