Variational Flow Matching for Graph Generation

Authors: floor  eijkelboom
grigory  bartosh
christian a naesseth
max  welling
jan-willem  van de meent
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:22:30

Abstract: We present a formulation of flow matching as variational inference, which we refer to as variational flow matching (VFM). We use this formulation to develop CatFlow, a flow matching method for categorical data that is easy to implement, computationally efficient, and achieves strong results on graph generation tasks. In VFM, the objective is to approximate the posterior probability path, which is a distribution over possible end points of a trajectory. VFM admits both the original flow matching objective and the CatFlow objective as special cases. We also relate VFM to score-based models, in which the dynamics are stochastic rather than deterministic, and derive a bound on the model likelihood based on a reweighted VFM objective. We evaluate CatFlow on one abstract graph generation task and two molecular generation tasks. In all cases, CatFlow exceeds or matches performance of the current state-of-the-art models.

Main paper: 
 introduction: 
 In recent years, the field of generative modeling has seen notable advancements. In image generation [41, 45] , the development and refinement of diffusion-based approaches -specifically those using denoising score matching [59] -have proven effective for generation at scale [15, 54] . However, while training can be done effectively, the constrained space of sampling probability paths in a diffusion requires tailored techniques to work [52, 64] . This is in contrast to more flexible approaches such as continuous normalizing flows (CNFs) [8] , that are able to learn a more general set of probability paths than diffusion models [53] , at the expense of being expensive to train as they require one to solve an ODE during each training step (see e.g. [5, 46, 14] ).Recently, Lipman and collaborators [27] proposed flow matching (FM), an efficient and simulationfree approach to training CNFs. Concretely, they use a per-sample interpolation to derive a simpler objective for learning the marginal vector field that generates a desired probability path in a CNF. This formulation provides equivalent gradients without explicit knowledge of the (generally intractable) marginal vector field. This work has been extended to different geometries [7, 23] and various applications [60, 9, 13, 24] . Similar work has been proposed concurrently in [32, 1] . This paper identifies a reformulation of flow matching that we refer to as variational flow matching (VFM). In flow matching, the vector field at any point can be understood as the expected continuation toward the data distribution. In VFM, we explicitly parameterize the learned vector field as an expectation relative to a variational distribution. The objective of VFM is then to minimize the Kullback-Leibler (KL) divergence between the posterior probability path, i.e. the distribution over possible end points (continuations) at a particular point in the space, and the variational approximation.We show that VFM recovers the original flow matching objective when the variational approximation is Gaussian and the conditional vector field is linear in the data. Under the assumption of linearity, a solution to the VFM problem is also exact whenever the variational approximation matches the marginals of the posterior probability path, which means that we can employ a fully-factorized variational approximation without loss of generality.While VFM provides a general formulation, our primary interest in this paper is its application to graph generation, where the data are categorical. This setting leads to a simple method that we refer to as CatFlow, in which the objective reduces to training a classifier over end points on a per-component basis. We apply CatFlow to a set of graph generation tasks, both for abstract graphs [35] and molecular generation [40, 19] . By all metrics, our results match or substantially exceed those obtained by existing methods.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: Our paper follows the structure of the abstract and introduction directly.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
background 2.1 transport framework for generative modeling and cnfs: 
 Common generative modeling approaches such as normalizing flows [42, 38] and diffusion models [15, 54] parameterize a transformation φ from some initial tractable probability density p 0 -typically a standard Gaussian distribution -to the target data density p 1 . In general, there is a trade-off between allowing φ to be expressive enough to model the complex transformation while ensuring that the determinant term is still tractable. One such transformation is a continuous normalizing flow (CNF).Any time-dependent 1 vector field v t : [0, 1] × R D → R D gives rise to such a transformation -called a flow -as such a field induces a time-dependent diffeomorphism φ t : [0, 1] × R D → R D defined by the following ordinary differential equation (ODE):EQUATIONIn CNFs, this vector field is learned using a neural network v θ t . Through the change of variables formula p t (x) can be evaluated (see appendix E) and hence one could try and optimize the empirical divergence between the resulting distribution p 1 and target distribution. However, obtaining a gradient sample for the loss requires ones to solve the ODE induced during training, making this approach computationally expensive.In flow matching [27] , the aim is to regress the underlying vector field of a CNF directly on the interval t ∈ [0, 1]. Flow matching leverages the fact that even though we do not have access to the actual underlying vector field -which we denote as u t -and probability path p t , one can construct a per-example formulation by defining conditional flows, i.e. the trajectories towards specific datapoints x 1 . Concretely, FM sets:EQUATIONwhere u t (x | x 1 ) is the conditional trajectory. The most common way to define u t (x | x 1 ) is as the straight line continuation from x to x 1 , implying one can obtain samples x ∼ p t (x | x 1 ) simply by interpolating samples x 0 ∼ p 0 for some p 0 and x 1 ∼ p 1 ,EQUATIONCrucially, the flow matching objectiveEQUATION(4) 1 Time dependence is denoted through the subscript t throughout this paper.is equivalent in expectation (up to a constant) to the conditional flow matching objectiveEQUATIONAn advantage of flow matching is that this conditional trajectory u t (x | x 1 ) can be chosen to make the problem tractable. The authors show that diffusion models can be instantiated as flow matching with a specific conditional trajectory, but also show that assuming a simple, straight-line trajectory leads to more efficient training. Note that in contrast to likelihood-based training of CNFs, flow matching is simulation free, leading to a scalable approach to learning CNFs.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We added a section addressing limitations in section 7.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
