Paper: 3D Autoencoding Diffusion Model for Molecule Interpolation and Manipulation (ICLR-2024)

Reviewer comment (Novelty):
The first major weakness of this paper is that it seems to ignore relevant literature in this area. Interpolation has been used and demonstrated for a long time [1, 2], the latent traversal has also been discovered before [3]. Despite this paper strikes the point that it focuses on diffusion models, but the authors should still consider proper comparisons or acknowledgments.

Author response:
1. Comparison with interpolation and latent space optimization for 2D molecules

Our focus is on semantics-guided diffusion in **3D molecular space**. We are aware of the existing works on interpolation and latent space optimization for molecules, but the majority of them are on 2D molecular graphs. The generation of 3D molecular structures poses different challenges such as preserving the symmetry, and incorporating semantic guidance into the generative process is a rarely explored topic. We will cite the mentioned papers, but it is not possible for direct comparison as these are different generation tasks.

Reviewer comment (Novelty):
Given this problem has been studied before, especially interpolation, the technical contribution of this paper is very limited. Another major contribution claimed in the paper (learning a semantics-guided autoencoding diffusion model) has also been proposed in [4]. The formulation looks exactly the same.

Author response:
2. Contribution of the paper

We have to disagree on this claim. Our main contribution is achieving semantic guidance on a **3D equivariant diffusion model**, which is a challenging task considering the complexity of the 3D space and the equivariance of the diffusion process. 3D interpolation is very different from that in 2D molecules or images. Also, learning informative semantics for 3D molecules is not trivial since we need a learning paradigm that can effectively leverage both the invariant and equivalent features. Please refer to the main rebuttal part 3 for extended discussion.

Reviewer comment (Reproducibility):
It is not surprising that linear interpolation could lead to improved properties as studied widely in images and some in molecules (e.g. [3]), but the experimental details need to be included for a fair comparison and evaluation.

Author response:
3. Interpolation and experimental details 

As we have mentioned above, the generation of 3D molecules is a challenging task. Also, experimental details for the interpolation have been given in 5.4.

Reviewer comment (Theory):
[Minor] The explanation of stability and validity seems to be different from the literature, see e.g. [5].

Author response:
5. Stability and validity metrics

We’d first like to clarify that our metric for “stability” is the same as Hoogeboom et al. (i.e. whether the allowed valency for the atoms is violated). However, as we are not modeling with explicit hydrogen, fragmented molecules can still be considered “stable” after adding the hydrogens back according to the allowed valency. This results in many false positives. Because of that, our “validity” metric adds another connectivity check, addressing the limitations of the previous metrics.

Furthermore, for better comparison, we now report the validity metrics from MiDi, which are not very different from ours. Please see the main rebuttal part 4.1 for more details.

Reviewer comment (Presentation):
Tables and figures are insufficiently described in the captions.

Author response:
Thanks for the comment. We have already discussed all tables and figures extensively in the Results section. We also add some descriptions to the captions in the revision.

Reviewer comment (Evaluation):
The paper only presents validity and stability metrics, but no similarity measures to examine the effectiveness of the template-guidance. Beyond that, the quality of the 3d geometries are not evaluated (e.g. energies and atomic forces).

Author response:
Thanks for the comment. We now also report similarity measures and energy estimated by xTB. Please find more details in the main rebuttal part 4.

Reviewer comment (Evaluation):
egDiffAE does not perform well in terms of stability (e.g. QM9 template-based generation 85.8% vs de-novo methods like EDM 90.7% [Hoogeboom et al] / Midi 97.5 [Vignac et al 23]). GEOM mol stability metrics are not given (only atom stability).

Author response:
We’d like to point out that the template-based (which we now call “semantics-guided” to avoid confusion) generation in our study is not directly comparable with *de novo* generation results in those works because:

(1) Our generation task is different from *de novo*. Please refer to the main rebuttal part 2 for more details.

(2) We do not explicitly include hydrogen. As mentioned in the main rebuttal part 4.1, the implicit hydrogen also leads to different validity metrics as in those works.

We’ve also shown *de novo* generation results in Appendix Table 7. We believe there is room for improvement in the quality of both tasks in future studies.

