Paper: Hierarchical Graph Latent Diffusion Model for Molecule Generation (ICLR-2024)

Reviewer comment (Experiments):
Although this work states that the (hierarchical) latent approach for graph generation provides a scalable solution for molecule generation, the provided experiments are limited to datasets (e.g., GuacaMol) in which previous diffusion models (e.g., GDSS and DiGress) are applicable. In order to justify the scalability of the proposed method, it should be evaluated in a larger dataset.

Author response:
Thank you for your suggestion. During the rebuttal period, we have added experimental results on a larger molecule dataset, MOSES dataset. The results are shown in the Table below. Our proposed method achieved better Valid, Novelty, Diversity and comparable Unique scores compared to the baseline methods on the MOSES dataset.

**Table 1: Unconditional Generation Results on MOSES Dataset**
| Methods | Valid $\uparrow$           | Unique$\uparrow$           | Novelty$\uparrow$          | Diversity$\uparrow$        | FCD$\downarrow$             | NSPDK MMD$\downarrow$      |
|-----------------|------------------------------------|------------------------------------|------------------------------------|------------------------------------|-------------------------------------|------------------------------------|
 | GDSS    | **1.000 $\pm$ 0.000** | 0.994 $\pm$ 0.003          | 0.999 $\pm$ 0.000          | 0.899 $\pm$ 0.000          | 21.265 $\pm$ 0.249          | 0.037 $\pm$ 0.005          |
| DiGress | 0.858 $\pm$ 0.005          | **1.000 $\pm$ 0.000** | 0.996 $\pm$ 0.001          | 0.886 $\pm$ 0.000          | **9.228 $\pm$ 0.081**  | **0.010 $\pm$ 0.000** |
 | PS-VAE  | **1.000 $\pm$ 0.000** | 0.999 $\pm$ 0.000          | **1.000 $\pm$ 0.000** | 0.905 $\pm$ 0.000          | 26.401 $\pm$ 0.078          | 0.079 $\pm$ 0.000          |
 | GLDM    | **1.000 $\pm$ 0.000** | 0.998 $\pm$ 0.000          | **1.000 $\pm$ 0.000** | 0.905 $\pm$ 0.000          | 26.365 $\pm$ 0.095          | 0.077 $\pm$ 0.001          |
 | HGLDM   | **1.000 $\pm$ 0.000** | 0.999 $\pm$ 0.000          | **1.000 $\pm$ 0.000** | **0.906 $\pm$ 0.000** | 25.815 $\pm$ 0.053          | 0.072 $\pm$ 0.000          |

Reviewer comment (Experiments):
The experimental setting for evaluating the computational efficiency is not clear. Is the training and sampling time measured in the same condition, e.g., training conducted via DDP and using the same number of V100 GPUs?

Author response:
Sorry for the confusion. The training and sampling time are measured under the same conditions. We have clarified this in the revised version.

Reviewer comment (Evaluation):
Generation performance on unconditional molecule generation tasks should be evaluated with more descriptive metrics, for example, FCD, Scaffold similarity [1], and Fragment similarity [1]. Reported metrics, i.e., validity, uniqueness, and novelty fail to measure how similar (e.g., chemical aspects) are the generated molecules to the molecules from the test set. In particular, under the current setting, GDSS seems to be showing comparable results in large datasets (ZINC250K and GuacaMol) with significantly fewer parameters.

Author response:
Thank you for your suggestion. In the revised version, we have added FCD and NSPDK to evaluate the graph distribution of the generated molecules. The detailed results are shown in the Tables 1-4 below. Although the proposed method does not achieve best results in FCD and NSPDK, this is mainly limited by the VAE backbone model, and we have improved upon PSVAE in terms of  FCD and NSPDK across all the datasets. Furthermore, unlike general graph generation that primarily focuses on learning graph distributions, we believe that comparing distribution similarity in molecular generation tasks without considering molecular properties lacks practical significance. Therefore, after verifying that the performance of the proposed method in terms of metrics such as validity and novelty are comparable, indicating that our model performs well in unconditional generation, we shift our focus to conditional molecular generation  that can incorporate more conditional information into the generation process.

Reviewer comment (Evaluation):
The quantitative results of Tables 2 and 3 show that the performances of GLDM and HGLDM on unconditional generation tasks are almost the same, whereas there is a significant improvement using the hierarchical approach for conditional generation tasks. What is the reason for the hierarchical approach only effective in conditional tasks?

Author response:
Because the properties of molecules are often determined by distinctive functional groups (subgraphs), HGLDM performs better than GLDM in conditional generation tasks by capturing subgraph information. We have inclued this explanation in the revised version.