variational flow matching for graph generation: 
 We derive CatFlow through a novel, variational view on flow matching we call Variational Flow Matching (VFM). The VFM framework relies on two insights. First, we can define the marginal vector field and its approximation in terms of an expectation with respect to a distribution over end points of the transformation. This implies that we can map a flow matching problem onto a variational counterpart. Second, under typical assumptions on the forward process, we can decompose the expected conditional vector field into components for individual variables which can be computed in terms of the marginals of the distribution over end points of the conditional trajectories. This implies that, without loss of generality, we can solve a VFM problem using a fully-factorized variational approximation, providing a tractable approximate vector field. For categorical data, the corresponding vector field can be computed efficiently via direct summation. This results in a closed-form objective to train CNFs for categorical data, which we refer to as CatFlow. We develop the theory of VFM in section 3 and we relate VFM to flow matching and score-based diffusion in section 4.In any flow matching problem, the vector field in eq. ( 2) can be expressed as an expectationEQUATIONwhere p t (x 1 | x) is the posterior probability path, the distribution over possible end points x 1 of paths passing through x at time t,EQUATIONThis makes intuitive sense: the velocity in point x is given by all the continuations from x to final points x 1 , weighted by how likely that final point is given that we are at x. Note that to compute u t (x), one has to evaluate a joint integral over D dimensions.This observation leads us to propose a change in parameterization of the learned vector field. Rather than predicting the components of the vector field directly, we can define an approximate vector field in terms of an expectation with respect to a variational distribution q θ t with parameters θ,EQUATIONClearly, in this construction v θ t (x) will be equal to u t (x) when q θ t (x 1 | x) and p t (x 1 | x) are identical. This implies that we can map a flow matching problem onto a variational inference problem.Concretely, we can define a variational flow matching problem by minimizing the Kullback-Leibler (KL) divergence from p t to q θ t , which we can express asEQUATIONwhere t ∼ Uniform(0, 1) and x, x 1 ∼ p t (x, x 1 ) (see appendix A.1 for derivations). This leads us to propose the variational flow matching (VFM) objectiveEQUATIONWhile this variational formulation of flow matching is promising, two potential drawbacks emerge in practical applications. First, although it is feasible to reformulate any flow matching problem as a variational inference task, doing so requires learning an approximation of a potentially complex, high-dimensional distribution p t (x 1 | x) (including any correlations between the different x d 1 ). Second, representing v θ t (x) as an expectation may pose intractability challenges. Interestingly, under typical assumptions about the conditional velocity field in flow matching, this objective simplifies significantly, making it no more computationally demanding than standard flow matching, a point we will further address in section 3.2.Decomposing the conditional vector field. At first glance, we do not seem to obtain much from this variational view due to the intractability of p t (x 1 | x) and v θ t (x). Fortunately, we can simplify the objective and the calculation of the marginal vector field under the typical case where the conditional vector field u t (x | x 1 ) is linear in x 1 , such as in straight line interpolations commonly used in flow matching. This leads to two simplifications. First, we notice that the expected value of a linear conditional velocity field simply equals the conditional velocity field towards the expectation/mean of the distribution we are considering, i.e.EQUATIONThis means that as long as our variational distribution has the same mean as p t (x 1 | x), we will learn the same underlying field. Second, we realize that the expected value of x d 1 only depends on p t (x d 1 | x), and thus as such as long as q θ t (x 1 | x) has the same marginal expectations E x d 1 | x , we will learn the same field as under p t (x 1 | x). In other words, it suffices to learn a fully-factorized approximation q θ t (x 1 | x), there is no need to fully characterize the covariance of the posterior probability path. Thich allows us to reduce a high-dimensional variational problem to a series of low dimension problems.Formally, the following holds: Theorem 1. Assume that the conditional vector field u t (x | x 1 ) is linear in x 1 . Then, for any distribution r t (x 1 | x) such that the marginal distributions coincide with those of p t (x 1 | x), the corresponding expectations of u t (x | x 1 ) are equal, i.e.EQUATION)We provide a proof in appendix A.3.It follows directly from theorem 1 that without loss of generality we can consider the considerably easier task of a fully-factorized approximationEQUATIONWe refer to this special case as mean-field variational flow matching (MF-VFM), and the VFM objective reduces toEQUATIONEven though we use a factorized model, it is worth emphasizing that due to the linearity of the conditional field, a mean-field variational distribution can learn the solution exactly.Computing the marginal vector field. To calculate the vector field v θ t (x), we can simply substitute the factorized distribution q θ t (x 1 | x) into eq. ( 8). However, this still requires an evaluation of an expectation. Fortunately, leveraging the linearity condition significantly simplifies this computation, since as long as we have access to the first moment of one-dimensional distributions q θ t (x d 1 | x), we can efficiently calculate v θ t (x) by simply considering the conditional field towards it. Note that for many distributions -e.g. Gaussian -learning its parameters is equivalent to learning its expected value. Note that the training procedure will differ for two distinct distributions -e.g. Gaussian versus Categorical -so the form of the distribution q θ t (x 1 | x) remains practically important, a flexibility provided through the variational view on flow matching. Figure 1 : Parameterization of the vector field in CatFlow. Given an interpolant x t = tx 1 + (1 -t)x 0 , CatFlow predicts a categorical distribution q θ t (x 1 | x t ) parameterized by a vector µ t (x t ). The resulting construction for the vector field v θ t (x t ) = (µ t (x t ) -x t )/(1 -t) ensures that trajectories converge to a point on the simplex at t = 1.If we now use the standard flow matching case of using a conditional vector field based on a linear interpolation, the approximate vector field can be expressed in terms of the first moment of the variational approximation:EQUATIONNote that this covers both the case of categorical data, which we focus on in this paper, and the case of continuous data, as considered in traditional flow matching methods. We provided the general algorithm for training and generation in appendix B.1.At first glance, the linearity condition of the conditional vector field u t (x | x 1 ) in theorem 1 might seem restricting. However, in most state-of-the-art generative modeling techniques, this condition is satisfied, e.g. diffusion-based models, such as flow matching [27, 2, 32] , diffusion models [15, 54] , and models that combine the injection of Gaussian noise with blurring [44, 16] , among others [10, 51] .The CatFlow Objective In CatFlow, we directly apply the VFM framework to the categorical case. Let our parameterised variational distribution q θ t (xd 1 | x) = Cat(x d 1 | θ d t (x)), and let us denote the parameters of this categorical asEQUATIONThen, the dth component of the learned vector field isEQUATIONIntuitively, CatFlow learns a distribution over the conditional trajectories to all corners of the probability simplices, rather than regressing towards an expected conditional trajectory.In the categorical setting, the MF-VFM objective can be written out explicitly. Writing out the probability mass function of the categorical distribution, we see thatEQUATIONAs such, we find that CatFlow objective is given by a standard cross-entropy loss:EQUATIONNote, however, that when actually computing v θ t , this can be done efficiently, sinceEQUATION) since µ d (x) := E q θ t (x d 1 |x) [u d t (x d 1 |x)] and the other terms are not in the expectation. Note that this geometrically corresponds to learning the mapping to a point in the probability simplex, and then flowing towards that. This procedure is illustrated in fig. 1 . Because of this, training CatFlow is no less efficient than flow matching. We provided the algorithm for training and generation for CatFlow and a Gaussian VFM objective in appendix B.2 and appendix B.3 respectively. More precisely, training CatFlow offers two key benefits over standard flow matching. First, given that CatFlow predicts points in the probability simplex and parametrizes the velocities to point into it, an inductive bias is introduced ensuring all generative paths align with realistic trajectories, hence avoiding misaligned paths. Second, using a cross-entropy loss instead of a mean-squared error improves gradient behavior during training. Both aspects enhance learning dynamics and speed up convergence, which is evaluated in section 6.3. Lastly, CatFlow's ability to learn probability vectors -rather than directly choosing classes as is common in discrete approaches -allows the model to express uncertainty about variables at a specific time. This is especially useful in complex domains like molecular generation, where initial uncertainty about components decreases as more structure is established, leading to more precise predictions.Permutation Equivariance. Graphs, defined by vertices and edges, lack a natural vertex order unlike other data types. This permutation invariance means any vertex labeling represents the same graph if the connections remain unchanged. Note that even though the following results apply to graphs, an unordered set of categorical variables can be described by a graph without edges. Under natural conditions -see appendix A -we ensure CatFlow adheres to this symmetry (see appendix A.4 for the proof).Theorem 2. CatFlow generates exchangeable distributions, i.e. CatFlow generates all graph permutations with equal probability.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We provide all code to the experiments. Moreover, all data is publically available, and will be provided in the repository, including data processing.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