Reviewer comment (Evaluation):
The property manipulation does not seem to meaningfully alter the template, but mostly move atoms around and break/form some bonds. As a result, the results of property manipulation do not appear to be in equilibrium, but arbitrary conformations. Energies and atomic forces need to be carefully checked, since the optimized molecular properties are only meaningful in equilibrium. Beyond that, as far as I can tell, the property models have been trained on equilibrium molecules, i.e. they might be out of their domain of applicability at arbitrary conformations.

Author response:
First of all, we have examples where atom types are modified in Fig 4. We also have discussed the limitations of the current results in Conclusions. 

As our model is trained on QM9, which only contains equilibrium molecules, we’d assume the generated molecules should resemble equilibrium ones. Thus, property models on equilibrium molecules could be applied to the generations, similar to EDM.

We use xTB to estimate the energy of the molecules. As shown in Appendix Table 8, for most properties, the energy distribution of the generated molecules shifts slightly higher from the test set (equilibrium), which means most of the generated molecules have near-equilibrium conformations.

Reviewer comment (Novelty):
What is the practical relevance of molecule interpolation?

Author response:
With interpolation, we could effectively explore and exploit the molecular space looking for novel “intermediate” molecules with properties of interest, which is an extremely challenging problem due to the huge space for 3D molecules.  For example, based on two different molecules with similar properties, we can find more novel candidates with similar properties and “intermediate” structures if we can achieve a smooth interpolation.

Reviewer comment (Evaluation):
How close are the generated / manipulated molecules to equilibrium?

Author response:
As mentioned above, we use xTB to evaluate the energy of the generated/manipulated molecules (Table 1, Appendix Table 8). We find most of the generated molecules have near-equilibrium conformations from the energy distribution.

Reviewer comment (Theory):
Why use MMD regularization instead of KL prior regularization as in VAE?

Author response:
Theoretically, KL regularization could achieve similar results, but it is known to be unstable to train and subject to problems like posterior collapse [1], and requires careful tuning. 

[1] Bowman, Samuel R., et al. "Generating sentences from a continuous space." arXiv preprint arXiv:1511.06349 (2015).

Reviewer comment (Reproducibility):
How large is the batch size? MMD usually requires large batch sizes.

Author response:
The batch size is 128 molecules, as stated in Appendix A.5.

Reviewer comment (Presentation):
From the figures, it seems that interpolation and manipulation is only possible for molecules with same number of atoms?

Author response:
As the diffusion model requires an initial noise with a specified atom number, it is more straightforward to use the same atom number, especially for interpolation where we also need to interpolate the initial noises. Nevertheless, we could also potentially work with different atom numbers, but it adds more to the complexity of the evaluation.

Reviewer comment (Experiments):
Why are there no figures showing results on GEOM? This might be important to evaluate how the generation performs on larger molecules.

Author response:
Thanks for the suggestion. We add GEOM generation results in Appendix Fig 7 and 10. Similar to ZINC, the generations on GEOM tend to be disconnected. This is a known limitation with models solely relying on coordinates of atoms like in EDM, as is discussed in Appendix A8. For future work, we could improve the generation quality on larger molecules by explicitly modeling the bonds, etc.

Reviewer comment (Writing):
The paper is not very well written. Many definitions are wrongly used or used too loosely (eg, “de novo” or “template-based”) and a lot of implementation details are missing, making reproduction difficult.

Author response:
1. Use of terms and experimental details

We believe *de novo* is properly used as we are randomly sampling molecules from the molecular space without any specification in the respective experiments. We do acknowledge that “template-based” may be misleading so we changed that wording to “semantics-guided”, reflecting the real nature of our experiment. Please see the main rebuttal part 1 for more details.

Also, we’ve included many implementation details in the Appendix, including on the data processing (A.3), model architecture (A.4), and model hyperparameters and training (A.5). We will make our references in the main text to the Appendix clearer. Also, if you could specify the missed implementation details, we will be happy to clarify and add them.

Reviewer comment (Novelty):
There are citations missing (many work do some kind of semantic guiding, eg Hoogeboom et al, 22, all the work on pocket-conditioning, linking, fragment-based generation, etc).

Author response:
2. Comparison with ligand and linker designs, and inpainting (also responding to question 2)