Reviewer comment (Evaluation):
The performance of GLDM (and HGLDM) comes from the effectiveness of using a latent representation of graphs compared to previous graph diffusion models, not from the diffusion processes. Thereby, analysis of the latent representation, e.g., interpolation in the latent space or clustering of the latent points with respect to certain conditions, would greatly strengthen this work.

Author response:
Thank you for your suggestion. During the rebuttal period, we have conducted a new experiment on the interpolation of latent embeddings. we randomly sample two latent embeddings, $\mathbf{z}^{G_1}$ and $\mathbf{z}^{G_2}$, and generate the interpolated embeddings $G'$ through $\mathbf{z}^{G'}=\lambda\mathbf{z}^{G_1}+(1-\lambda)\mathbf{z}^{G_2}$, where $\lambda \in (0, 1)$. Based on the visualized results of these latent embeddings (Figure 4 in the revised version), we can draw two key conclusions. Firstly, the clear boundaries between embeddings of different properties in the scatter plot, along with the high property values retained by the interpolated embeddings, suggest that **latent embeddings can effectively distinguish between molecules with varying property values**. This provides a solid foundation for further conditional generation using these latent embeddings. Secondly, **latent embeddings trained by HGLDM capture more distinct features of different property values compared to PS-VAE latent embeddings**, as evidenced by the farther distances between HGLDM latent embeddings with different property values and the higher property values of molecules decoded from HGLDM latent embeddings interpolation. This highlights the enhancement of the latent diffusion framework over VAE-based models (Please refer to General Response G1 and Section 5.2 in the revised version  for more details).

Reviewer comment (Novelty):
Missing references on related works:
  - Qiang et al., Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D, ICML 2023
  - Xu et al., Geometric Latent Diffusion Models for 3D Molecule Generation, ICML 2023

Author response:
Thank you for your suggestion. We have added these references to the Related Work section in the revised version.

Reviewer comment (Reproducibility):
Is the results of Table 2 from a single run or an average of multiple runs?

Author response:
The original results of Table 2 is from a single run. During the rebuttal period, we have conducted eight samplings with  different seeds and reported the mean and variance in the revised version. The new results are shown in the Tables 1-4 above.

Reviewer comment (Theory):
It is not clear why the authors select a way to generate graph embeddings and then decode them. Accessing the graph embedding could contain less information than accessing the subgraph- or node- and edge-level latent variables.

Author response:
We agree with your viewpoint that "Accessing the graph embedding could contain less information than accessing the subgraph- or node- and edge-level latent variables."  We need to clarify that **the main contribution of this paper is the HGLDM model, which includes graph embeddings, subgraph embeddings, and node embeddings.** The reason why edge-level embeddings are not included is that, as demonstrated by the baseline methods in the Section 5.1, the diffusion process on the edge matrix is computationally expensive. The graph embedding-based model GLDM proposed in this paper is just a straightforward attempt inspired by the latent diffusion model, as there exists the decoder that can generate molecules from graph embeddings. However, in this paper, we found that the hierarchical embeddings proposed in this paper are more effective than the graph embeddings, as the latter contain less information, as you mentioned.

Reviewer comment (Evaluation):
The authors report the validity, uniqueness, and novelty as the main results. However, these metrics seem to be restricted to only measure the sample diversity.

Author response:
For the validity, we need to clarify that during the decoding stage, the PS-VAE only performs link prediction to connect fragments and does not perform validity check. In the revised version, we have added FCD and NSPDK to evaluate the graph distribution of the generated molecules. The new results are shown in Tables 1-4 below. Although the proposed method does not achieve best results in FCD and NSPDK, this is mainly limited by the VAE backbone model, and we have improved upon PSVAE in terms of  FCD and NSPDK across all the datasets. **Furthermore, unlike general graph generation that primarily focuses on learning graph distributions, we believe that comparing distribution similarity in molecular generation tasks without considering molecular properties lacks practical significance**. Therefore, after verifying that the performance of the prpoposed method in terms of metrics such as validity and novelty are comparable, indicating that our model performs well in unconditional generation, we shift our focus to conditional molecular generation  that can incorporate more conditional information into the generation process.

Reviewer comment (Reproducibility):
In Section 4.1.1, to my understanding, the number of subgraph embeddings should not be the number of nodes. If so, how do you sample the number of subgraph embeddings at the beginning of the sampling stage?

Author response:
Because there is no node and subgraph alignment issue in the Decoder of PSVAE, during diffusion training, we replicate the subgraph embedding across all nodes within the subgraph to form the subgraph latent embedding. This ensures that the subgraph latent embedding has the same dimension as the node-level embedding. During sampling, we sample subgraph embeddings from the same dimension and match them to the embeddings of the subgraphs where each node belongs.

