Improving Graph Generation with Flow Matching and Optimal Transport

Q: Scalability

A: We conducted experiments to evaluate generation times across various graph scales using the Enzymes dataset. The GraphEvo model was configured with six layers, and generation time was measured for a single instance on an NVIDIA A100 80G GPU. The results are detailed in lines 1455-1457 of Appendix J.3 in our latest submission.

Q: Lack of novelty

A: The DeFoG model was released after our submission, so we believe it does not diminish the novelty of our work. Furthermore, our main contributions include the integration of optimal transport into flow matching, the development of a new architecture, and the implementation of goal-directed conditional generation. We are confident that our method will make a contribution to the community.

Q: Missing Explicit References in Method Presentation

A: Thank you for your suggestion. We have added the relevant section in lines 861-870 of Appendix B.2 in our latest submission.

Q: Minior1

A: We have corrected this issue.

Q: Minior 2

A: Due to limited time and computing resources, we were unable to reproduce the training instability of diffusion models. However, several references discuss this phenomenon [1].

[1] Karras, Tero, et al. "Analyzing and improving the training dynamics of diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

We sincerely appreciate the reviewer’s constructive feedback. We have carefully addressed each comment and made the necessary revisions accordingly.

Q: Validity of Theorem 1

A: In our method, we assume that all pairs between the source and target data share the same node order, as the interpolation between them relies on pairs with identical node order. During the training procedure, our network $\hat{p}^1(G^1|G^0,G^t)$ maintains permutation equivariance, such that $\hat{p}^1(G^1|\pi G^0,\pi G^t) = \pi \hat{p}^1(G^1| G^0, G^t)$ for any permutation $\pi$ to approximate $\pi G^1$. Additionally, the interpolation $\pi G^t = t\pi G^0 +(1-t)\pi G^1$ necessitates that the $G^0$ and $G^1$ share identical permutations, thus we prove the invariance of optimal transport under identical permutations, i.e. $\phi(G^0,G^1)=\phi(\pi G^0,\pi G^1)$.

Since the prior distribution is permutation invariant and our GraphEvo model is permutation equivariant, the generated samples also maintain permutation invariance, as established in Lemma 1, which adapts Theorem 5 from CatFlow [1] to our framework in Appendix C.4 of our latest submission. We have added these clarifications in lines 986-992 of Appendix C.4.

Moreover, our model adopts the Hamming distance to balance the efficiency and performance of optimal transport. The results in our paper can demonstrate the utility of our optimal transport and we additionally tested the computation time of optimal transport on different datasets to demonstrate its efficiency. These results are included in lines 1084-1101 of Appendix E in our latest submission.

Q: Results in Enzymes dataset

A: We have included the results for the Enzymes dataset in Table S4 of Appendix H in our latest submission. These results further demonstrate GGFlow's superiority in larger graph datasets.

Q: Overfitting Concerns on Ego Small Dataset

A: We acknowledge that many methods, including ours, may face overfitting risks on the Ego Small dataset due to its simple graph structure and limited data. However, the superior performance of our method on larger graph datasets, such as Enzymes, Planar, and Grid, demonstrates its effectiveness.

Q: Ablation Study

A: In Table 4 of the initial manuscript, GGFlow (w/o GraphEvo) and GGFlow (w/o both), which shares the same architecture, extra features, and hyperparameters as DiGress, demonstrate superior performance, highlighting the advantages of the flow matching framework over the diffusion framework.

Following your suggestion, we conducted an additional ablation study on the Planar dataset. The results indicate that all components contribute to our strong performance, which are included in Table 4 of Section 4.4 in our latest submission.

Additionally, we focus on results without extra features, labeled GGFlow (w/o extra), on the Zinc250k dataset due to the limited time and computational resources, demonstrating that these features enhance performance. DiGress also demonstrated the utility of these features in Table 6 of their paper [2].

[1] Eijkelboom, Floor, et al. "Variational Flow Matching for Graph Generation." arXiv preprint arXiv:2406.04843 (2024).

[2] Vignac, Clement, et al. "DiGress: Discrete Denoising diffusion for graph generation." The Eleventh International Conference on Learning Representations.

Dear Reviewer,

Thank you very much for your thoughtful review. We have addressed the concerns you raised and look forward to your constructive feedback.

Thank you for your reply, and we apologize for any unclear presentation.

