Diffusion Twigs with Loop Guidance for Conditional Graph Generation

Thank you very much for your many excellent suggestions! We have acted on all of these, and additionally, address all your questions and comments below.

RE: Twigs cannot guide unseen conditions (W1) (Q2)

Thanks for raising this point. Based on your suggestion we have added an experiment to tackle the issue that you raised. The Twig framework is very flexible, and it can be trained with unconditional setting, and combined with a classifier, to achieve generalization, as you mentioned. We provide here an example.

Specifically, we first train an unconditional version of Twigs, which consists of a diffusion model over $k$ graph properties. Secondly, we train a separate discriminator for the new property to achieve prediction guidance. This idea is the same as in classifier guidance, and in MOOD (Lee et al 2023), which allows the model to generalize to properties not seen during the training of the diffusion model.

The table below shows the unconditional version of Twigs trained with three properties (excluding the unseen property) and a separate property predictor (for the unseen property), on the community-small dataset. The results show that Twigs receives benefits from having the diffusion over multiple properties.

RE: experimental settings and evaluation pipeline are unclear (W2)

Thank you, based on your comment, we have added the hyperparameters used in the experiments, which we will integrate into the main text. In particular, we will clarify the following settings. For Sections 4.1 and 4.2 we follow the same hyperparameters from Huang et al (2023). For Section 4.3 we follow the hyperparameters from Lee et al (2023), for the MOOD baseline, we explore the OOD coefficient between $0.01$ and $0.09$. For Section 4.4 we follow the hyperparameters from Jo et al (2022).

Guidelines checklist (W3): Yes, we will bring it back.

RE: how step sizes affect generation. Q1

For step-size Sections 4.1 and 4.2, we follow the same hyperparameters from Huang et al (2023), including step-sizes of the sampling procedure. For Section 4.3 we follow the hyperparameters from Lee et al (2023), which leverages GDSS.

RE: can we have different main and secondary process stepsizes. Q1 To generalize to unseen properties, we do not need to separate the diffusion, we show above that we can train a separate predictor on unseen properties. Therefore we can maintain the same stepsizes in the diffusion flows.

RE: provide the training time comparison in multiple diffusion (Q3)

We show the impact of multiple diffusion flows on the community-small and Enzymes datasets. Specifically, we report the average time (in seconds) for the training of a single epoch for Twigs with one and three secondary diffusion flows. We observe that our models encounter a small overhead compared to GDSS and Digress, however, we believe it is a good tradeoff for performance.

Final comment

We are grateful for your thoughtful review. We hope our response has addressed your questions and concerns, and will appreciate the same being reflected in your stronger support for this work.

Thank you for the additional results. Given the above discussion, seems the proposed method provides limited performance improvement while introducing small computational overheads. Can you elaborate more on what is the advantage of the proposed method compared to MOOD and JODO?

Many thanks!

RE: computational cost (W1)

In Table 12 of paper, we present the inference times of our model compared to MOOD, showing that while we encounter a small overhead, we achieve superior performance.

In addition, we report below the average time (in seconds) required to train a single epoch for the community-small and Enzymes. For Twigs, we consider one property (density).

baseline methods (W2)

We provide a list of the recent methods (after 2023) in our experiments. In Sections 4.1 and 4.2, we compare Twigs with the following methods on QM9: TEDMol (Luo et al 2024), Jodo (Huang et al 2023), EEGSDE (Bao et al 2023), GeoLDM (Xu et al 2023), EquiFM (Song et al 2023). In Section 4.3, we compare with Lee et al (2023). In Section 4.4, we compare with Lee et al (2023), and Vignac et al (2023) on the community-small dataset.

Hyperparameters (W3)

Here we report the hyperparameters used in the experiments. For Sections 4.1 and 4.2 we follow the same hyperparameters from Huang et al (2023). For Section 4.3 we follow the hyperparameters from Lee et al (2023), for the MOOD baseline, we explore the OOD coefficient between $0.01$ and $0.09$. For Section 4.4 we follow the hyperparameters from Jo et al (2022).

Ablation study (W3)

We have added a new study to understand the impact of multiple properties over the community-small dataset. We report the results in the "Additional experiments with multiple properties" part of the global comment.

Refs

Huang et al (2023). Learning joint 2d & 3d diffusion models for complete molecule generation.

Lee et al (2023). Exploring chemical space with score-based out-of-distribution generation.

Jo et al (2022). Score-based generative modeling of graphs via the system of stochastic differential equations.

Bao et al 2023. Equivariant Energy-Guided SDE for Inverse Molecular Design.

Luo et al 2024. Text-guided diffusion model for 3d molecule generation.

Xu et al 2023. Geometric latent diffusion models for 3d molecule generation.

Song et al 2023. Equivariant flow matching with hybrid probability transport for 3d molecule generation.

Vignac et al 2023. Digress: Discrete denoising diffusion for graph generation

properties close to the guide (Q1)

The generated molecules' properties closely align with the target values, as evidenced by the MAE values presented in Tables 3, 5, and 8. The MAE quantifies the deviation between the desired ground truth properties and those of the generated molecules.

