Unifying Generation and Prediction on Graphs with Latent Graph Diffusion

References

[1] C. Vignac, I. Krawczuk, A. Siraudin, B. Wang, V. Cevher, and P. Frossard. DiGress: Discrete denoising diffusion for graph generation. In International Conference on Learning Representations, 2022.

[2] Y. Jang, D. Kim, and S. Ahn. Graph generation with K2-trees. In The Twelfth International Conference on Learning Representations, 2023.

[3] Jaehyeong Jo, Dongki Kim, and Sung Ju Hwang. Graph generation with diffusion mixture., ICML 2024.

[4] Liu, H., Feng, J., Kong, L., Liang, N., Tao, D., Chen, Y., and Zhang, M. One for all: Towards training one graph model for all classification tasks. In The Twelfth International Conference on Learning Representations, 2024

We thank the reviewer for the useful comments. We now address your concerns as follows.

1. Experiments (weakness 1).

• Originally, LGD is designed to be a single framework that can do both generation and regression/classification, where we put our experimental focus on demonstrating that latent diffusion models can achieve state-of-the-art regression/classification results. Now, we have rerun our generation experiments on QM9 and (i) compare our results with state-of-the-art diffusion-based graph generation methods including DiGress [1], HGGT [2] and GruM [3]; (ii) use metrics including FCD and NSPDK to better evaluate the quality of generation. In our new experiments, we train a more powerful encoder/decoder by incorporating positional encodings and more complex reconstruction tasks to obtain better latent graph representations, and carefully choose the hyper-parameters of the diffusion model. The results are shown in Table 1 of the pdf in the global response. It is remarkable that our new results are highly competitive or superior than the baselines in nearly all metrics, which reveals the strong generation ability of LGD. Especially, LGD achieves a better FCD score (in favor of generating graphs similar to the training data) than the three strong recent baselines [1,2,3] while sacrificing the least novelty.

• We are conducting new experiments on the larger dataset MOSES. Since the time is limited, the results are not yet ready. We will update the results as soon as possible, using comments and the link in the global response. We promise to include the new results in our camera-ready paper. It is also notable that most recent studies including HGGT [2] and GruM [3] also do not include these datasets at scale.

2. Clarifications on treating classification and regression tasks as generation tasks (weakness 2). Actually, in Appendix B of our paper, we already provided solid theoretical analysis and clarifications on why diffusion models can perform better than traditional deterministic models in classification and regression tasks. In detail, Theorem B.4. provides error bounds of diffusion models on regression tasks, equation (52) explicitly states the condition where diffusion models outperform regression models, and Corollary B.6. provides guarantees that there exists at least one latent diffusion model whose MAE is not larger than the autoencoder. Combining all these pieces, we can use arbitrary existing prediction models as the autoencoder, and the diffusion model trained on top this autoencoder is able to perform better than this prediction model. To empirically verify this, we utilize the ZINC dataset and train a regression model with our proposed graph attention architecture (the same as the denoising model) using the masked global features (regression target) as the supervising signal. This prediction model has $0.084\pm 0.004$ test MAE on ZINC, while the diffusion model trained upon the autoencoder with the same architecture achieves $0.065\pm0.003$ MAE. This results verifies that we are not complicating simple tasks---instead, diffusion models can provably outperform traditional prediction models.

3. Clarification on training (weakness 3). First of all, LGD is a unified framework (and not necessarily a single model) for all levels of graph learning tasks. As explained in both the limitation section and the experiment section of our paper, we currently train our models separately for each task. However, we are indeed able to train a single model for all tasks using LGD. For example, we can use OFA [4] as the autoencoder, which is a single model that embeds graphs from various domains into a unified latent space through a powerful LLM-GNN mixture architecture. We can then train a diffusion model in this latent space to tackle tasks of all levels and all types, utilizing the unified formulation in our paper. We believe the unified formulation of our paper is a significant contribution and a necessary stage towards graph foundation models. As for the model architectures, details are provided in Appendix D.1 and D.3.1. As a reference, the total number of parameters of LGD (including both the autoencoders and the diffusion model) is 1,825,481 for the regression task on ZINC, and 6,605,545 for the generation task on QM9.

