Paper: Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation (ICLR-2024)

Reviewer comment (Novelty):
The de novo generation of 3D molecular structures, as currently presented, appears to lack inherent value.

Author response:
Our study extensively examines various model design choices for 3D molecular data. Initially, we focused on unconditional 3D molecule generation to gain relevant insights into essential design decisions, such as tailoring fast and powerful equivariant deep learning models, incorporating categorical diffusion on discrete data modalities, and a task-specific loss weighting function. Our findings show new state-of-the-art results by a large margin while being significantly faster in training and inference (Appendix Fig. 9 and Fig. 10). Thus, this work helps build the foundation for further investigations in generative chemistry. 

A critical aspect of it is, e.g., structure-based conditional molecule generation. Here, sampling molecules with specific quantum- and physicochemical properties in the presence of known protein receptors necessitates accurate, reliable, and fast 3D molecular modeling.

For conditional molecule design, we can use a trained unconditional EQGAT-diff model together with classifier-guidance, as proposed in [1], to steer the generation of samples using the gradient of an external classifier/regressor during the reverse sampling trajectory from noise to data. As a proof of concept, we explored classifier-guidance to generate molecules optimizing for low/high polarizability showing promising results.

To further highlight the importance of our findings, we provide results on structure-based de novo ligand generation. Here, we evaluated EQGAT-diff on the Crossdocked dataset [2] and found that SNR-t weighting and finetuning are crucial components confirming our previous results.

Reviewer comment (Novelty):
The work does not introduce new technical advancements.

Author response:
In this work, we provide a profound contribution by rigorous experimental design containing many important ablations. 
We think that in science, experimental validation and exploration are at least as important as technical novelty. Especially in ML, we find that proper evaluation by in-depth experimentation and ablation is often harmed when focusing on delivering theoretical or architectural novelty. This work's main idea and driver was the belief that current state-of-the-art 3D diffusion models were not sufficiently examined, resulting in significant flaws, e.g., in the design of larger molecular structures. Thus, we focus on a concise experimental investigation of well-known theoretical findings. We show that by exploring the diffusion design space, it is possible to vastly increase the model's generative performance - e.g., compared to the latest SOTA[3], our model increases the validity on sampled molecules on average by roughly 25\%; we increase the molecular stability (3D metric) by roughly 240\% compared to EDM [4].

Reviewer comment (Presentation):
The referencing is inconsistent, with many citations missing journal or conference details. There are incorrect uses of quotation marks, and figure clarity is compromised, diminishing the manuscript's overall professionalism.

Author response:
We checked the references carefully, added missing journal and conference details, and addressed the concerns raised about incorrect uses of quotation marks. We updated the figures to improve the overall readability of the manuscript.

Reviewer comment (Experiments):
It would be beneficial for the authors to compare their 3D molecule generation model with a 2D molecule generation, specifically using OpenEye’s OMEGA, which is cited as the ground truth for conformation generation. Would the EQGAT-diff model maintain its efficiency advantage in such a comparison?

Author response:
EQGAT-diff is designed and trained to perform de novo molecule generation. Our model is not conditioned on a molecular topology and composition and, thus, is not intended to sample conformers.
Despite this, to our knowledge, there is no SMILES strings or graph-based generative model to which we could compare our results on GEOM-Drugs. 
The task conditional sampling of molecules, where the condition is inherently connected to the 3D structure, e.g., a quantum property or a binding affinity of a ligand to a receptor, would not be feasible with the 2D graph and conformer sampling approach. With our work, we want to lay a solid foundation of selected design components to tackle molecule generation in  3D space, with the further goal to apply such in the structured-based conditional setting. 

In light of the mentioned improvements, we ask the reviewer to consider raising the score.

Reviewer comment (Experiments):
The work proposes to navigate through the design space of equivariant diffusion model for de novo 3D molecule design. However, there are still other design choices that are omitted, e.g., different diffusion kernel and backbones. I know it's infeasible to exhaust the whole design space in one paper. I want to put a note here about the possible extensions.

Author response:
Indeed, the design space of (equivariant) diffusion models is vast. Still, we have covered many choices through the experimental ablations while using one backbone neural network architecture.