We acknowledge that your statements (2-5) are correct. However, in our method settings, we only need to ensure optimal transport invariance under identical permutations, i.e. $\phi (G^0,G^1) = \phi (\pi G^0,\pi G^1)$ for any permutation $\pi$. We would like to clarify the following points:

1. Purpose of Theorem 1: Our Theorem 1 aims to prove the permutation invariance of the optimal transport map for the same permutation. Given our objectives, the theorem holds true.

2. Reasons for Same Permutation: We will explain why our method demonstrates permutation invariance specifically for the same permutation.

Reasons:

The goal of optimal transport is to pair source and target data points with minimal cost during training, which is beneficial for our interpolation [1,2]. Thus, we design our optimal transport approach from the perspective of interpolation. The interpolation of source and target data is defined as $G^t = tG^0 + (1-t)G^1$.

(What is the node order?) It is important to clarify the node order of the graph $G$: If the node set of $G$ is $\\{A,B,C\\}$, the possible node orders include $(A,B,C)$ or $(B,A,C)$ or $(C,B,A)$ and is the order for our nodes and edges in matrix representation.

1. In the interpolation process, we transform the graph representation to a matrix representation before performing interpolation. For example, for source data $G^0=(V^0, E^0), V^0\in \mathbb{R}^{a\times n}, E^0\in \mathbb{R}^{a\times a\times m}$ and target data $G^1=(V^1, E^1), V^1\in \mathbb{R}^{a\times n}, E^1\in \mathbb{R}^{a\times a\times m}$, where $a$ is the number of nodes, $n$ is the class number of nodes, and $m$ is the class number of edges, the node orders of $G^0$ and $G^1$ have been fixed. Therefore, interpolation is performed directly on these fixed node orders.

2. (On what relies this assumption?) The optimal transport aims to find pairs with the minimum cost for interpolation, and interpolation is conducted on a fixed node order. Additionally, during optimal transport calculations, we also utilize the matrix representation of these graphs and our prior distribution is permutation invariant. Therefore, our objective is to match source data with the target data $G^1$ whose node order is fixed, to achieve minimal transport cost. Furthermore, we assume that all pairs of source and target data share the same node order during optimal transport, which also facilitates the identification of pairs with minimal cost.

3. (Why we need to prove Theorem 1?) Regarding the permutation of the intermediate graph $G^t$, we have $\pi G^t = t\pi G^0 + (1-t)\pi G^1$, where $G^0$ and $G^1$ share an identical permutation. Our network $\hat{p}^1(G^1|G^0,G^t)$ needs to maintain permutation equivariance, such that $\hat{p}^1(G^1|\pi G^0,\pi G^t) = \pi \hat{p}^1(G^1| G^0, G^t)$ for any permutation $\pi$ to approximate $\pi G^1$. So we prove the invariance of optimal transport under identical permutations, i.e. $\phi(G^0,G^1)=\phi(\pi G^0,\pi G^1)$.

4. (What is Theorem 1) Considering the invariance of Hamming distance to the same permutation, i.e. $H(G^0, G^1) = H(\pi G^0, \pi G^1)$ for any permutation $\pi$, our optimal transport is permutation invariance to same permutation i.e. $\phi(G^0,G^1)=\phi(\pi G^0,\pi G^1)$.

(Why Choose This Method Setting?) Moreover, we have tested the cost with invariance to any permutation in our initial experiments. However, computational methods such as the GOT Distance [3] and FGW distance [4] have prohibitive computational costs. Therefore, we opted to utilize the Hamming distance as our cost metric, enabling us to achieve our objectives with acceptable computational efficiency.

In summary, Theorem 1 is designed to ensure the invariance of optimal transport under identical permutations. According to our experimental results, this invariance benefits our training process and generation results.

We hope this explanation clarifies Theorem 1 and addresses your concerns.

[1] Bose, Avishek Joey, et al. "SE(3)-stochastic flow matching for protein backbone generation." arXiv preprint arXiv:2310.02391 (2023).

[2] Song, Yuxuan, et al. "Equivariant flow matching with hybrid probability transport for 3d molecule generation." Advances in Neural Information Processing Systems 36 (2024).

[3] Chen, Liqun, et al. "Graph optimal transport for cross-domain alignment." International Conference on Machine Learning. PMLR, 2020.

[4] Titouan, Vayer, et al. "Optimal transport for structured data with application on graphs." International Conference on Machine Learning. PMLR, 2019.

'Theorem 1: If the distributions $p(G_0)$ and $p(G_1)$ are permutation invariant and the cost func- tion maintains this invariance, then the optimal transport map ϕ also respects this property, i.e., $ϕ(G_0, G_1) = ϕ(πG_0, πG_1)$, where π is a permutation operator.'