flow matching and score matching: bridging the gap: 
 In this section, we relate the VFM framework to existing generative modeling approaches. First, we show that VFM has standard flow matching as a special case when the variational approximation is Gaussian. This implies that VFM provides a more general approach to learning CNFs. Second, we show that through VFM, we are not only able to compute the target vector field, but also the score function as used in score-based diffusion. This has two primary theoretical implications: 1) VFM simultaneously learns deterministic and stochastic dynamics -as diffusion models rely on stochastic dynamics, and 2) VFM provides a variational bound on the model likelihood.Relationship to Flow Matching. VFM admits FM as a special case, under certain assumptions on u t (x | x 1 ), when the variational approximation is Gaussian. Formally, the following holds (see theorem 3 in appendix A.2 for the proof): Theorem 3. Assume the conditional vector field u t (x | x 1 ) is linear in x 1 and is of the formEQUATIONwhereA t (x) : [0, 1]×R D → R D ×R D and b t (x) : [0, 1]×R D → R D . Moreover, assume that A t (x) is an invertible matrix and q θ t (x 1 | x) = N (x 1 | µ θ t (x), Σ t (x)), where Σ t (x) = 1 2 (A ⊤ t (x)A t (x)) -1. Then, VFM reduces to flow matching.Relationship to Score-Based Models. Flow matching [27] is inspired by score-based models [54] and shares strong connections with them. This leads us to two observations. The first is that it should be possible to define a variational parameterization of the score function in diffusion models that is analogous to the one in VFM. The second is that we can build on existing results that show that many diffusion model objectives, including the standard flow matching objective with a linear interpolant, can be expressed as special cases of a general weighted loss function [22, 21] . Here we simililarly define a bound on the log-likelihood in terms of a reweighted VFM objective.In score-based models, the objective is to approximate the score function ∇ x log p t (x) with a function s θ t (x). A connection to VFM becomes apparent by observing that the score function can also be expressed as an expectation with respect to p t (x 1 | x) (see appendix A.5 for derivation):EQUATIONwhere ∇ x log p t (x | x 1 ) is the tractable conditional score function. Similarly, we can parameterize s θ t (x) in terms of an expectation with respect to a variational approximationEQUATIONIt is now clear thats θ t (x) = ∇ x log p t (x) when q θ t (x 1 | x) = p t (x 1 | x). This suggests that there exists a variational formulation of score-based models that is entirely analogous to VFM. Indeed, existing work on continuous diffusion for categorical data [11] defines a parameterization of the score function of this form (see section 5 for a more detailed discussion).Following [54, 3] , we can construct stochastic generative dynamics dx = ṽθ t (x)dt + g t dw to approximate the true dynamics dx = ũt (x)dt + g t dw (see details in appendix A.6), withEQUATION)Here g t : [0, 1] → R + is a scalar function, and w is a standard Wiener process.This connection has two important implications. First, it shows that learning a variational approximation q θ t (x 1 | x) can be used to define both deterministic and stochastic dynamics, whereas flow matching typically considers deterministic dynamics only (as flows are viewed through the lens of ODEs). Second, it enables us to show that a reweighted version of the VFM objective provides a bound on the log-likelihood of the model. This result, inspired by [21] , provides another theoretical motivation for learning using the VFM objective. Theorem 4. Rewrite the Variational Flow Matching objective as follows:L VFM (θ) = E t,x L θ (t, x)whereEQUATIONThen, the following holds:EQUATIONwhere λ t (x) is a non-negative function and C is a constant.We provide a proof in appendix A.7. We further note that that applying the same linearity condition that we discussed in section 3.2 to the conditional score function maintains all the same connections with score-based models. 