Keeping the same network throughout the experiments is crucial for better comparison,  e.g., applying the $\epsilon$-parameterization vs. $x_0$-parameterization for Gaussian diffusion.
Regarding the choice of diffusion kernels, for the discrete diffusion, there is an alternative way to include an additional mask/absorbing state token that can be regarded as an augmented feature for a node. Having such a token in the discrete diffusion kernel might enable practitioners to model the deletion of a node. For example, assuming that we sample $N$ nodes by design, all $N$ nodes need to be populated and filled with one atom class. Similar to mask text in language modeling with an absorbing state feature/token, the model could perform the reverse sampling trajectory and potentially assign nodes in the point cloud of the absorbing state feature. We did not explore this in our work. We will mention possible extensions in our updated manuscript.

Reviewer comment (Experiments):
Will different settings (like discrete vs continuous) lead to significantly different training time till convergence as well as inference time?

Author response:
We evaluated our models during training every 20 epochs. Here, we noticed that the model with $x_0$-parameterization EQGAT-disc-x0 and SNR(t) loss weighting converged significantly faster than its continuous counterparts with and without loss weighting. The inference times for both parameterizations are very similar to each other. However, we also noticed that when using EQGAT-disc-x0 with loss weighting, we can reduce the number of diffusion timesteps to 100 without sacrificing much molecule stability and validity, respectively. Surprisingly, we find that the performance of EQGAT-disc-x0 without loss weighting trained on 500 diffusion steps is worse than with loss weighting on 100. Thus, we see a five-fold decrease in inference time when using the proposed framework (Appendix Fig. 10).