Reviewer comment (Reproducibility):
How do you get the subgraph embeddings from the PS-VAE architecture?

Author response:
We have followed the approach of PSVAE. First, we use the principal subgraph extraction algorithm in PS-VAE to construct the subgraph vocabulary. Then, we obtain the subgraph embeddings through the embedding layer of the encoder in PS-VAE.

Reviewer comment (Evaluation):
In Table 3, how did you measure the mean absolute error (MAE) on the unconditional setting? Does it mean training on the selected 100 molecules without the conditions?

Author response:
Yes, we train the unconditional generation structure based on the selected molecues and measure the MAE between the generated molecules and the selected molecules.

Reviewer comment (Evaluation):
Why are some reported values different from the original papers? For example, the novelty and uniqueness of PS-VAE on the QM9 dataset and the validity of GDSS on the ZINC250k dataset.

Author response:
For the PS-VAE model, we found that lowering the weight of the KL loss can lead to better results than those reported in the original paper. We explain this phenomenon in Appendix B. The reason for the differences in the GDSS results is that we reproduced the GDSS code and discovered that we achieved better  performances. As a result, we reported the improved results.

Reviewer comment (Evaluation):
I'd like to see error bars indicating variance accross multiple seeds if possible.

Author response:
Thank you for your suggestion. During the rebuttal period, we have conducted eight samplings with  different seeds and reported the mean and variance in the revised version. The new results are shown as below:

 **Table 1: Unconditional Generation Results on QM9 Dataset**
 Methods | Valid $\uparrow$           | Unique$\uparrow$           | Novelty$\uparrow$          | Diversity$\uparrow$        | FCD$\downarrow$             | NSPDK MMD$\downarrow$      |
|-----------------|------------------------------------|------------------------------------|------------------------------------|------------------------------------|-------------------------------------|------------------------------------|
| GDSS    | 0.957 $\pm$ 0.000          | 0.982 $\pm$ 0.003          | 0.988 $\pm$ 0.001          | **0.925 $\pm$ 0.000** | 2.959 $\pm$ 0.040           | 0.003 $\pm$ 0.000          |
| DiGress | 0.992 $\pm$ 0.003          | 0.960 $\pm$ 0.001          | 0.391 $\pm$ 0.001          | 0.920 $\pm$ 0.000          | **2.123 $\pm$ 0.033**  | **0.001 $\pm$ 0.000** |
| PS-VAE  | **1.000 $\pm$ 0.000** | 0.981 $\pm$ 0.002          | 0.996 $\pm$ 0.000          | 0.881 $\pm$ 0.000          | 16.877 $\pm$ 0.059          | 0.059 $\pm$ 0.000          |
| GLDM    | **1.000 $\pm$ 0.000** | 0.982 $\pm$ 0.001          | 0.996 $\pm$ 0.000          | 0.884 $\pm$ 0.000          | 14.829 $\pm$ 0.169          | 0.050 $\pm$ 0.001          |
| HGLDM   | **1.000 $\pm$ 0.000** | **0.985 $\pm$ 0.002** | **0.997 $\pm$ 0.000** | 0.884 $\pm$ 0.000          | 14.576 $\pm$ 0.183          | 0.047 $\pm$ 0.001          |

**Table 2: Unconditional Generation Results on ZINC250K Dataset**
| Methods | Valid $\uparrow$           | Unique$\uparrow$           | Novelty$\uparrow$          | Diversity$\uparrow$        | FCD$\downarrow$             | NSPDK MMD$\downarrow$      |
|-----------------|------------------------------------|------------------------------------|------------------------------------|------------------------------------|-------------------------------------|------------------------------------|
| GDSS    | **1.000 $\pm$ 0.000** | 0.997 $\pm$ 0.001          | **1.000 $\pm$ 0.000** | 0.902 $\pm$ 0.000          | 16.086 $\pm$ 0.071          | **0.018 $\pm$ 0.000** |
| DiGress | 0.565 $\pm$ 0.005          | **1.000 $\pm$ 0.000** | **1.000 $\pm$ 0.000** | 0.882 $\pm$ 0.000          | **13.042 $\pm$ 0.164** | 0.031 $\pm$ 0.001          |
 | PS-VAE  | **1.000 $\pm$ 0.000** | 0.993 $\pm$ 0.001          | **1.000 $\pm$ 0.000** | 0.912 $\pm$ 0.000          | 20.386 $\pm$ 0.061          | 0.085 $\pm$ 0.001          |