related work: 
 Diffusion Models for Discrete Data. Several approaches to diffusion models have been developed for graph generation. In [58] , the authors define a Markov process that progressively edits graphs by adding or removing edges and altering node or edge categories and is trained using a graph transformer network, that reverses this process to predict the original graph structure from its noisy version. This approach breaks down the complex task of graph distribution learning into simpler node and edge classification tasks. Moreover, [20] proposes a score-based generative model for graph generation using a system of stochastic differential equations (SDEs). The model effectively captures the complex dependencies between graph nodes and edges by diffusing both node features and adjacency matrices through continuous-time processes. These are non-autoregressive graph generation approaches that perform on par with autoregressive ones, such as in [31, 26, 36] . Other non-autoregressive approaches worth mentioning are [28, 34, 33] .There is also work on diffusion-based models for discrete data in the form of text and other sequential data [4, 33, 61, 18] . Though not explicitly formulated in terms of a variational perspective, the work on continuous diffusion for categorical data [11] arrives at a an approach that is closely related to that of CatFlow. This work defines a diffusion in the embedding space of a transformer-based language 23), which leads to an expression for the learned score function in terms of a mean embedding, analogous to the one we obtain in eq. (15) . Where this approach differs from CatFlow, other than in that it defines a flow in the embedding space of language models, is in how the objective is defined. The authors also minimize a cross-entropy loss, but employ a time-warped objective, similar to the general weighted objective proposed in [22] , where the warping is optimized to ensure that the entropy of predictions decreases linearly when moving from noise to data in uniform time.Flow-based methods for Discrete Data. Recently, two flow-based methods for discrete generative modeling have been proposed, which differ both in terms of technical approach and intended use case from the work that we present here. 2In [56] , a Dirichlet flow framework for DNA sequence design is introduced, utilizing a transport problem defined over the probability simplex, similar to diffusion on simplices proposed in [43] . This approach differs from CatFlow in that it represents the conditional probability path p t (x | x 1 ) using a Dirichlet distribution. This implies that points x are constrained to the simplex, which is not the case for CatFlow. Dirichlet Flows have not been evaluated on graph generation, but we did carry out preliminary experiments based on the released source code. We compare to this approach in our experiments.In [6] , Discrete Flow Models (DFMs) are introduced. DFMs use Continuous-Time Markov Chains to enable flexible and dynamic sampling in multimodal generative modeling of both continuous and discrete data. Though sharing a goal, this approach differs significantly from CatFlow as in the end the resulting model does not learn a CNF, but rather generation through sequential sampling from a time-dependent categorical distribution. As in the case of Dirichlet flows, no evaluation on graph generation was performed.The switch to the variational perspective is inspired by [58] , showing significant improvement through viewing the dynamics as a classification task over end points. However, CatFlow is still a continuous model, and integrates -rather than iteratively samples -during generation. 