In the graph learning community, permutation-invariance refers to invariance under any permutation of the input node ordering.

Regarding Theorem 1, the following statements seem implied:

1. The cost function is permutation-invariant with respect to any permutation of the input node ordering.
2. Consequently, the optimal transport map is also permutation-invariant with respect to any permutation of the input node ordering.
3. An optimal transport map $\phi$ is permutation-invariant with respect to any permutation of the input node ordering, if $\phi(G_0, G_1) = \phi(\pi G_0, \pi G_1)$, where $\pi$ is a permutation.

However, as acknowledged, these three statements are incorrect. The current formulation is, at the very least, misleading.

Additionally, the confusion is compounded by the preceding statement: "To optimize the computational efficiency of OT and preserve permutation invariance, we define the distance..."

Theorem 1 should be reformulated to eliminate confusion, as it is incorrect to state that the cost function and the transport map are permutation-invariant (implicitly understood with respect to any permutation of the input node ordering).

Thank you for your patient and constructive feedback. We have made the necessary revisions to Theorem 1 and its clarifications in accordance with your suggestions to avoid any misleading information. Theorem 1 is now presented as follows:

If the distributions $p(G^0)$ and $p(G^1)$ are permutation invariant, and the cost function maintains invariance under identical permutations, i.e., $H(G^0, G^1) = H(\pi G^0, \pi G^1)$ for any permutation $\pi$, then the optimal transport map $\phi$ also exhibits invariance under identical permutations, such that $\phi(G^0, G^1) = \phi(\pi G^0, \pi G^1)$.

Additionally, we have revised the sentences in lines 212 and 275-282 of our latest submission to avoid any potential confusion for readers. We have also added clarifications from our discussion regarding our method settings in lines 989-1024 of Appendix C.4 to provide a more detailed explanation of the rationale behind Theorem 1. We also hope these clarifications will also inspire other researchers in the field. All of these revisions are marked in red in our latest submission.

We appreciate your valuable feedback and suggestions, which have significantly improved the quality of our manuscript. We sincerely hope that our new submission and revisions have addressed your concerns.

We sincerely appreciate the reviewer’s constructive feedback. We have carefully addressed each comment and made the necessary revisions accordingly.

Q: Ablation studies of Planar dataset

A: Following your suggestion, we conducted an additional ablation study on the Planar dataset. The results indicate that all components contribute to our strong performance, which have been included in Table 4 of Section 4.4 in our latest submission.

Q: Training stability

We provide training loss and average values on community small datasets, comparing our method with DiGress, which shares the same training objectives. For fair comparisons, we analyze GGFlow (w/o both) and GGFlow (w/o Evo) to highlight the advantages of the flow matching framework and optimal transport. These results are detailed in our submission (lines 1447-1457) and illustrated in Figure S3 of Appendix J.2.

Our results in Figure S3 indicate that GGFlow (w/o both) and GGFlow (w/o Evo) achieve faster and more effective convergence than DiGress using the same architecture. Furthermore, GGFlow (w/o Evo) surpasses both GGFlow (w/o both) and DiGress in average metrics, demonstrating the benefits of flow matching and optimal transport. These results suggest enhanced stability during training.

Dear Reviewer,

Thank you very much for your thoughtful review. We have addressed the concerns you raised and look forward to your constructive feedback.

Q: Ablation studies of optimal transport

A: The y-axis in Figure 3 has been scaled by DiGress and GDSS, making it difficult to observe the difference between GGFlow (w/o OT) and GGFlow. To address this, we have included an additional figure in Figure S2 (a) of Appendix J.2, which focuses solely on GGFlow and GGFlow (w/o OT). This figure demonstrates that GGFlow consistently achieves better average performance and narrower confidence intervals than GGFlow (w/o OT) across nearly all inference steps.

Additionally, we conducted experiments on the Planar dataset to further illustrate the sampling efficiency of optimal transport in Figure S2 (b) of Appendix J.2. The results show that GGFlow outperforms GGFlow (w/o OT) in V.U.N. performance at all inference steps and can achieve comparable performance with fewer steps, reinforcing the superiority of our optimal transport approach.

We sincerely appreciate the reviewer’s constructive feedback. We have carefully addressed each comment and made the necessary revisions accordingly.

Q: Comparison with CatFlow