| GLDM    | **1.000 $\pm$ 0.000** | 0.994 $\pm$ 0.001          | **1.000 $\pm$ 0.000** | 0.913 $\pm$ 0.000          | 20.444 $\pm$ 0.088          | 0.086 $\pm$ 0.001          |
 | HGLDM   | **1.000 $\pm$ 0.000** | 0.997 $\pm$ 0.001          | **1.000 $\pm$ 0.000** | **0.914 $\pm$ 0.000** | 19.913 $\pm$ 0.091          | 0.084 $\pm$ 0.000          |

**Table 3: Unconditional Generation Results on Guacamol Dataset**
 | Methods | Valid $\uparrow$           | Unique$\uparrow$           | Novelty$\uparrow$          | Diversity$\uparrow$        | FCD$\downarrow$             | NSPDK MMD$\downarrow$      |
|-----------------|------------------------------------|------------------------------------|------------------------------------|------------------------------------|-------------------------------------|------------------------------------|
| GDSS    | **1.000 $\pm$ 0.000** | 0.986 $\pm$ 0.001          | 0.996 $\pm$ 0.001          | 0.892 $\pm$ 0.000          | 40.291 $\pm$ 0.072          | 0.058 $\pm$ 0.000          |
| DiGress | 0.875 $\pm$ 0.005          | **1.000 $\pm$ 0.000** | **0.999 $\pm$ 0.001** | 0.904 $\pm$ 0.000          | **12.069 $\pm$ 0.051** | **0.018 $\pm$ 0.000** |
 | PS-VAE  | **1.000 $\pm$ 0.000** | 0.998 $\pm$ 0.000          | 0.998 $\pm$ 0.000          | **0.905 $\pm$ 0.000** | 24.105 $\pm$ 0.082          | 0.090 $\pm$ 0.000          |
 | GLDM    | **1.000 $\pm$ 0.000** | 0.998 $\pm$ 0.000          | 0.998 $\pm$ 0.000          | 0.904 $\pm$ 0.000          | 23.879 $\pm$ 0.041          | 0.095 $\pm$ 0.000          |
| HGLDM   | **1.000 $\pm$ 0.000** | 0.999 $\pm$ 0.001          | **0.999 $\pm$ 0.000** | **0.905 $\pm$ 0.000** | 23.845 $\pm$ 0.098          | 0.095 $\pm$ 0.001          |

Reviewer comment (Experiments):
While QM9 and ZINC250k performance is imprressive, these graphs are kind of solved. What is the performance on MOSES or shapenet?

Author response:
Thank you for your suggestion. During the rebuttal period, we have added experimental results on the MOSES dataset. The new results are shown in Table 4. Our proposed model achieved better Valid, Novelty, Diversity and comparable Unique scores compared to the baseline methods on the MOSES dataset. In the newly added FCD and NSPDK metrics, although our proposed HGLDM does not outperform diffusion-based methods, it shows further improvement compared to the backbone model PSVAE. Additionally, we believe that molecular generation tasks, which focus more on task-specific molecular properties, have greater practical value than general graph generation tasks that primarily emphasize learning graph distributions. Therefore, during the rebuttal period, we included a latent embedding interpolation experiment to further illustrate the advantages of the latent diffusion model framework in molecular generation tasks. For more details, please refer to General Response G1 and Section 5.2 in the revised version. In addition, the ShapeNet dataset we retrieved is a 3D point cloud data, while our work mainly focuses on 2D molecular generation, so it is not applicable for us to use this dataset.

Reviewer comment (Experiments):
To clarify, you are evaluating all datasets without hydrogens?

Author response:
Yes, we removed the hydrogens.

Reviewer comment (Theory):
Equation (2) is not correct. The probability for each node is computed twice. I think the correct definition should be $\prod p(x_i) \prod \prod (e_{ij})$.

Author response:
Thank you for your correction. We have revised it to the correct form in the revised version.

Reviewer comment (Theory):
It is not reasonable to design a hierarchical diffusion model. There is no need to sample three variables $z^x, z^M, z^G$ at the same time. As a hierarchical model, the decoding process of PS-VAE is $ G \sim q(G|z^M)q(z^M|z^G)$. That is, decoding a subgraph from a graph-level vector by a GRU, and then predicting the connection for the subgraphs. So actually, we only need to define a diffusion model over graph-level $z^G$. During the sampling process, we first sample $z^G$ from the diffusion model, and then use decoding of PS-VAE to get the graph. The current framework actually learns the decoding part twice, during the training of PS-VAE, the relationship between each level has been learned already. However, the diffusion model learns it one more time.