It is noteworthy that those works do not have the same “guidance” as ours. We consider those works as more of an externally **conditioned** generation, where a molecule (or a part of it) is generated conditioned on known pockets/fragments. There is no explicit control over the “semantics” (i.e. the composition and shape) of the generated molecule itself in those problems. Here our semantics contains direct information about the entire molecule, rather than the condition. Please refer to the main rebuttal part 2. We’ve also discussed those works in the Related Works.

We do believe these are interesting applications for 3D diffusion models, and that our framework could be applied to them by leveraging both the condition and the explicit control of certain parts of the structure, and compared with those papers. However, it may require major modifications to the current training and generation framework to perform inpainting, so we’d like to leave it for further studies.

Reviewer comment (Theory):
The idea of conditional generation in some semantic space is interesting. However, the way it has been proposed seems ad-hoc to me. It is unclear what the “semantic encoding” is learning if the only loss is MMD with Gaussian.

Author response:
3. The loss for the semantics embedding

We’d like to clarify that the semantic embedding is used as an input to guide the diffusion process. Thus, the diffusion model’s noise estimation loss (equivalent to the reconstruction loss in traditional autoencoder settings) is also back-propagate to the encoder, and the MMD loss is an additional regularization term. Thus, as is proved in 4.3, it also contains information about the input molecule. This is similar to an autoencoder framework where the diffusion model effectively functions as a “decoder”.

Reviewer comment (Experiments):
Many work propose some kind of inpainting for scaffold/linking (eg, DiffHopp (Torge et al 23), SBDDDiff (Schneuing et al 22), DiffLinker (Igashov et al 22), etc). It would be good to compare the performance of the model with some inpainting baseline. How does the proposed method differs from those approaches?

Author response:
2. Comparison with ligand and linker designs, and inpainting (also responding to question 2)

It is noteworthy that those works do not have the same “guidance” as ours. We consider those works as more of an externally **conditioned** generation, where a molecule (or a part of it) is generated conditioned on known pockets/fragments. There is no explicit control over the “semantics” (i.e. the composition and shape) of the generated molecule itself in those problems. Here our semantics contains direct information about the entire molecule, rather than the condition. Please refer to the main rebuttal part 2. We’ve also discussed those works in the Related Works.

We do believe these are interesting applications for 3D diffusion models, and that our framework could be applied to them by leveraging both the condition and the explicit control of certain parts of the structure, and compared with those papers. However, it may require major modifications to the current training and generation framework to perform inpainting, so we’d like to leave it for further studies.

Reviewer comment (Evaluation):
The biggest concern is on the evaluation of the proposed method, as there are no baselines that are directly comparable for the main experiments (Table 1, Table 2, Table 3), and the used evaluation metrics are sometimes insufficient to prove the value of the proposed method.

Author response:
We acknowledge that there aren’t many existing baselines to compare with our method, except for interpolation where we could directly compare with a simple DDIM. For example, regarding molecular properties, most existing studies rely on conditioned generation given the properties, with little to no control over other aspects of the generated molecule. In this study, we seek finer control on not only the properties but also the composition, overall shape, etc. with the semantics embedding. Thus, for most parts of our work, we are trying to prove the generated molecules are of good quality and could achieve the conditioning we want. 

As mentioned in the main rebuttal part 3, here we try to demonstrate the usefulness of semantics guidance for the generation on a very well-established basis and dataset. We believe this framework could be extended to achieve finer guidance on more complex generative tasks, such as linkers and pocket-based ligands, by incorporating some recent progress in 3D molecule modeling dedicated to those tasks.

Reviewer comment (Evaluation):
In Table 2, the used evaluation metrics are somewhat confusing. For `smoothness` and `midpoint similarity`, it is unclear how the similarity scores are calculated — whether cosine similarity on embeddings is used or other similarity metrics are applied directly on molecules.

Author response:
The similarity scores are defined as the cosine similarity between the embeddings. For an interpolation trajectory, our desired characteristics are:

(1) Overall smooth transitions between neighboring points (low smoothness mean)

(2) Near-uniform steps taken across the trajectory (low smoothness std)

(3) The midpoint being a proper mixture of the two endpoints, rather than something entirely different (high midpoint similarity) 

Specifically for (3), we can observe in the interpolation with standard DDIM that the midpoint is a new molecule that does not resemble either of the endpoints. Thus, it is not a smooth trajectory from one end to the other. Please find more detailed descriptions in 5.2 and 5.4.