3D Autoencoding Diffusion Model for Molecule Interpolation and Manipulation

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.

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.

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.

4. Property manipulation

We’d like to clarify that:

(1) In fact, we are manipulating given molecule templates (which we now call “source” to avoid confusion). As described in 5.5, our property manipulation experiment takes the semantic embedding of a given source molecule, uses a simple linear model to alter its properties, and utilizes the altered embedding to generate a new molecule. We show in Fig 4 and Appendix Table 8 that molecules generated in this way largely resemble the source, while also gaining the altered property.

(2) Our framework is based on DDIM, which allows for deterministic sampling. In other words, every embedding+noise combination could be deterministically mapped to a data point in the molecular space. Thus, the manipulated embedding also maps to a specific molecule.

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.

1. Captions

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.

2. Similarity and energy metrics

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.

3. Comparison with de novo generation

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.

4. Metrics for GEOM

Our model is based on the framework of EDM, which only models the atom types and coordinates, and tends to produce disconnected structures for large molecules (Appendix A8). Thus, only atom stability is reported for GEOM, as in EDM.

5. Property manipulation

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.

6. Retrieval

We’d like to clarify that the “retrieval” method uses the semantic embedding learned from egDiffAE to directly select molecules from the training set with the desired property. This has been explained in 5.5. It is another way of obtaining molecules with desired compositions and properties using our semantics embedding. Its high performance is an expected result because:

(1) It searches for real molecules, so validity is guaranteed.

(2) It shows our semantic embedding contains sufficient information for the property-determining components, so if a known molecule has a similar embedding as z’, it means that the molecule should carry the same property. This further proves the effectiveness of the semantics embedding.

On the other hand, this baseline does not meaningfully generate any new molecules but only searches among the training set. Thus, it mostly offers a sanity check and an approximate upper bound of the performance. We have further clarified that in Table 3 and Section 5.5.

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.

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.

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”.

4. Evaluation metrics

4.1 Validity

Thanks for the suggestion. We now report the same validity metric as in MiDi. Interestingly, both MiDi’s valid metric and our previous valid are both stability (valency not violated) plus connectivity (successfully sanitized by RDKit), the results are almost the same as before.

4.2. Atom stability

For larger molecules, because we are building upon EDM’s framework (see main rebuttal part 2), we encounter the same problem where the generations are fragmented. Thus, we also rely on only atom similarity like in EDM. We could improve the generation quality on those datasets by incorporating some more recent advances in the 3D modeling of molecules (e.g. explicitly generating the bonds). We’d leave it for future studies.

4.3 Interpolation

We have explained in detail the smoothness metrics in 5.4. Basically, 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.

5. Novelty and uniqueness

Thanks for the comment. We now report the uniqueness and novelty using the same metrics as in MiDi (see Table 1 and Appendix A.7.3).

As shown in Table 1, we expect molecules generated with the same semantics embedding to have lower diversity than de novo generations. This is the exact purpose of our framework: to enforce sufficient control on the composition, shape, etc. of the molecule so that the same semantic embedding will result in mostly similar molecules regardless of the random noise. Then we can explore the molecular space with finer control over the direction.

If the goal is to generate diverse molecules, we could randomly sample a semantic embedding and use it for generation, as in Appendix Table 7.

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.

1. Evaluation metrics

Thanks for the suggestion. We’ve included similarity and uniqueness (i.e. diversity) in the results. Please refer to the main rebuttal part 4.3 for more details.

2. Interpolation metrics

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.