Reviewer comment (Experiments):
Torsional diffusion (https://arxiv.org/abs/2206.01729) introduces a GEOM-XL dataset including large molecules (>100 atoms). It may also be interesting to see how the proposed method transfer to such OOD datasets.

Author response:
The Torsional Diffusion model is a 3D conformer generator that gradually changes torsion angles from a given input 3D structure (e.g., obtained through the RDKit). Hence, the molecule and its atom composition and bond topology do not change. In our scenario of de novo, unsupervised generative modeling, we do not see a direct way to evaluate the GEOM-XL dataset. However, as shown in Figure 3 of our manuscript, our model can sample larger molecules comprising 80-100 atoms but with some decreased performance regarding molecule stability. This is expected, as the diffusion model, which learns the distribution of atoms in a point cloud, was trained on a dataset with an average of around 45 atoms (see Figure 5 in the Appendix).

We thank the reviewer for the feedback. 

A critical aspect of drug discovery is structure-based conditional molecule generation. Here the task is to sample molecules with specific quantum- and physicochemical properties in the presence of known protein receptors.
To further highlight the importance of our findings, we provide results on structure-based de novo ligand generation. We evaluated EQGAT-diff on the Crossdocked dataset \cite{crossdocked} and found that SNR-t weighting and finetuning are crucial components confirming our previous results.

We believe that the revision of the manuscript significantly improves the content and quality of our paper.
Thus, we ask the reviewer to consider raising the score.

Reviewer comment (Presentation):
The fonts in all the figure labels are unreadably (unpublishably?) small. If space is the main concern, then there certainly exist small chunks of text (e.g. the last paragraph of section 4) and some figure panels (e.g. most from Fig 2) that could move to the supplement, and some parts that could be rewritten concisely (e.g. intro to sec 5, which could also be clarified as it was not clear what "these three aspects" meant in paragraph 2, top of page 5).

Author response:
We updated all our figures with increased font size and line widths for improved readability and moved details from the end of Section 4 into the Appendix as suggested by the reviewer. Furthermore, we clarified the introduction of Section 5 as marked in blue color.

Reviewer comment (Experiments):
The dataset is a weakness for this paper and most other papers in this area. The reviewer disagrees with the premise of the authors that the availability of molecular data is not as abundant, though agrees that the community has been stuck on trivial and no longer meaningful datasets. The authors mention the catalogue of Enamine Real in the supplement which scales to the 10s of billions of molecular graphs (albeit limited in complexity), but there are also computational enumerations of valid molecules that can be much bigger. The published patent literature probably scales to the multiple 10s of millions. QM methods of relatively reasonable accuracy exist that are not exceedingly expensive and datasets of 10s of millions of published 3D structures with simpler methods also exist.

Author response:
We disagree on the weakness of GEOM-Drugs, as this dataset covers an interesting area of the chemical space using a well-established semi-empirical level of theory (xTB-GFN2) and has previously been hard to learn (EDM and similar models failed to learn the data distribution).
In principle, there is a large amount of publicly available SMILES strings and molecular graphs. But, for structure-based tasks or tasks that involve experimental measurements, the data is by far not as abundant. This is precisely why we show that pre-training on large databases for 3D conformers like PubChem3D can help with that significantly, even though they have been generated only with empirical force fields. 
We show that for pre-training, the level of theory and the presence of explicit hydrogens is not decisive. Specifically, as shown in Fig. 2 in the main text and Fig. 8 in the Appendix, bond distance distributions do not match between GEOM-Drugs and PubChem3D. Although explicit hydrogens are also missing, fine-tuning is very effective even on small subsets of the target distribution with faster convergence.

Reviewer comment (Theory):
The paper uses E3 symmetries which includes the reflection group and can therefore mix chiralities, in principle changing the input distribution of molecular graphs. The authors seem to not address the possible distribution shifts due to chiral transformations.

Author response:
Our model does not consider chirality for single molecules since the unconditional distribution is symmetric under reflection. The reflection symmetry is only broken for use cases where an external reference system, e.g., a target protein, is introduced. When this is encoded in the model condition, for example, when the model is to be trained to generate a 3d ligand into a 3d pocket, the chirality of the molecule is correctly reflected in the conditional distribution.
If one aims to represent chirality directly, we believe implementing a chirality-sensitive function on torsion angles would be possible, as described in [2]. For diffusion-based de-novo design, however, it is not straightforward to enumerate the corresponding torsion angles since the molecular graph changes over time.

Reviewer comment (Experiments):
Small molecules with 3D coordinates often have specific chirality. Although this chirality is specified in the full set of 3D positions, the authors' model is using E3 symmetries in its core, so there is a concern it will mix chiral molecules in ways that would result in a distribution shift. The reviewer asks whether the authors have checked the chirality of the generated molecular distributions by keeping only one kind of stereo isomer and validating whether the model would construct both. The reviewer also asks whether implementing the SE3 group would allow the authors to include chiral flags on the bonds, similar to the typical 2D representations of molecular graphs.

Author response:
Our model does not consider chirality for single molecules since the unconditional distribution is symmetric under reflection. The reflection symmetry is only broken for use cases where an external reference system, e.g., a target protein, is introduced. When this is encoded in the model condition, for example, when the model is to be trained to generate a 3d ligand into a 3d pocket, the chirality of the molecule is correctly reflected in the conditional distribution.
If one aims to represent chirality directly, we believe implementing a chirality-sensitive function on torsion angles would be possible, as described in [2]. For diffusion-based de-novo design, however, it is not straightforward to enumerate the corresponding torsion angles since the molecular graph changes over time.

Reviewer comment (Theory):
It was not clear why the inclusion of the eps-parameterization does not use a composite loss as did the x0-parameterization. The reviewer questions whether it is not a priori clear that the addition of a categorical loss for categorical data will outperform a mean squared error on categorical data, and suggests that perhaps the authors can explain this and possibly reformulate the early parts of section 5.

Author response:
The diffusion model aims in minimizing the KL-divergence $D_\text{KL}(q(x_{t-1}|x_t, x_0) || p_\theta(x_{t-1}|x_t) )$, while for continuous data $x$, the tractable reverse distribution $q(x_{t-1}|x_t, x_0)$ is a multivariate normal distribution [Ho et al. 2020], and for discrete data $x$, a categorical distribution [Hoogeboom et al. 2021, Austin et al. 2021]. 
The best choice to implement $p_\theta$ is to predict the corresponding distribution parameters of $q(x_{t-1}|x_t, x_0)$, which, for the continuous case, leads to predicting the mean of $q(x_{t-1}|x_t, x_0)$ or equivalently the noise $\epsilon$ that was used when perturbing $x_0$ to $x_t$, overall minimizing a mean squared error loss, since the Gaussian includes the quadratic term.
In the discrete case, the \textit{derivation} for the objective function of the KL-divergence results in the cross-entropy loss $L_\text{CE} = -\sum_{k}^K p_k \log(\hat{p}_k)$ as discussed in [1].  
In both cases, the inherent connection to maximum likelihood and assumption of the data type results in the mean-squared error or cross-entropy loss. 
However, it is true that one could utilize a mean squared error on the predicted class probabilities resulting in the loss term $\sum_k (p_k - \hat{p}_k)^2$. Since the softmax cross-entropy loss has been established in the supervised learning formulation, especially in natural language and image processing, we believe it is the first choice to implement, next to our initial motivation to use the proper loss function when deriving from the KL divergence minimization.

Reviewer comment (Reproducibility):
The lack of the reproduction of the MiDi code benchmarks that the authors report in the supplement is interesting, though the reviewer thinks it probably would not change the results in a qualitative way. The reviewer asks whether the authors contacted the authors of MiDi to resolve this issue and whether they received any feedback.

Author response:
We contacted the authors of MiDi, but we did not get enlightening answers to our questions. As we used not only the provided checkpoint but also re-trained the MiDi model using the provided code and could in both cases not verify or reproduce the results, we assume that the authors made a mistake in reporting. Nevertheless, it is true there are no qualitative changes for the MiDi publication.

Reviewer comment (Evaluation):
The observation by the authors regarding deteriorating performance upon pretraining on explicit hydrogen on a larger molecular set is an interesting point that gets buried in the supplement. The reviewer asks whether there is a generalizable lesson and possibly a better or different way to overcome this problem than dropping the hydrogens, or if this observation is too close to the level of noise.

Author response:
We do not have a clear answer for the inferior performance when pretraining on PubChem3D with explicit hydrogens. One hypothesis is that the diffusion model has a simpler task during finetuning to learn to add hydrogens onto a molecule as opposed to relearn a distribution shift of the attached hydrogens.
To investigate this, we subsampled 1M molecules from PubChem3D and GEOM-Drugs and enumerated all bonds that involve hydrogens. We computed distances and noticed that the hydrogen-oxygen distances distribution in PubChem3D seems to have a smaller variance compared to GEOM-Drugs as illustrated in Fig. 8 in the Appendix of the updated manuscript. 
Nevertheless, our results suggest that the diffusion model can effectively learn a distribution shift underpinning the importance of pre-training on large 3D databases independent of the level of theory and the presence of explicit hydrogens. The pre-trained model yields superior performance while accurately learning the underlying data distribution faster from less data of the target distribution (Fig. 2, Appendix Fig. 8).

Reviewer comment (Novelty):
The novelty is limited. Most explored designs are easy to think about, like time-dependent loss weight, modeling discrete or continuous atom/bond types, parameterizing to match noise or x0, etc. The authors didn't come up with new designs.

Author response:
Our extensive analysis and benchmark with many ablations of different design choices on our EQGAT-diff model provides a profound contribution to the machine learning community, especially for researchers with applications in the bio-chemistry domain. Some implemented design choices have been initially proposed in the image and language processing community. Nevertheless, We argue that our work still contains highly relevant in-depth experimentation. We set a new state-of-the-art emphasizing important design choices crucial for successfully training and validating diffusion models on molecular data. Our findings lead to faster training convergence and increased inference speed while providing significantly higher stability and validity on larger molecular structures.

Reviewer comment (Experiments):
One important baseline MolDiff [1] is not compared with, although it has been cited. 

[1] Peng, X., Guan, J., Liu, Q., & Ma, J. MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation. ICML 2023.

Author response:
MolDiff utilizes a different pre-processing and splitting of the GEOM-Drugs dataset compared to EDM and MiDi while using an enhanced post-processing pipeline that fixes valency and aromaticity for generated molecules before computing validity and success rate. We adopted the data and evaluation pipeline from EDM and MiDi, which is not directly comparable to the results presented by MolDiff.

In the updated manuscript, we added results following the evaluation pipeline from MolDiff's GitHub repository (commit \#861e82c) and evaluated our generated samples. We report the mean validity and mean success rate in Tab. 10 [edited] in the Appendix.

Our proposed EQGAT-diff model achieves superior performance over MolDiff in generating chemically valid molecules, although MolDiff additionally utilizes uncertainty guidance via a separate bond-classifier model. 

We have carefully considered and addressed the concerns you raised. Building on well-known technical facts, we developed a state-of-the-art framework based on thorough experimental results. We considered many novel aspects in diffusion-based 3D modeling, like additional chemical features, denoising diffusion pre-training, and timestep-depending loss weighting. Our model improved upon standard metrics in this domain by large margins compared to established diffusion models like EDM, MiDi, and MolDiff. Furthermore, we have extended our framework to include not only classifier-guided molecular design but also structure-based ligand generation and find that our findings are also applicable in these situations, underpinning the importance of this work 

Thus, we ask the reviewer to consider raising the score.