4. Decoding graph structure (weakness 4). As clearly explained in our main text, the latent space contains both the structural information and edge features, since the encoder treats absence of an edge as a special feature. We have a latent representation $W\in \mathbb R^{n\times n\times d}$ for adjacency matrix $A$ (line 145), which is exactly the diffusion target. After obtaining the generated latent edge-augmented representations $\hat W\in \mathbb R^{n\times n\times d}$, we directly apply the pretrained decoder (which is actually a linear layer) to reconstruct both the graph structure ($A$) and edge features. In summary, LGD is a one-shot generation method that generates graph structures and features simultaneously.

5. Figure (weakness 5). Figure 1 is presented in the Appendix due to limited space. We will move the figure to the main context in the camera ready version.

6. Equivariance (question 1). The encoder and decoder are indeed permutation equivariant. We can either use existing models such as MPNN and graph Transformers, or use our specially designed graph attention mechanisms as encoder/decoder. All of the above models are permutation equivariant. Actually, we also use MPNN and graph transformers for the denoising network, making our diffusion model also permutation equivariant (in terms of each denoising step).

7. Source code (question 2). We provide our code through the anonymous link in the global response. We will definitely release all the source code and checkpoints in the final version.

Dear reviewer, While we have addressed all your concerns and answered all your questions, we are looking forward to your reply. As shown in the global response, we rerun the experiments on QM9 and introduce FCD and NSPDK metrics, showing that LGD is competitive with all the strong baselines you mentioned in terms of all metrics. We also conduct experiments on large scale dataset MOSES and show that LGD outperforms DiGress. The answers to your other questions are included as well in our response. While other reviewers have replied positively, we are looking forward to your reply and we are happy to answer any further question you may have. Thank you.

We thank the reviewer for the comments and acknowledgement.

1. Writing (weakness 1). We will refine the writing of the paper and provide as much important information as possible in the main text.

2. Novelty (weakness 2). We believe that our unified formulation is novel and beneficial. Regarding model design, we design novel self-attention and cross-attention mechanisms for graphs with augmented edges, as shown in Appendix D.1. For the diffusion model, we explained in our main text that we use the latent diffusion to tackle the problems brought by the discrete structure, feature misalignment and task variety. It is novel to use a latent diffusion model to perform graph learning tasks of all levels and all types. Furthermore, while architecture and diffusion model design are not our main focus, our LGD framework is extremely flexible, being able to incorporate various model designs and diffusion processes (e.g. DDPM, DDIM, Rectified flow etc.).

3. Source code (weakness 3). We provide the code through the anonymous link in the global response. We will definitely release all the code and checkpoints in the camera-ready version.

4. Scalability of LGD (question 1 and 3). We have explained the scalabiltiy of LGD in multiple places in appendix, such as architecture design in Appendix D.3, theoretical complexity in Appendix E.2, and experimental verifications in Appendix D.2 (line 887-891). We do not have to always model the dense graph where no link is a special type of edge---actually when we do not need to generate structures, we can utilize the sparse graph and only consider the real edges. We can thus use efficient MPNN based models to scale LGD to extremely large graphs such as OGBN-Arxiv, see Table 3 in the paper.

5. Node classification (question 2). In classification and regression tasks, generating the entire graphs is not our main focus---LGD mainly aims to generate the class labels of nodes. However, since LGD actually refines the representation in the latent space, it is still able generate the whole graph. LGD outputs the refined representation, which contains the information of both input (graph structures and known features) and class labels to be predicted, thus it could generate new graphs similar to the training sets.

We sincerely thank the reviewers for the constructive comments and suggestions. We address your concerns as follows.

1. Clarification on the introduction section (weakness 1). The introduction section mainly states the motivation to build a graph foundation model, but we do not claim that we have already done so (even in lines 28-29). We will surely rewrite the paper to avoid misunderstanding. Actually, LGD is indeed a unified framework to handle multiple tasks. Although we do not train a single model currently (as mentioned in our paper), we can indeed build a single graph foundation model under the LGD framework. For example, we can use OFA [1] as the autoencoder, which is a single model that embeds graphs from various domains into a unified latent space through a powerful LLM-GNN mixture architecture. We can then train a diffusion model in this latent space to tackle tasks of all levels and all types, utilizing the unified formulation in our paper. We believe the unified formulation of our paper is a significant step and a necessary stage towards graph foundation models.

