Title: GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation
Authors: Minkai Xu
Lantao Yu
Yang Song
Chence Shi
Stefano Ermon
Jian Tang
Venue: ICLR 2022
Keywords: molecular conformation generation
deep generative models
diffusion probabilistic models

Abstract: Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative models. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to a noise distribution, in this paper, we propose a novel generative model named GeoDiff for molecular conformation prediction. GeoDiff treats each atom as a particle and learns to directly reverse the diffusion process (i.e., transforming from a noise distribution to stable conformations) as a Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be roto-translational invariant. We theoretically show that Markov chains evolving with equivariant Markov kernels can induce an invariant distribution by design, and further propose building blocks for the Markov kernels to preserve the desirable equivariance property. The whole framework can be efficiently trained in an end-to-end fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments on multiple benchmarks show that GeoDiff is superior or comparable to existing state-of-the-art approaches, especially on large molecules. 


Review discussion:
## Pros
* GeoDiff is an end-to-end model for generating Cartesian coordinates from a molecular graph.
* GeoDiff shows very good performance on standard test sets.
* GeoDiff outperforms various state-of-the-art methods for generating molecular conformations.
## Cons
* GeoDiff's model seems to involve a large number of parameters.
## Comments and Questions
* What is the computation time for generating a structure with GeoDiff? How does this compare to other ML approaches or RDKit?
* What is the accuracy of the structures at intermediate stages of the diffusion process? Perhaps this could provide some insight into whether $T=5000$ diffusion steps are needed to reach the reported performance.
* Have you also run tests with a smaller number of diffusion steps?
* You show that GeoDiff's model and objective functions are invariant under rigid transformations.
However, what if the training data are not invariant? Doesn't this still fix some arbitrary orientation in the trained model?
* Which threshold $\delta$ did you use in the computation of the precision and recall measures?
* What is the total number of model parameters that are trained?
* How is GeoDiff +FF implemented? What is GeoDiff's performance in property prediction without FF?
## Language
* Page 3: "It has been shown effective"
* Page 9: "can stay the superior diversity"
* Page 9: "yeild" instead of "yield"
NO.
Thank you for your constructive comments and suggestions! The response to your questions are listed below:
**Q1: The computational time for generation.**
- Our generation time is in line with the most competitive existing method ConfGF (which also requires 5000 sampling steps), and requires around 8500 and 11500 seconds to decode the entire QM9 and Drugs test sets respectively. The sampling time varies among different molecules, so here we report the sampling time of the whole test set instead. This sampling time is around 100x times compared to more previous models as well as RDKit.
- We admit that sampling inefficiency is a typical bottleneck for diffusion-based models. And as stated in the conclusion section, we leave it as an important future work about utilizing the most recent progress of diffusion methods [1] to accelerate the sampling process.
Following this concern, we also further test an additional setting with fewer diffusion steps. Please refer to our response to your Q3 for details.
**Q2: About the structures at intermediate stages.**
Thanks for your suggestion and we have evaluated the intermediate structures. The results are as below:
| Number of steps | COV-R(Mean) | COV-R(Median) | MAT-R(Mean) | MAT-R(Median) | COV-P(Mean) | COV-P(Median) | MAT-P(Mean) | MAT-P(Median) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|4000 | 89.33 | 97.62 | 0.8542 | 0.8251 | 61.63 | 64.33 | 1.1709 | 1.1358|
|3000 | 71.34 | 85.39 | 1.1055 | 1.0901 | 33.57 | 25.98 | 1.5256 | 1.4822|
|2000 | 32.40 | 22.49 | 1.3812 | 1.3756| 9.75 | 3.78| 1.8977 | 1.8495|
|1000 | 0.65 | 0.00 | 1.9854 | 1.9965 | 0.09 | 0.00 | 2.4904 | 2.4629|
By comparing with Tab.1, we can observe that running 4000 steps is already good enough and in the last 1000 steps (from 4000 to 5000) the performance only changes slightly. This interesting result indicates 4000 is enough to reach the desired performance, and we leave more detailed studies of the decoding process as future work.
**Q3: Tests with fewer diffusion steps.**
Thanks for your suggestion and we tested the setting with T=1000 on GEOM-Drugs. We set $\beta_1$ as 1e-7 and $\beta_T$ as 9e-3. The results are:
|COV-R(Mean) | COV-R(Median) | MAT-R(Mean) | MAT-R(Median) | COV-P(Mean) | COV-P(Median) | MAT-P(Mean) | MAT-P(Median)|
|--- |--- | --- | --- | --- | --- | --- | --- |
|82.96 | 96.29 | 0.9525 | 0.9334 | 48.27 |  46.03 | 1.3205 | 1.2724|
Compared with Tab.1, we can see T=1000 shows slightly weaker results than T=5000, but already outperforms all existing methods. Note that, the most competitive baseline ConfGF also requires 5000 sampling steps, which indicates that our model can achieve better performance with fewer computational costs compared with the state-of-the-art method.
**Q4: What if the training data are not invariant?**
Since our model and objective functions preserve the invariance of transformations, any arbitrary orientation will have no effect on the trained model.
Therefore, even when the training data is not invariant, our model can still enjoy the invariance property.
**Q5: Which threshold $\delta$ did you use?**
Thanks a lot and we have missed this important information! We follow existing works (CGCF, ConfVAE and ConfGF methods mentioned in the paper) to set $\delta=0.5$ for QM9 and $\delta=1.25$ for drugs. We have added this to the paper.
**Q6: The total number of model parameters.**
Our current configuration holds 795858 parameters, which lies in the typical shape of graph neural networks for molecular geometric systems.
**Q7: How is GeoDiff +FF implemented? What is GeoDiff's performance in property prediction without FF?**
GeoDiff+FF is implemented by first drawing samples from GeoDiff as illustrated in Sec 4.4, and then further optimizing the structures with an empirical force field.
Specifically, we use the API from RDKit that “from rdkit.Chem.rdForceFieldHelpers import MMFFOptimizeMolecule”.
*We hope the above response and the updated draft could address your questions!*
**References**: [1] Nichol, Alex, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." arXiv preprint arXiv:2102.09672 (2021).
**Originality:** The idea of the work is novel - **GeoDiff** is the first generative model for molecular conformation generation based on a diffusion framework.
**Clarity and Quality:** The authors clearly provide motivation and state the differences of the proposed model from other approaches. The architecture is described in an easy-to-follow and intuitive way. The paper provides good **Related Work** and **Results & Discussion** sections.
Still, it would be great if authors add a clear description of contributions and additional figures of architecture and an aligned cloud of conformations to visually evaluate the diversity of conformations.
**Significance:** The contribution of the paper to the field is significant. The experiments show a significant metrics gap between the proposed method and baselines, still, it would be great if authors will add standard deviation of metrics to evaluate the stability of the method.
**Drawbacks / questions**:
1. In 3.1 Authors say, they are interested in generating stable conformations. Authors provide energy metrics for the generated ensemble on GEOM-QM9 split, which covers only 30 molecules.
Still, they don't provide energy metrics for medium size molecules from GEOM-DRUGS dataset.
2.
Despite the proposed model being compared against a variety of baselines and that the authors mentioned the reason why they are not compared with the GeoMol, it requires a more detailed explanation of GeoMol limitations, since the GeoMol is the current state of the art approach in conformation generation task.
3. Diffusion models suffer from high computational complexity. Based on the parameters provided by the authors in appendix B, it's logical to assume that 5000 iterations per molecule are very computationally expensive. The article lacks an estimate of the generation time for one molecule.
NO.
Thank you for your constructive comments and suggestions! The response to your questions are listed below:
**Q1: About generating stable conformations.**
- First of all, one could note that the GEOM dataset provides stable conformations with local minimum energy for all molecules.
Therefore, by saying “generating stable conformations”, we mean our model can fit the distribution of the dataset and generate conformations close to the dataset, which is verified by experiments in Sec 5.2.
- Besides, one should notice that the energy prediction task is not used to test “whether generated conformations are stable with low enough”, but to verify “whether the energy ensemble of generated ones is near to reference ones”.
Therefore, the energy metrics are not directly related to the structure stability.
- **For the dataset setup of energy prediction task**: We just directly follow previous work ConfGF to compare over 30 molecules of GEOM-QM9, which is a standard benchmark. We have also followed your suggestion to test on GEOM-Drugs, but the results turned out to be meaningless with extremely high variance. We think this is because small changes of large molecules can lead to a significant difference of energy ensemble, and thus not suitable for this benchmark.
**Q2: Detailed explanation of GeoMol limitation.**
GeoMol requires that the molecular graph should contain a special substructure of the dihedral pattern.
Specifically, in the official implementation, the pattern is set with open-sourced RDKit tool as “dihedral_pattern = Chem.MolFromSmarts('[\*]\~[\*]\~[\*]\~[\*]')”, which filtered out about one-third of the dataset and made it incomparable with other models.
**Q3: An estimation of the generation time.**
- Our generation time is in line with the most competitive existing method ConfGF (which also requires 5000 sampling steps), and requires around 8500 and 11500 seconds to decode the entire QM9 and Drugs test sets respectively. The sampling time varies among different molecules, so here we report the sampling time of the whole test set instead.
- We admit that sampling inefficiency is a typical bottleneck for diffusion-based models. And as stated in the conclusion section, we leave it as an important future work about utilizing the most recent progress of diffusion methods [1] to accelerate the sampling process.
- Following your concern, we also test an additional setting with T=1000 on GEOM-Drugs. We set $\beta_1$ as 1e-7 and $\beta_T$ as 9e-3. The results are:
|COV-R(Mean) | COV-R(Median) | MAT-R(Mean) | MAT-R(Median) | COV-P(Mean) | COV-P(Median) | MAT-P(Mean) | MAT-P(Median)|
| --- | --- | --- | --- | --- | --- | --- | --- |
| 82.96 | 96.29 | 0.9525 | 0.9334 | 48.27 |  46.03 | 1.3205 | 1.2724 |
- This can further speed up GeoDiff by 5 times. Compared with Tab.1, we can see T=1000 shows slightly weaker results than T=5000, but already outperforms all existing methods. Note that, the most competitive baseline ConfGF also requires 5000 sampling steps, which indicates that our model can achieve better performance with fewer computational costs compared with the state-of-the-art method.
We hope the above response and the updated draft could address your questions!
**Reference**: [1] Nichol, Alex, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." arXiv preprint arXiv:2102.09672 (2021).
The authors clarified all unclear parts, accelerated the sampling process, and provided an estimate of the sampling time. I will raise my score from 6 to 8.
This paper brings together recent ideas and methods (e.g. diffusion, SE(3) equivariance) to the established task of molecular conformation generation with impressive empirical results. A few more experiments are necessary. The paper is well written. I am only tangentially in this area so I am not intimately familiar with the prior works and benchmarks, and may have missed something in the empirical evaluation.
Comments:
* The authors state that roto-translation invariance is a critical inductive bias for this problem.
But I am not convinced that it is critical for this setup and model class. It would be nice to include a baseline of the diffusion model without SE(3) invariance.
* I am curious about the sensitivity of the performance of the model to the hyperparameters of the diffusion process. Can the authors describe how they reached the current settings given in Table 5.
In particular, T=5000 steps seems quite high and I was curious about the performance with lower T say around 1000. As a minor note, there appears to be an inconsistency in the reported batch size (128 in the text vs. 32/64 in Table 5)
* What is the generation time with 5000 diffusion steps? Does that become a bottleneck in conformation generation?
* What happened to the other methods on the geom-drugs dataset for Table 3?
* The authors should mention recent point cloud diffusion papers in their related works as it is addressing a very similar problem
* What are the reverse process sigmas set to?
Minor comments:
* Typo: translatoinal
* Typo: sophisticate
* I noticed a reference that seemed incorrect: Xu et al 2021a as the reference for molecular generation with VAEs in the second paragraph. The authors should double check their references
NO.
Thank you for your constructive comments and suggestions! The response to your comments are listed below:
**Q1: Not convinced about the equivariant setup and it would be nice to include a baseline without SE(3) invariance.**
Thanks for your suggestion, and we have designed an additional non-SE(3) equivariant baseline. We use the same networks architecture to obtain atom embeddings $h$ (as in Eq.6), and then directly output the atomic $\epsilon$ from $h$ by MLPs.
Therefore, $\epsilon$ is non-equivariant w.r.t $C$. The numerical results on the challenging GEOM-Drugs dataset are as below:
|COV-R(Mean) | COV-R(Median) | MAT-R(Mean) | MAT-R(Median) | COV-P(Mean) | COV-P(Median) | MAT-P(Mean) | MAT-P(Median)|
| --- | --- | --- | --- | --- | --- | --- | --- |
|0.0 | 0.0 | 54.8588 | 55.0188 | 0.0 | 0.0 | 634.4928 | 625.7850 |
As expected, the non-equivariant model shows extremely poor performance, which is in line with the observation in previous works [1].
**Q2: How to select the hyper-parameters. How about lower diffusion steps T like 1000.**
- The majority of hyper-parameters ($T$, $\tau$, batch size, and training iters) are selected by following the previous work ConfGF. For the beta scheduler, our strategy is to ensure the variances from 0 to T go exponentially from around 0.01 to 10, which has also shown effectiveness in ConfGF.
- The most competitive baseline ConfGF also uses 5000 sampling steps, where they set 50 noise scales and conduct 100 optimization steps for each noise scale.
Therefore we also choose T=5000. We have followed your suggestion to test the setting with T=1000 on GEOM-Drugs. We set $\beta_1$ as 1e-7 and $\beta_T$ as 9e-3. The results are:
|COV-R(Mean) | COV-R(Median) | MAT-R(Mean) | MAT-R(Median) | COV-P(Mean) | COV-P(Median) | MAT-P(Mean) | MAT-P(Median)|
| --- | --- | --- | --- | --- | --- | --- | --- |
| 82.96 | 96.29 | 0.9525 | 0.9334 | 48.27 |  46.03 | 1.3205 | 1.2724 |
- Compared with Tab.1, we can see when T=1000, though slightly weaker than T=5000, it already outperforms all existing methods. Note that, the most competitive baseline ConfGF also requires 5000 sampling steps, which indicates that our model can achieve better performance with fewer computational costs compared with the state-of-the-art method.
**Q3: About the generation time.**
Our generation time is in line with the most competitive existing method ConfGF, and it requires around 8500 and 11500 seconds to decode QM9 and Drugs test sets respectively. This is a typical bottleneck for diffusion-based models. And as stated in the conclusion section, we leave it as an important future work about utilizing the most recent progress of diffusion methods [2] to accelerate the sampling process.
**Q4: About other methods on Tab.4.**
This table mainly aims to compare GeoDiff with RDKit, and thus we omit other models. We have additionally tested the performance of several other methods. The results are as follows:
|Method | COV-R(Mean) | COV-R(Median) | MAT-R(Mean) | MAT-R(Median) | COV-P(Mean) | COV-P(Median) | MAT-P(Mean) | MAT-P(Median)|
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|GraphDG+FF | 90.71 | 100.00 | 0.8839 | 0.8514 | 57.21 | 60.81 | 1.4549 | 1.4268|
|CGCF+FF | 90.56 | 100.00 | 0.7827 | 0.7573 | 74.06 | 81.88 | 1.1736 | 1.1230|
|ConfVAE+FF | 91.43 | 100.00 | 0.7749 | 0.7575 | 75.74 | 83.32 | 1.1391 | 1.1007|
|ConfGF+FF | 92.54 | 100.00 | 0.7696 | 0.7528 | 76.34 | 85.86 | 1.1374 | 1.0886|
Compared with the results in Tab.3,  we can see that when further optimized with FF, some baselines can achieve similar performance considering the Recall (-R) metrics, but still behind GeoDiff with a large margin for the Precision (-P) metrics.
**Q5: Mention related work for the point cloud generation.**
Thanks for the suggestion. We agree that the problem is related and have added them now. We would highlight that for existing point cloud generative models,  typically the data is not in graph structures with atom and bond information, and the equivariance is not considered, making them fundamentally different from our method.
**Q6: What are sigmas?**
In this work we set $\sigma_t = 1 - \frac{\alpha_t}{\alpha_{t-1}}$.
We hope the above response and the updated draft could address your questions!
**References:**
- [1] Mansimov, Elman, et al. "Molecular geometry prediction using a deep generative graph neural network." Scientific reports 9.1 (2019): 1-13.
- [2] Nichol, Alex, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." arXiv preprint arXiv:2102.09672 (2021).
We would like to thank all the reviewers and public readers for your constructive feedback. We’ve updated the paper according to your suggestions. More specifically, we have made the following changes:
- We add several additional contents to the paper, including extra experiments and related works. We have also added several additional experiment details. These contents are all marked in blue color. Due to the page limits, most contents are added to the Appendix part.
- We have updated several typos.
*We hope the above response and the updated draft could address your questions!*

