Variational Flow Matching for Graph Generation

Summary

This paper proposed a variant of flow matching for graph generation from a variational perspective, called CatFlow. The main idea is not to estimate the marginal velocity, but instead to estimate the posterior first and then evaluate the marginal velocity as the expectation of the conditional velocity w.r.t. the posterior. In the case of linear conditional velocity and Gaussian variational approximation, the equivalence to vanilla flow matching has been established and connection to score-based models have been discussed. A simplified mean-field variational approach has also been proposed based on the observation that marginal posterior matching (for individual dimensioins) is enough for posterior mean estimation for linear conditional velocity. Experiments on several graph/molecule generation tasks demonstrate the effectiveness of the CatFlow.

Strengths

The paper is written clearly. The variational formulation and the connection to standard flow matching is new. The authors have also shown that a weighted average of the training objective of VFM provides a bound on the log-likelihood of the model.

Weaknesses

The main weakness of the paper is that the core idea of learning the posterior first (i.e., the variational formulation) and evaluating the marginal velocity using the expected conditional velocity w.r.t. the posterior has already been proposed in Dirichlet Flow [1]. Although the author claimed that Dirichlet Flow uses the specific dirichlet distribution as the conditonal probability, this is mainly a design choice which does not affect the generality of Dirichlet Flow for modeling discrete objectives such as graphs and molecules. Given this, the lack of an appropriate comparison to Dirichlet Flow in the experiments represents a significant drawback of the current paper. The paper also lacks a brief introduction to the baseline methods, which would help contextualize the results and better highlight the contributions of the proposed approach. [1] Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.05841, 2024.

Questions

1. Although it is reasonable to match the marginal posterior using mean-field approximation, it is still not clear if this simple mean-field approximation is good enough to capture the posterior mean. Can the author provide some intuition when the mean-field approximation is good enough? 2. In flow matching, an error bound for the generated distribution in term of flow matching error is provided. Can a similar result be provided here in terms of the variational approximation error? 3. How would VFM perform when the conditional velocity is not linear?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations has been adequately addressed.

Summary

This work presents variational flow matching framework VFM and introduces CatFlow, an application to generation of categorical data such as graph. The paper reformulates flow matching using variational perspective and with linear assumption on the conditional vector field achieves tractable objective. The authors derive the connection with the score-based models and show that the framework provides variational bound on the model likelihood.

Strengths

- To the best of my knowledge, the variational framework for flow matching is novel and the provided theoretical works seem solid. - Connection with the score-based model provide intuitions on the method. - Application to discrete data seems reasonable and the benefits over standard flow matching is well explained.

Weaknesses

- The reason CatFlow outperforms previous diffusion models is unclear. Is it because VFM provides variational bound on the model likelihood? Ablation studies on why CatFlow outperforms other diffusion models would strengthen the work. - CatFlow should be validated on larger datasets used in recent works [1, 2], for example Planar and SBM, since community-small and ego-small datasets consists of very small graphs which is not suitable for evaluating generative models. - In particular, validity metric should be used to evaluate the method, instead of only relying on MMD. Evaluating only with MMD is not appropriate as MMD may fail to catch important graph characteristics, for example, small MMD does not guarantee the generated graphs is actually a community graph. Evaluating on Planar or SBM dataset and comparing the results of V.U.N. (valid, unique, and novel) is necessary to strongly argue the advantage of the proposed method. - Comparison with discrete diffusion model like DiGress should be done on more datasets other than QM9. This also leads to experiments on Planar or SBM datasets or in molecular generation tasks, like GuacaMol dataset. Minor correction: citation in line 148 should be fixed

Questions

I would appreciate it if the authors could address the weakness above. Also, I recommend more explanation on Figure 1 as it would privde the readers to get better understanding on CatFlow. In summary, I believe the paper provides theoretical contributions, but the experimental validation needs more room to improve.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have addressed the limitations at the end of the paper.

Summary

This paper introduces a new variational inference framework of flow matching with a focus on applying the framework on discrete data generation. Instead of using the squared norm in standard flow matching, the paper proposes a variational distribution to the conditional path, which is used in the vector field. The paper shows that when the variational distribution is identical to the conditional path, the approximate vector field equals to the target one. The paper then uses mean-field factorisation of the variational distribution and applies linear conditional vector field. Based on the proposed framework, the paper mainly applies it for discrete data generation including the tasks of abstract graph generation and molecular generation. Experiments show that the proposed method outperforms other diffusion-based or flow-matching-based methods.

Strengths

1. The technical contribution and significance of the paper are good. The paper provides a novel angle of formulating flow matching as a variational inference problem, which is theoretically sound and conceptually intuitive. 2. The paper derives a general framework first and then implement it on the task of discrete data generation. The paper also provides comprehensive theoretical analysis on the connections to related methods. 3. The experimental results of the proposed method are convincing. 4. The paper is well-written in general.

Weaknesses

In general, the quality of the paper is good. But I have a few questions on clarity. Please see below.

Questions

but not all of them are weaknesses. 1. In standard flow matching, one needs to minimise the squared norm between the parameterised vector field and the true vector field. That's why one needs to compute the (conditional) vector field. But in variational flow matching, it seems that the final objective is the maximisation of the data likelihood in terms of $\theta$. I wonder how the vector fields come into play in this case. Are they used for computing the variational distribution somehow? 2. Related to the previous question, it would be good to have more details of the parameterisation of $q$ or the implementation of $\theta$. 3. Mean-field approximation of the variational distribution and linear formulation of the conditional vector field are used, more for efficiency consideration. Can the authors provide more discussions on when these approximation and assumption will be the bottleneck of performance? 4. Although the code is provided, it might be good to provide an algorithm of pseudo code in the paper. 5. It would be good to have more explanations of Figure 1.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors adequately addressed the limitations