2. Clarification on experiments (weakness 2). Appendix C actually describes the potential to build a complete pretrain, finetune and in-context learning framework using LGD. The contents are mainly discussions on future work about LGD, where we have not fully implemented them. We will move Appendix C to the discussion and future work section in the revised version of our paper and make it clearer.

• For your question 2.(1) regarding the conditional generation task on QM9, the encoder is pretrained through the reconstruction loss of nodes and edges as described in the main context, without targets as supervision signals. The diffusion model still generates the molecules, and we follow the prior practice to use another pretrained model to evaluate the properties of all generated molecules, so our method does not introduce any unfairness.

In addition, we would like to emphasize that even the encoders are trained in a supervised manner, it still does not introduce unfairness. Denote the conditions (properties) as $y$ and the molecules as $x$, there is a ground truth distribution $p(x|y)$. In training we sample $x\sim p(x|y)$ and denote it as a probabilistic mapping $f(y)$. The traditional conditional generation methods model $\hat x:=\theta(y)$ with a supervision signal $f(y)$. In comparison, our latent diffusion model learns the distribution of latent representation $\mathcal E(\hat x, \hat y):=\theta(y)$ with a supervision signal $\mathcal E(x, y)=\mathcal E(f(y), y)$. We can observe that both schemes include a supervision signal containing $y$ for generative models, except that we use an encoder to better represent the desired property. This is actually an advantage: utilizing representation learning, latent diffusion can better learns conditional distribution without introducing unfairness.

• For your questions 2.(2), we train a separate encoder/decoder and diffusion model for every dataset for now. However, encoders and diffusion models trained across multiple datasets may have better performance and generalization, which we leave as future work.

3. Clarification on some statements. We thank the reviewer for the advice. There are indeed some statements that can be refined - we will correct them in the revision.

(1) It is true that 3D generation is not necessarily easier. Our recent experiments also find that 3D generation can sometimes be harder due to the difference in data distribution and internal network architectures. We will correct related discussions in the paper. We also acknowledge the practical value of 3D generation and leave it as future work. Regarding the evaluation metrics, we additionally introduce the persuasive FCD and NSPDK metrics in our new experiments for better evaluation, and LGD achieves highly competitive performance. We do not include QED and TPSA for the QM9 dataset, since drug-likeliness is not suitable for the small molecules in QM9.

(2) We are aware of PRODIGY and did not mention it since it works in a slightly different manner compared with the in-context learning in NLP: PRODIGY works by measuring the "similarity" between test and context examples, which is different from the conditional generation as in NLP. However, it is true that PRODIGY is also one sort of ICL, and we will make sure to include this citation and discussion in the revised version.

4. We will add the discussion concerning generalization ability of generative models, including LLM (Llama-3 and GPT-4 etc.) and diffusion models [2-3].

5. Answers to questions. (1) We train diffusion models separately as explained in 2. (2) Feature misalignment is indeed one of the obstacles to build a graph foundation model. LGD provides the potential simultaneously perform generation and regression/classification as long as graphs with different features are mapped to a unified latent space. Actually, OFA [1] is a single model that embeds graphs from various domains into a unified latent space through a powerful LLM-GNN mixture architecture. Combining LGD and OFA, we can then train a generative model in the latent space to overcome the difficulty of feature misalignment and task unification problems, which is left for future work. Overall, we believe that LGD is a significant contribution to the community and could inspire future work on graph foundation models.

[1] Liu, H., Feng, J., Kong, L., Liang, N., Tao, D., Chen, Y., and Zhang, M. One for all: Towards training one graph model for all classification tasks. ICLR, 2024.

[2] Zahra Kadkhodaie, Florentin Guth, Eero P Simoncelli, and Stephane Mallat. Generalization in diffusion models arises from geometry-adaptive harmonic representations. ICLR, 2024.

[3] Puheng Li, Zhong Li, Huishuai Zhang, and Jiang Bian. On the generalization properties of diffusion models. NeurIPS, 2023.

We thank the reviewer for the useful comments and suggestions. As presented in the global response, we have conducted extensive new experiments to improve the paper. Here we address your concerns as follows.