experiments: 
 We evaluate CatFlow in three sets of experiments. First, we consider an abstract graph generation task proposed in [35] , where the goal of this task is to evaluate if CatFlow is able to capture the topological properties of graphs. Second, we consider two common molecular benchmarks, QM9 [40] and ZINC250k [19] , consisting of small and (relatively) large molecules respectively. This task is chosen to see if CatFlow can learn semantic information in graph generation, such as molecular properties. Finally, we perform an ablation comparing CatFlow to standard flow matching, specifically in terms of generalization. The experimental setup and model choices are provided in appendix D.Note that we treat graphs as purely categorical/discrete objects and do not consider 'geometric' graphs that are embedded in e.g. Euclidean space. Specifically, for some graph with K v node classes and K e edge classes, we process the graph as a fully-connected graph, where each node is treated as a We first evaluate CatFlow on an abstract graph generation task, including synthetic and real-world graphs. We consider 1) Ego-small (200 graphs), consisting of small ego graphs drawn from a larger Citeseer network dataset [48] , 2) Community-small (100 graphs), consisting of randomly generated community graphs, 3) Enzymes (587 graphs), consisting of protein graphs representing tertiary structures of the enzymes from [47] , and 4) Grid (100 graphs), consisting of 2D grid graphs. We follow the standard experimental setup popularized by [62] and hence report the maximum mean discrepancy (MMD) to compare the distributions of degree, clustering coefficient, and the number of occurrences of orbits with 4 nodes between generated graphs and a test set. Following [20] , we also use the Gaussian Earth Mover's Distance kernel to compute the MMDs instead of the total variation.The results of the Ego-small and Community-small tasks are summarized in table 1, and additional results (and error bars) are provided in appendix C. The results indicate that CatFlow is able to capture topological properties of graphs, and performs well on abstract graph generation.Molecular generation entails designing novel molecules with specific properties, a complex task hindered by the vast chemical space and long-range dependencies in molecular structures. We evaluate CatFlow on two popular molecular generation benchmarks: QM9 and ZINC250k [40, 19] .We follow the standard setup -e.g. as in [49, 33, 58, 20] -of kekulizing the molecules using RDKit [25] and removing the hydrogen atoms. We sample 10,000 molecules and evaluate them on validity, uniqueness, and Fréchet ChemNet Distance (FCD) -evaluating the distance between data and generated molecules using the activations of ChemNet [39] . Here, validity is computed without valency correction or edge resampling, hence following [63] rather than [49, 33] , as is more reasonable due to the existence of formal charges in the data itself. We do not report novelty for QM9 and ZINC250k, as QM9 is an exhaustive list of all small molecules under some chemical constraint and all models obtain (close to) 100% novelty on ZINC250k. 3 3 CatfFlow obtains 49% novelty on QM9. The results are summarized in table 2 and samples from the model are shown in fig. 2 . CatFlow obtains state-of-the-art results on both QM9 and ZINC250k, virtually obtaining perfect performance on both datasets. It is worth noting that CatFlow also converges faster than flow matching and is not computationally more expensive than any of the baselines either during training or generation.To understand the difference in performance between CatFlow and a standard flow matching formulation we perform ablations. we focus on generalization capabilities, and as such consider ablations that the number of parameters in the model and the amount of training data.In fig. 3 we report a score, which is the percentage of generated molecules that is valid and unique.CatFlow not only outperforms regular flow matching in the large model and full data setting, but is also significantly more robust to a decrease in model-size and data. Moreover, we observe significantly faster convergence (curves not shown). We hypothesize this is a consequence of the optimization procedure not exploring 'irrelevant' paths that do not point towards the probability simplex.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide the hyperparameters and experimental detail appendix D. Moreover, we took all of our choices from previous works of [20] , which is also publicly available. Last, our code will be made public.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 We have introduced a variational reformulation of flow matching. This formulation in turn informed the design of a simple flow matching method for categorical data, which achieves strong performance on graph generation tasks. Variational flow is very general and opens up several lines of inquiry. We see immediate opportunities to apply CatFlow to other discrete data types, including text, source code, and more broadly to the modeling of mixed discrete-continuous data modalities. Additionally, the connections to score-based models that we identify in this paper, suggest a path towards learning both deterministic and stochastic dynamics.Limitations. While the VFM formulation that we identify in this paper has potential in terms of its generality, we have as yet only considered its application to the specific task of categorical graph generation. We leave other use cases of VFM to future work. A limitation of CatFlow, which is shared with related approaches to graph generation, is that reasoning about the set of possible edges has a cost that is quadratic in the number of nodes. As a result CatFlow does not scale well to e.g. large proteins of 10 4 or more atoms.Ethics Statement. Graph generation in general, and molecular generation specifically, holds great promise for advancing drug discovery and personalized medicine. However, this technology also poses ethical concerns, such as the potential for misuse in creating harmful substances. In terms of technology readiness, this work is not yet at a level where we foresee direct benefits or risks.Acknowledgements. This project was supported by the Bosch Center for Artificial Intelligence.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We provide error bars and how they are computed in appendix C.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not publish such models and as such do not require safeguards for them.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite the authors of the relevant assets where used, see section 6. All relevant information about licenses is publicly available and respected in our code.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators.  
 
 Appendix: We derive eq. ( 9) that states the equivalence of the optimization of the VFM objective and minimization of the KL divergence between the true endpoint distribution p t (x 1 | x) and the variational approximation q θ t (x 1 | x). Note thatEQUATIONwhere t ∼ Uniform(0, 1), x ∼ p t (x) and x 1 ∼ p data (x 1 ).First, we rewrite the KL divergence as combination of entropy and cross-entropy KL (p || q) = H(p, q) -H(p):EQUATIONEQUATIONWe observe that the entropy term does not depend on the parameters θ. Consequently we disregard it when optimising the variational distribution q θ t (x 1 | x). Second, we rewrite the cross-entropy term:EQUATIONTherefore, the second part of eq. ( 27) corresponds to negative cross-entropy and minimisation corresponds to maximisation of negative cross-entropy.Theorem 3. Assume the conditional vector field u t (x | x 1 ) is linear in x 1 and is of the formEQUATIONwhere A t (x) : [0, 1]×R D → R D ×R D and b t (x) : [0, 1]×R D → R D . Moreover, assume that A t (x) is an invertible matrix and q θ t (x1 | x) = N (x 1 | µ θ t (x), Σ t (x)), where Σ t (x) = 1 2 (A ⊤ t (x)A t (x)) -1 . Then, VFM reduces to flow matching.Proof. Let us substitute the assumed form of q θ t (x 1 | x) into the VFM objective:EQUATION= -E t,x,x1 log (2π)EQUATION(32)EQUATIONwhich is what we wanted to show.Theorem 1. Assume that the conditional vector field u t (x | x 1 ) is linear in x 1 . Then, for any distribution r t (x 1 | x) such that the marginal distributions coincide with those of p t (x 1 | x), the corresponding expectations of u t (x | x 1 ) are equal, i.e.EQUATIONProof. Applying the linearity condition, we rewrite the conditional vector field u t (x|x 1 ) as such:EQUATIONwhereA t (x) : [0, 1] × R D → R D × R D and b t (x) : [0, 1] × R D → R D .Then, we substitute it into equation:EQUATIONThen, we know, that for any distribution r such that the marginal distributions coincide with those of p the following holds:EQUATIONApplying this fact to eq. ( 38) we obtain:EQUATIONwhich is what we wanted to show.Let G = (V, E) be a graph with node and edge features given by H n ∈ R |V|×dn and H e ∈ R |V|×|V|×de respectively, and let x denote the graphs and its features. Moreover, let π ∈ S |V| be a permutation and P its associated permutation matrix, such that the action of the group is defined as• π • H n := PH n ,• π • H e := PH e P ⊤ .To simplify notation, we will simply write π • x to denote the above operation. Lemma 1. If µ t (x) is permutation equivariant w.r.t S |V| , then so is v t .Proof. Note thatv t (π • x) = µ t (π • x) -π • x 1 -t = π • µ t (x) -π • x 1 -t ,where the last step follows from permutation equivariance. Moreover, since π acts on x through permutation matrices, we can leverage the distributive property of linear operators, i.e. we conclude that π• µ t (x) -π • x 1 -t = π • (µ t (x) -x) 1 -t = π • v t (x),finishing the proof.Theorem 5. Let p 0 be an exchangeable distribution -e.g. a standard normal distribution -and µ t (x) be permutation equivariant. Then, all permutations of graphs are generated with equal probability.Proof. By result 1, we know that µ t (x) being permutation equivariant implies that v t (x) is permutation equivariant. Moreover, if we let Γ(x) := x + 1 0 v t (x t )dt, the for all π ∈ S n we have thatΓ(π • x) = π • x + 1 0 v t (π • x t )dt = π • x + 1 0 π • v t (x t )dt = π • Γ(x),where again the last step follows by basic properties of linear operators. Therefore, since p 0 assigns equal density of all permutations of x, the resulting distribution p 1 preservers this property, which is what we wanted to show.In this subsection we want to derive the equation that allows us to express the score function ∇ x log p t (x) in terms of the posterior probability path p t (x | x 1 ). Notice thatEQUATIONIn this subsection, we discuss how the VFM framework can be applied to construct stochastic generative dynamics and how it relates to score-based models.First, let us consider the marginal vector field u t (x). It provides the deterministic dynamic that can be written as the following ordinary differential equation (ODE):EQUATIONFor the vector field u t (x) we know that -starting from the distribution p 0 (x) -it generates some probability path p t (x). However, as we know from [54, 3] , if we have access to the score function ∇ x log p t (x) of distribution p t (x), we can construct a stochastic differential equation (SDE) that, staring from distribution p 0 (x), generates the same probability path p t (x):EQUATIONwhere g : R → R ≥0 is a scalar function, and w is a standard Wiener process.Given thatEQUATIONwe can rewrite eq. ( 50) in the following form: dx = ũt (x)dt + g t dw, where ( 52)EQUATION)Importantly, the function g t does not affect the distribution path p t (x): it only changes the stochasticity of the trajectories. In the extreme case when g t ≡ 0, the SDE in eq. ( 50) coincides with the ODE in eq. ( 49).Moreover, if we construct the stochastic dynamics in this way, we obtain score-based models as a special case. In score-based models [54] the deterministic process that corresponds to probability path p t (x) has the following form:EQUATIONwhere f t (x) is tractable. Substituting this u t (x) into eq. ( 50) we obtain:EQUATION)The SDE in eq. ( 55) only depends on the score function, so in score-based models, the aim is to learn the score function.Now, let us note that similarly to vector field u t (x), the drift term ũt (x) in eq. ( 52) can be expressed in terms of an end point distribution p t (x 1 | x). Consequently, having a variational approximation of end point distribution q θ t (x 1 | x), allows us to construct an approximation of the drift term ũt (x):EQUATIONwhereEQUATIONThus, we may define the following approximated SDE:EQUATIONThen, eq. ( 58) is not just simply a new dynamic, it is a family of dynamics that admits deterministic dynamics as a special case when g t ≡ 0. Importantly, the only thing we need to construct the stochastic process in eq. ( 58) is an approximation of the end point distributions q θ t (x 1 | x). Additionally we know that when p t (x 1 | x) = q θ t (x 1 | x), ũt (x) = ṽθ t (x), the processes coincide for all functions g t . Therefore, we can train the model with the same objective as in VFM.In this subsection, we leverage the connections of VFM with stochastic processes to show that a reweighted integral over the point-wise VFM objective defines a bound on the data likelihood in the generative model. Theorem 4. Rewrite the Variational Flow Matching objective as follows:L VFM (θ) = E t,x L θ (t, x)whereEQUATIONThen, the following holds:EQUATIONwhere λ t (x) is a non-negative function and C is a constant.Proof. Let us consider the two stochastic processes, that we discussed in appendix A.6:EQUATION)Note that they both start from the same prior distribution p 0 (x). The first one, by design, generates probability path p t (x) and ends up in the data distribution p data (x) = p 1 (x). The second process generates some probability path q θ t (x) that depends on the variational distribution q θ t (x 1 |x). We want to find a variational bound on KL divergence between p 1 (x) and q θ 1 (x). We start by applying the result from [1] (see Lemma 2.22):EQUATION(60)EQUATION)EQUATIONEQUATION)Algorithm 1 Variational Flow MatchingSample x 1 from data and x 0 ∼ p 0 (x 0 ) Sample t ∼ U(0, 1)Compute x t = tx 1 + (1 -t)x 0 Compute loss L = -log q θ t (x 1 | x t ) Backpropagate. # Generation Sample x 0 ∼ p 0 (x 0 ) Solve ODE x 1 = x 0 + t=1 t=0 E q θ t [x1|xt]-xt 1-t+ε dt Return x 1 B.2 Categorical Variational Flow Matching (CatFlow) Algorithm 2 Categorical Variational Flow Matching (CatFlow) # Assume q t (x 1 | x t ) = D d=1 Cat(x d 1 | µ d t (x t ))where µ is learnable# Training Sample x 1 from data and x 0 ∼ p 0 (x 0 ) Sample t ∼ U(0, 1) Compute x t = tx 1 + (1 -t)x 0 Compute cross-entropy loss L = - D d=1 K d k=1 I[x d 1 = k] log µ dk t (x) Backpropagate. # Generation Sample x 0 ∼ p 0 (x 0 ) Solve ODE x 1 = x 0 + t=1 t=0 µt(xt)-xt 1-t+ε dt Return x 1 B.3 Gaussian Variational Flow Matching Algorithm 3 Gaussian Variational Flow Matching # Assume q t (x 1 | x t ) = D d=1 N (x d 1 | µ d t (x t ), I) where µ is learnable # Training Sample x 1 from data and x 0 ∼ p 0 (x 0 ) Sample t ∼ U(0, 1) Compute x t = tx 1 + (1 -t)x 0 Compute mean squared error loss L = 1 2 D d=1 ||µ d t (x t ) -x d 1 || 2 Backpropagate. # Generation Sample x 0 ∼ p 0 (x 0 ) Solve ODE x 1 = x 0 + t=1 t=0 µt(xt)-xt 1-t+ε dt Return x 1Here, we provide the results for CatFlow with standard deviations, as computed as in [20] through multiple seeds.Community-small 4 ≤ |V| ≤ 18 12 ≤ |V| ≤ 20 Degree ↓ Clustering ↓ Orbit ↓ Degree ↓ Clustering ↓ Orbit ↓ 0.To ensure a comparison on equal terms to baselines, we employ the same graph transformer network as proposed in [12] , which was also used in [58] , along with the same hyper-parameter setup. We summarize the parametrization of our network here.Just as done in DiGress, our graph transformer takes as input a graph (H n , H e ) and predicts a distribution over the clean graphs, using structural and spectral features to improve the network expressivity, which we denote as H g . Each transformer layer does the following operations:1. Node Features H n and Edge Features H e :(a) Linear Transformation: Apply linear transformations to both H n and H e . (b) Outer Product and Scaling: Compute the outer product of the transformed features and apply scaling.2. Node Features H n :(a) Feature-wise Linear Modulation (FiLM): Apply FiLM to the transformed node features using global features H g .Edge Features H e :(a) Feature-wise Linear Modulation (FiLM): Apply FiLM to the transformed edge features using global features H g .Self-Attention Mechanism: Here, the FiLM operation is defined as:FiLM(M 1 , M 2 ) = M 1 W 1 + (M 1 W 2 ) ⊙ M 2 + M 2for learnable weight matrices W 1 and W 2 , and PNA is defined as: PNA(X) = cat(max(X), min(X), mean(X), std(X))W.We report the hyperparameters here:All models were trained until convergence. Furthermore, all data splits are kept the same as in [20] , and hidden dimensions are kept the same as [58] . All experiments were run on a single NVIDIA RTX 6000 and took about a day to run. To compute the resulting distribution p t for CNF, one can use the change of variables formula:EQUATIONThis induces a probability path, i.e. a mappingp t : [0, 1] × R D → R >0 such that p t (x)dx = 1 for all t ∈ [0, 1].We say that v t generates this probability path given a starting distribution p 0 . In theory, one could try and optimize the empirical divergence between the resulting distribution p 1 and target distribution, but obtaining a gradient sample for the loss requires us to solve the ODE at each step during training, making this approach computationally prohibitive.One way to assess if a vector field generates a specific probability path is using the continuity equation, i.e. we can assess whether v t and p t satisfyEQUATIONwhere ∇ is the divergence operator. Note that by sampling x 0 ∼ p 0 , a new sample from p 1 can be generated by following this ODE, i.e. integratingEQUATION)NeurIPS Paper ChecklistFlow matching and diffusion models have also been proposed for geometric graph generation, e.g. in[23,55] and[17,57] respectively, but since these approaches are continuous (as they generate coordinates based on some conformer) they consider a fundamentally different task than the ones we consider here. 


Decision: The work suggests the use of variational framework in flow matching, a natural step in an interesting topic. This work is timely and important to the community and will be an excellent addition to NeurIPS.

Reviews: review: paper_summary: 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.
strength_weakness: strength_weakness: In general, the quality of the paper is good. But I have a few questions on clarity. Please see below.
questions: 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.
limitations: limitations: The authors adequately addressed the limitations
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: The authors have addressed the limitations at the end of the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: The limitations has been adequately addressed.

Comments: comment: Thanks for the response. I am happy to keep my original rating of the paper.
Dear reviewer vH9B, we wondered if the added results and comparison to Dirichlet FM addressed your main concerns regarding our work. Thanks a lot for your review once more.