A: Since CatFlow [1] has not made its source code available, we did not include it in our experiments for fairness. Instead, we clarified the differences between our methods in Appendix B.3.

The tables below present results directly from their manuscript, where our methods demonstrate superior performance.

Three reasons contribute to our better performance:

1. Model Architecture: CatFlow employs a graph transformer similar to DiGress, whereas our method utilizes an edge-augmented graph transformer to improve edge communication.
2. Prior Distribution: CatFlow uses a standard normal distribution as its prior, while our reference distribution is specifically designed to closely approximate the true data distribution, addressing graph sparsity.
3. Source and Target Coupling: CatFlow implements independent source and target coupling, while our method employs optimal transport to extend coupling to more general scenarios. This approach results in straighter generative trajectories, thereby enhancing performance.

Q: Results of Planar dataset and validity of Grid dataset

A: We have included the Planar dataset in our initial manuscript, presenting the V.U.N. results in Table S2 of Appendix G. These results demonstrate the superiority of our models compared to other methods, including the V.U.N. metric.

Previous works [2, 3, 4, 5] did not test validity for the Grid dataset, so we didn't report the metric in our initial submission. Following your suggestion, we calculated the validity metric for the Grid dataset using the method from the DiGress GitHub repository (https://github.com/cvignac/DiGress/blob/9f526a05c8c55715d5c92b6a35d32ed758a9098b/src/analysis/spectre_utils.py#L581). The results, shown in the following table, indicate that our method achieves comparable validity than other approaches. However, the rationale for the validity calculation method remains unproven, as it relies on searching for isomorphic graphs across all grid data, including the training dataset, rather than checking features (such as planarity) of the generated graph to assess validity, as is done with the Planar dataset.

Q: Ablation studies in Planar dataset

A: Following your suggestion, we conducted an additional ablation study on the Planar dataset and reported the V.U.N. metric. The results indicate that all components contribute to our strong performance, which have been included in Table 4 of Section 4.4 in our latest submission.

Q: Results of molecule datasets using the same architecture

A: Given the limited time and computational resources, we focused on results from the Zinc250k dataset, employing the same model architecture as DiGress. In the table, GGFlow (w/o both) outperforms DiGress across nearly all metrics, demonstrating the superiority of our flow matching framework.

[1] Eijkelboom, Floor, et al. "Variational Flow Matching for Graph Generation." arXiv preprint arXiv:2406.04843 (2024).

[2] Jo, Jaehyeong, Seul Lee, and Sung Ju Hwang. "Score-based generative modeling of graphs via the system of stochastic differential equations." International conference on machine learning. PMLR, 2022.

[3] You, Jiaxuan, et al. "Graphrnn: Generating realistic graphs with deep auto-regressive models." International conference on machine learning. PMLR, 2018.

[4] Yan, Qi, et al. "Swingnn: Rethinking permutation invariance in diffusion models for graph generation." arXiv preprint arXiv:2307.01646 (2023).

[5] Luo, Tianze, Zhanfeng Mo, and Sinno Jialin Pan. "Fast graph generation via spectral diffusion." IEEE Transactions on Pattern Analysis and Machine Intelligence (2023).

Dear Reviewer,

Thank you very much for your thoughtful review. We have addressed the concerns you raised and look forward to your constructive feedback.

We sincerely appreciate the reviewer’s constructive feedback. We have carefully addressed each comment and made the necessary revisions accordingly.

Q: Reward Function Definition

A: Thank you for your feedback. We acknowledge the importance of the reward function in training. In our initial submission, we used $|\mu - \mu^*|$ as the reward, aligning with our evaluation metric (MAE). Objectives relevant to the metric can serve as effective rewards to achieve the desired outcomes.

In response to your suggestion, we also implemented $(\mu - \mu^*)^2$ as a reward function. Our experiments showed that this alternative improved the guidance effect, yielding results comparable to our initial approach.

In many application scenarios, using evaluation metrics or related functions as reward signals can effectively achieve the desired objectives [1, 2, 3]. Our findings align with existing literature.

Q: Training Stability of GGFlow

A: We provide training loss and average values on Community-small datasets, comparing our method with DiGress, which shares the same training objectives. For fair comparisons, we analyze GGFlow (w/o both) and GGFlow (w/o Evo) to highlight the advantages of the flow matching framework and optimal transport. These results are detailed in our latest submission (lines 1426-1434) and illustrated in Figure S3 of Appendix J.2.