meanings of the properties (Q1)

The mentioned properties are standard quantum properties used in the QM9 dataset for modeling molecules (Hoogeboom et al 2022). The molecule properties are obtained from the data.

• $\alpha$ Polarizability: Tendency of a molecule to acquire an electric dipole moment when subjected to an external electric field.

• $\epsilon_{\text{HOMO}}$: Highest occupied molecular orbital energy.

• $\epsilon_{\text{LUMO}}$: Lowest unoccupied molecular orbital energy.

• $\Delta \epsilon$ Gap: The energy difference between HOMO and LUMO.

• $\mu$: Dipole moment.

• $C_v$: Heat capacity at 298.15K.

Refs:

Hoogeboom et al 2022. Equivariant Diffusion for Molecule Generation in 3D. ICML.

related works (Q2)

Our approach differs from the mentioned methods in two key ways, which we will elaborate on in the related works section of our paper.

Framework: The first two studies (Nisonoff et al. 2024, Gruver et al. 2023) are based on discrete frameworks. The third study (Klarner et al. 2024) adopts a plug-and-play approach using a diffusion model, which is agnostic to the underlying framework. Differently from the above, our method operates within a continuous framework utilizing stochastic differential equations (SDEs).

Contributions: Nisonoff et al. 2024 primary contribution is to provide a framework for diffusion guidance within discrete spaces. On the other hand, the main contribution of Gruver et al. (2023), addresses the challenges of optimizing discrete sequences by a novel sampling technique. Klarner et al. 2024 propose to learn two models one diffusion without labels and an additional discriminator model with property labels. Our contribution is separate from all the above methods, as we uniquely leverage multiple diffusion flows, defined over two distinct types of flows within a hierarchical structure. Specifically, we define a primary flow for capturing structure and a secondary flow for modeling properties.

Refs

Nisonoff et al (2024). Unlocking Guidance for Discrete State-Space Diffusion and Flow Models.

Gruver et al (2023). Protein Design with Guided Discrete Diffusion.

Klarner et al (2024). Context-Guided Diffusion for Out-of-Distribution Molecular and Protein Design.

Clarification on Figs (Q3-Q4)

Figure 2: Depits uncurated samples obtained via our model conditioned on $C_v$ property. The performance of conditioning is given in Table 3 First column ($C_v$). In addition, to validate the legitimacy, those molecules are representative of the "Molecule quality" results from Table 4.

Figure 3. It shows uncurated samples of molecules produced under two desired properties ($C_v$ and $\mu$). The performance corresponding to Fig 3 is represented in Table 5 first two columns ($C_v$ and $\mu$).

Note that in this setup, our goal is not to reconstruct a specific molecule there could be multiple different molecules that have the desired properties and do not reconstruct the training set. In this sense, the target properties should be considered separately from the molecule.

We would be grateful if this could be reflected in an increased score for our work. Thank you!

Many thanks.

Re: related works (W1)

The suggested works, including GDSS (Jo et al. 2022), GraphMaker (Li et al. 2024), and EDGE (Chen et al. 2023), are relevant to our research and will be included in our paper. Below we report the differences with our method. Refer to the table of the global comment for a summary.

First, these methods use different hierarchical notions in their frameworks. In EDGE and GDSS, edges and nodes are modeled symmetrically. In contrast, we introduce a unique hierarchical approach by learning multiple asymmetric diffusion flows. This distinction is crucial: while the mentioned methods use diffusion flows in the same roles, our method uses a secondary process to learn interactions with properties, feeding back into the main process to learn the graph structure.

Secondly, our methodology is designed specifically for generating graphs with desired conditional properties, which the mentioned papers do not focus on. GDSS and EDGE do not consider conditional modeling, and while GraphMaker explores node-label conditioning, this differs from our goal of generating graphs constrained by specific properties. Our method learns a joint distribution of graph structures (nodes, edges) and properties, unlike the mentioned methods that only learn a joint distribution of nodes and edges.

RE: conditional independence (W2)

Assuming conditional independence among the properties $\alpha$, $\epsilon_{\text{HOMO}}$, $\epsilon_{\text{LUMO}}$, $\Delta \epsilon$, $\mu$, and $C_v$ given the molecular graph can simplify the modeling process. This assumption leverages the fact that the molecular graph captures the essential structural dependencies, allowing us to treat the properties as independent for computational efficiency and ease of interpretation, even if slight interdependencies exist.

RE: Modeling multiple properties (W3)

We show here two additional results on 2 and 3 properties. We study the models over the community-small dataset. Specifically, our method (Twigs) is trained using multiple secondary processes (one for each property).

Two properties

We show the MAE results for property pairs for density. Our model achieves the lowest MAE in both cases. We notice that while the properties may present some form of correlation, our model can achieve a good performance in generating the graphs with desired properties.

Three properties

The Table below shows that our method achieves the lowest MAE values (the lower the better) across all three required properties.

RE: Hierarchical conditional diffusion (W4.1)