Author response:
Sorry for the misunderstanding, we need to make some clarifications here. 

**Clarification 1**: **The idea of "we first sampling $z^G$ from the diffusion model, and then using decoding of PS-VAE to get the graph" is actually the GLDM model that we proposed in this paper.** However, due to its unsatisfactory experimental results, we further propose the hierarchical model HGLDM to capture the node-level and structural information. Our experimental results demonstrate the improvement of introducing node-level and subgraph-level embeddings on generating graph embeddings. Furthermore, we argue that during the training of PS-VAE, the graph level embedding is obtained through sum pooling, the relationships between each level have not been well learned. **Through the hierarchical denoising neural network proposed in this paer, we learn the relationships between each level to ensure that the latent embeddings generated through diffusion contains more node-level and subgraph information than the graph latent embedding encoded by PS-VAE.** During the rebuttal period, we conducted a new experiment on the interpolation of latent embeddings, which further demonstrates the enhancement of the HGLDM over PSVAE. we randomly sample two latent embeddings, $\mathbf{z}^{G_1}$ and $\mathbf{z}^{G_2}$, and generate the interpolated embeddings $G'$ through $\mathbf{z}^{G'}=\lambda\mathbf{z}^{G_1}+(1-\lambda)\mathbf{z}^{G_2}$, where $\lambda \in (0, 1)$. Based on the visualized results of these latent embeddings (Figure 4 in the revised version), we can draw two key conclusions. Firstly, the clear boundaries between embeddings of different properties in the scatter plot, along with the high property values retained by the interpolated embeddings, suggest that **latent embeddings can effectively distinguish between molecules with varying property values**. This provides a solid foundation for further conditional generation using these latent embeddings. Secondly, **latent embeddings trained by HGLDM capture more distinct features of different property values compared to PS-VAE latent embeddings**, as evidenced by the farther distances between HGLDM latent embeddings with different property values and the higher property values of molecules decoded from HGLDM latent embeddings interpolation. This highlights the enhancement of the latent diffusion framework over VAE-based models (Please refer to General Response G1 and Section 5.2 in the revised version  for more details).

**Clarification 2**: "The current framework actually learns the decoding part twice" is incorrect. As described in Section 4.3 and Algorithm 1 in the Appendix, **the decoder is only trained once**. When training the diffusion model, we only predict the noise added to the latent embedding and do not need to train the decoder.

Reviewer comment (Evaluation):
The results lack MMD metrics. I think it is very important to check the distribution of the graphs.

Author response:
Thank you for your suggestion. We have included the NSPDK MMD metric in the revised version. NSPDK (Neighborhood subgraph pairwise distance kernel) MMD (Costa & De Grave, 2010) is the MMD between the generated molecules and test molecules which takes into account the node and edge features for evaluation. The new evaluation results are shown below. Although the proposed method does not achieve best results in NSPDK MMD, this is mainly limited by the VAE backbone model, and we have improved upon PSVAE in terms of  NSPDK MMD across all the datasets. **Furthermore, unlike general graph generation that primarily focuses on learning graph distributions, we believe that comparing distribution similarity in molecular generation tasks without considering molecular properties lacks practical significance**. Therefore, after verifying that the performance of the prpoposed method in terms of metrics such as validity and novelty are comparable, indicating that our model performs well in unconditional generation, we shift our focus to conditional molecular generation  that can incorporate more conditional information into the generation process.

Reviewer comment (Theory):
"However, these approaches sacrifice the random exploration ability to ensure that the final noisy data conforms to the appropriate discrete category distribution." Why do you make such claims? The definition for the distribution of the discrete variables is different. And people can also define a discrete diffusion process over them such as [1] and many other works. The performance is also very good and I think people should select models based on the specific problem. There is no any conclusion to support that continuous features is better than discrete process.

Author response:
Thank you for your question. We make this claim based on the discrete diffusion process in [1] you mentioned. It can be seen from Eq. (6) in [1] that their transition matrix is only related to the marginal distribution of the dataset without adding random noise.  Therefore, we claim that such methods sacrifice the ability of random exploration. From the experimental results, we can also observe that when conducting experiments on datasets with simpler data distributions such as QM9, the Novelty of DiGress decreases significantly compared to other diffusion methods. We agree that ''people should select models based on the specific problem, and there is no any conclusion to support that continuous features is better" because both methods perform good in unconditional task. However, **proving that the reverse generation process of a custom discrete diffusion satisfies the inverse process of its diffusion process is not easy, while this is easier to achieve in continuous Gaussain diffusion and can be more easily extended for conditional generation tasks.**