Meta-review:
role: PC
data: comment: The authors focus on the conditional generation of molecular conformations (i.e. 3D cartesian atom positions) from a given molecular graph. They formulate the generation via diffusion probabilistic models.  Conformations are generated by a reverse diffusion process from isotropic Gaussian noise to molecular conformations. This diffusion process is learned from data using a SE(3) invariant formulation of the diffusion process. The authors work directly with atomic positions (i.e. a point cloud) instead of interatomic distances or an intermediate bond geometry representation. Experimental evaluations show state-of-the-art results according to COV/MAT metrics on GEOM-Drugs and GEOM-QM9 datasets.

Strengths:

- High technical novelty: first generative model for molecular conformation generation based on a diffusion framework
- Very clearly written paper.
- Impressive empirical results with state-of-the-art results on GEOM-Drugs and GEOM-QM9 datasets.

Weaknesses:

- Most of the weaknesses reported by the reviewers seem to have been addressed in the rebuttal.

The idea of the work is highly novel. The authors propose the first generative model for molecular conformation generation based on a diffusion framework. This paper brings together recent ideas and methods (e.g. diffusion, SE(3) equivariance) to the established task of molecular conformation generation with impressive empirical results. All the reviewers agree on acceptance with high scores. I recommend the authors to look at the reviewers' comments to improve the paper for the camera-ready version

Decision: Accept (Oral)