Here's a clarification of "Hierarchical Conditional Diffusion" and the distinction from "unconditional hierarchical models", which will be integrated into the paper.

In lines 39-43, we define "Hierarchical Conditional Diffusion" as follows: ".. rather than treating heterogeneous structural and label information uniformly within the hierarchy, we advocate for the co-evolution of multiple processes with distinct roles. These roles encompass a primary process governing structural evolution alongside multiple secondary processes responsible for driving conditional content."

To distinguish our method from "unconditional" hierarchical models by Jin et al. (2020) and Qiang et al. (2023), we label those models as "Hierarchical Modeling" in Table 1. While they model hierarchical structures, our approach is unique in leveraging a hierarchy of branching diffusion processes for conditional generation based on desired properties.

References

Jin et al. 2020. Hierarchical generation of molecular graphs using structural motifs.

Quiang et al. 2023. Coarse-to-fine: a hierarchical diffusion model for molecule generation in 3d.

RE: Variable $y_s$ (W4.2)
In the variable $y_{s,t}$ t indicates time. We will add to the main text.

RE: Dimension of graph (W4.3)

We have two cases:

The 3D case in (Appendix B.1), we denote the variable $y_s$ as a 3D graph $G = (A, x, h)$, with node coordinates $x= (x^1, \ldots, x^N) \in \mathbb{R}^{N \times 3}$, node features $h = (h^1, \ldots,h^N) \in \mathbb{R}^{N \times d1}$, and edge information $A \in \mathbb{R}^{N \times N \times d2}$.

The 2D case in Appendix B.2: we denote $y_s$ as a 2D graph with $N$ nodes we consider the variable $y_s=(X,A)\in \mathbb{R}^{N\times F}\times\mathbb{R}^{N\times N}$, where $F$ is the dimension of the node features, $X\in\mathbb{R}^{N\times F}$ are node features, $A\in\mathbb{R}^{N\times N}$ is weighted adjacency matrix. We define the perturbed property $y_i \in \mathbb{R}$ and the (fixed) property $y_C \in \mathbb{R}$.

RE: EDM and Tedmol results (Q1)

Apologies for the misleading results presentation. Due to space constraints, we placed the results for EDM and Tedmol in Table 10 of Appendix C.1. We will move it to the main text. These methods were included in the appendix because they are less competitive compared to the other baselines.

Final comment

We hope our response addresses your concerns and reinforces your support for this work.

Thank you for the detailed response, and I've read it carefully. Still, my main concern is the lack of studies on modeling >2 properties for real-world datasets like QM9, where the assumption of conditional independence can be insufficient.

Thank you so much for your thoughtful comments. We address all your concerns, as described below.

RE: Section 3 is written too generally and does not address the specific characteristics of the graph (W1)

We maintain the section in a general format to accommodate multiple cases, specifically one for 2D graphs and another for 3D graphs. We achieve the desired flexibility by introducing a variable $y_s$ that encompasses both node features and the adjacency matrix (2D case), as well as the coordinates (3D case). Below are the detailed descriptions of the $y_s$ variable.

RE: Dimension of the graph structure in Eqn (1) (W2)

We address two cases:

3D Case (Appendix B.1): We denote the variable $y_s$ as a 3D graph $G = (A, x, h)$, where node coordinates are represented as $x = (x^1, \ldots, x^N) \in \mathbb{R}^{N \times 3}$, node features as $h = (h^1, \ldots, h^N) \in \mathbb{R}^{N \times d_1}$, and edge information as $A \in \mathbb{R}^{N \times N \times d_2}$.

2D Case (Appendix B.2): Here, we denote $y_s$ as a 2D graph with $N$ nodes. The variable is defined as $y_s = (X, A) \in \mathbb{R}^{N \times F} \times \mathbb{R}^{N \times N}$, where $F$ represents the dimension of the node features. In this case, $X \in \mathbb{R}^{N \times F}$ are the node features, and $A \in \mathbb{R}^{N \times N}$ is the weighted adjacency matrix. Additionally, we define the perturbed property $y_i \in \mathbb{R}$ and the fixed property $y_C \in \mathbb{R}$.

RE: External Context $y_C$ and Dependent Graph Properties (Q1)

Indeed, the external context $y_C$ can be one or more of the $k$ dependent graph properties. Specifically, $y_C$ represents the graph property (or properties) upon which we are conditioning. For instance, in the context of conditional generative modeling for drug design, we might seek a molecule with a specific $\epsilon_{\text{LUMO}}$ value. In this scenario, $\epsilon_{\text{LUMO}}$ serves as the $y_C$ variable. The properties listed in Equation (2) encompass all relevant characteristics of the molecule, such as $\alpha$, $C_v$, $\epsilon_{\text{HOMO}}$, and others. Our model aims to perform multiple $k$ diffusion processes for each property, with each process conditioned on the context $y_C$ and the property $y_k$.

Final Comment

Thank you so much for your constructive feedback. If you believe we sufficiently addressed your concerns, we would appreciate an increase in your score for this paper.

Thanks for your response. My questions are well-resolved. I will keep my score.