Our results in Figure S3 indicate that GGFlow (w/o both) and GGFlow (w/o Evo) achieve faster and more effective convergence than DiGress using the same architecture. Furthermore, GGFlow (w/o Evo) surpasses both GGFlow (w/o both) and DiGress in average metrics, demonstrating the benefits of flow matching and optimal transport. These results suggest enhanced stability during training.

Additionally, Figure 3 of our initial submission presents sampling efficiency data, showing that GGFlow has narrower confidence intervals than DiGress and GDSS, indicating improved sampling stability. Our method achieves superior performance with fewer inference steps, further demonstrating its sampling efficiency.

Q: Optimal Transport Justification

A: As noted in FM [4], the source and target coupling can be extended to more general cases, and we choose the optimal transport as our coupling. Optimal transport enables our model to pair the nearest source and target data points, ensuring that flow trajectories are optimal [5], which is particularly applicable to graph-structured data. This results in more direct generative paths, enhancing the overall performance of our model.

Q: Complexity Analysis of Optimal Transport

A: During inference, we do not need to calculate the optimal transport map as outlined in Equation 6 of our manuscript.

We compared the training times of our method with DiGress using identical architectures on an NVIDIA A100 80G GPU. Considering model size, batch size, and number of nodes, we measured the duration of single training steps across three datasets. Our results show that the time required for optimal transport accounts for only 5% of the total training time, demonstrating its efficiency. This complexity analysis is detailed in our latest submission, specifically in lines 1084-1101 of Appendix E.

[1]Liu, Yijing, et al. "Graph diffusion policy optimization." arXiv preprint arXiv:2402.16302 (2024).

[2] Fan, Ying, et al. "Reinforcement learning for fine-tuning text-to-image diffusion models." Advances in Neural Information Processing Systems 36 (2024).

[3] Mnih, Volodymyr, et al. "Human-level control through deep reinforcement learning." nature 518.7540 (2015): 529-533.

[4] Gat, Itai, et al. "Discrete Flow Matching." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

[5] Liu, Xingchao, and Chengyue Gong. "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow." The Eleventh International Conference on Learning Representations.

Dear Reviewer, Thank you very much for your thoughtful review. We have addressed the concerns you raised and look forward to your constructive feedback.

We will describe our model results from model performance, inference efficiency and training stability to demonstrate that flow matching and optimal transport contribute to our results. Furthermore, we will discuss the issues related to the extra features.

Community Small Dataset

Planar Dataset

Firstly, we provided detailed ablation studies on both the Community Small and Planar datasets. In the settings of our ablation experiments, GGFlow (w/o Evo) and GGFlow (w/o both) share the same architecture which incorporates extra features, and the same hyperparameter settings as DiGress [1]. GGFlow (w/o both) adopts only the flow matching framework without optimal transport, and GGFlow (w/o Evo) incorporates both the flow matching framework and optimal transport.

The results in Table 4 of our submission demonstrate that GGFlow (w/o both) outperforms DiGress across most metrics, highlighting the advantages of the flow matching framework over the diffusion framework in graph space. Moreover, GGFlow (w/o Evo) also outperforms GGFlow (w/o both) in nearly all metrics, indicating the effectiveness of optimal transport. Additionally, according to Figure S2 in Appendix J.2, we also observed that optimal transport enhances the inference efficiency of GGFlow.

Then, we also provided training analysis of using optimal transport in Appendix J.2. According to Figure S3, GGFlow (w/o both) and GGFlow (w/o Evo) show better training convergence and stability than DiGress. Moreover, GGFlow (w/o Evo) demonstrates better and more stable generation results than GGFlow (w/o both) during the training procedure, indicating that optimal transport significantly improves the training stability and efficiency of our model.

Finally, regarding the extra features in our model, we conducted experiments on the Zinc250k dataset in our first rebuttal to demonstrate their utility. It is noted that these extra features are not a primary contribution of our models; we clarified in our initial submission that they come from DiGress [1] and also cited DiGress in line 247 of our initial submission. Therefore, we did not provide ablation studies of these extra features in our intial submission. The original DiGress paper provides ablation studies of these features, which can be found in Table 6 of their paper.

In summary, we conducted comprehensive ablation studies to demonstrate the effectiveness of our proposed flow matching method and optimal transport. We think these results are sufficient to support the effectiveness of our method and the contributions of our work.

Thank you for your attention and reading.

[1] Vignac, Clement, et al. "DiGress: Discrete Denoising diffusion for graph generation." The Eleventh International Conference on Learning Representations.