1. General graph generation (weakness 1). As explained in our paper, LGD is capable of generating arbitrary graphs through the powerful continuous latent space, hence is able to handle various types of graph data including generic and synthetic graphs. We conduct generic graph generation experiments and report the results in the global response. We use SBM, Planar and Community-20 datasets for evaluation. LGD performs better than most traditional graph generation methods on these datasets (although it is not always the best). However, other methods often have hard set rules or structural information in the diffusion process, while LGD performs a pure diffusion process in the Euclidean latent space without input structural information, requiring the encoder to be powerful enough to encode structural information. While MPNN and graph transformers are known to have limited expressivity, they may fail to encode enough structural information like degree, clusters, orbits and eigenvalues. Moreover, the autoencoder is trained through reconstruction loss, which has limited supervision on high-order structures. Therefore, if expressive higher-order GNNs or proper structural encodings are used, or more pretraining tasks related to complex structures are adopt, we will be able to obtain a more powerful encoder and thus better generation qualities, which we leave as future work. Please note that LGD has superior performance in real-world tasks where node/edge features are often more important than high-order structures, making it easier to obtain more powerful and suitable encoders for diffusion.

2. Baselines and metrics (weakness 2). We rerun our experiments on QM9 and compare our results with state-of-the-art methods including DiGress [1], HGGT [2] and GruM [3]. We also include the FCD and NSPDK metrics to better evaluate the quality of generation. The results are shown in Table 1 of the pdf in the global response. In our new experiments, we train a more powerful encoder/decoder by incorporating positional encodings and more complex reconstruction tasks to obtain better latent graph representations, and carefully choose the hyper-parameters of the diffusion model. It is remarkable that our new results now are highly competitive or superior than the baselines in nearly all metrics. Especially, it achieves a better FCD score (in favor of generating graphs similar to the training data) than the three strong recent baselines [1,2,3] while sacrificing the least novelty.

3. Larger datasets (weakness 3). We are conducting new experiments on the larger benchmark dataset MOSES. The results are not yet ready due to the limited time. We will update the results as soon as possible, using comments and the link in the global response. We will include the new results in our revised paper. It is also notable that most recent studies including HGGT [2] and GruM [3] also do not include these datasets at scale.

4. Model architecture details (weakness 4). Actually, we have included our model architecture details in the appendix of our submission. Please refer to Appendix D.1 for our architecture design and implementation details, and Appendix D.3.1 for the detailed architectures utilized in each experiment.

5. Open-source implementation (weakness 5). We provide the code through the anonymous link in the global response. We will definitely include the open-source code in the final version of our paper.

6. Computation efficiency (question 1). We have actually included comprehensive analysis on computation efficiency in Appendix E.2 of our submission. In summary, the complexity of LGD depends on the specific architecture of the autoencoder and the denoising network, which is generally $O(n^2)$ for generation tasks and $O(n)$ for regression/classification tasks. The training stage of the diffusion model is fast, and the inference stage can also be accelerated through various methods (such as DDIM, which we already implemented). As a reference, training LGD requires about $0.2$s/epoch on Cora and about $16$s/epoch on ZINC.

[1] C. Vignac, I. Krawczuk, A. Siraudin, B. Wang, V. Cevher, and P. Frossard. DiGress: Discrete denoising diffusion for graph generation. In International Conference on Learning Representations, 2022.

[2] Y. Jang, D. Kim, and S. Ahn. Graph generation with K2-trees. In The Twelfth International Conference on Learning Representations, 2023.

[3] Jaehyeong Jo, Dongki Kim, and Sung Ju Hwang. Graph generation with diffusion mixture., ICML 2024.

Dear reviewers, We now update the results on the large MOSES dataset for molecule generation at scale.

The results show that LGD has superior performance compared with DiGress in terms of validity and novelty metrics. Apart from DiGress and ConGress, we are the only model to scale to this large dataset, bridging the gap between diffusion-based one-shot graph generation model and other traditional methods like SMILES-based, fragment-based, and autoregressive models. Combining with other new results in our global response and the extensive experiments in the original submission, it is convincing that LGD reveals strong capability in both generation and prediction tasks.

In addition, we sincerely appreciate reviewer suUP and a5WK's positive reply. We are looking forward to reviewer XBTL and PVEb's reply. We are happy to answer any further questions regarding the paper. Thank you.