Fisher Flow Matching for Generative Modeling over Discrete Data

We thank the reviewer for their time and detailed feedback. We focus here on addressing the key clarification points raised in this review while the global response contains new experiments on molecules and language modelling with additional baselines. ## Riemannian OT is of low-novelty We acknowledge the reviewer’s concern that OT for Riemannian manifolds has previously been considered in [1]. Unfortunately, the general theory of optimal transport on manifolds Villani 2003 [2] (Theorem 2.47) does not give a prescription for a fast computation implementation of OT on manifolds—i.e., analytic expressions for the geodesic cost $c(x_0, x_1)$. While [1] found a closed form solution to this cost for $SE(3)$, it cannot be generally extended to all Riemannian manifolds, since, for arbitrary manifolds, one must measure the geodesic by simulating the expensive Euler-Lagrange equations, which are 2nd order PDEs.

Our work with Fisher-Flows pushes the advantage of Riemannian OT within generative modelling by showing that: 1.) The probability simplex and the sphere both have closed form expressions for the OT cost $c(x_0, x_1)$. 2.) We demonstrate that OT via Fisher-Flows can be used for discrete data. We highlight that this is the first use of OT for generative modelling over discrete data. 3.) As proved in Appendix D, OT over the probability simplex elicits new benefits, such as the fact that $p_t$ becomes the Wasserstein geodesic between $p_0$ and $p_1$, and the associated velocity field $v_t$ minimises the kinetic energy. These two points were not rigorously proven in [1]. ## Clearer experimental demonstration over the benefits of the Fisher-Rao metric and Sphere map ### OT in Fig 2 With regards to Fig. 2, we find that both c) and e) to be of comparable visual quality, and, more importantly, of a higher one than that of their non-OT counterparts. As this is a small toy synthetic density estimation task, the benefits of the Fisher-Rao metric and the added numerical benefits of the sphere map are less pronounced.

Comparison of FR on Fig 3

Since the sphere is more numerically stable and, for practical problems—where the number of dimensions can greatly increase—we expect Fisher-Flows on the sphere to perform the best—regardless of OT. We see this in Fig 1 of the 1pg rebuttal PDF, Dirichlet Flow Matching is often worse than FF-no OT on the sphere as we increase the number of categories. With regards to OT, our experimental findings are in line with previous use of OT in generative models, where lowering variance in training and shorter paths during inference (fewer ODE integration steps [3,4]) are observed. ### The other two experiments only show FF-OT sphere. Our new experiments on molecule generation and language modelling (see 1pg rebuttal PDF) used Fisher-Flows on the sphere and without OT. In our language modelling results, we found that Fisher-Flow slightly outperforms discrete diffusion approaches (e.g., D3PM, SEDD) while marginally outperforming concurrent SOTA work on Masked Discrete Diffusion (MDLM [5]) released after the submission deadline. This suggests that geometry plays a key role in improvement as we compare against conventional discrete diffusion methods on higher dimensional problem setting ~800k categories (text). ### No ablation on other potential metrics. Thank you for this suggestion! We have included new ablations in our 1 pg rebuttal PDF that use the Euclidean metric (Linear Flow Matching) on the simplex for both the synthetic data and DNA Promoter and Enhancer experiments. We observe that the Fisher-Rao metric is still on par and better than Dirichlet FM, in terms of empirical performance. Finally, considering other metrics beyond the Fisher-Rao and Euclidean is an interesting thought; however, it is unclear whether it is possible to easily obtain closed-form geodesics, which are needed for simulation-free training as mandated in the flow matching framework.

Q1. Use of Riemannian OT vs. Prior Work

Please see our detailed response under the heading “Riemannian OT is of low-novelty” for a deeper discussion.

Q2. Flexibility of source distribution in Campbell et. al 2024?

We note that Campbell et. al 2024 explicitly state that the “conditional flows we use in this paper linearly interpolate towards $x_1$ from a uniform prior or an artificially introduced mask state, M.” They do not make any theoretical claims about the generality of their starting source distribution unlike we do for Fisher-Flows. While it is potentially possible to accommodate for a more flexible source distribution in their CTMC framework, it is not immediately clear how to do this easily. We will refine our discussion over Campbell et. al 2024 in our paper to highlight this aspect.

Q3. Ablation on geometry vs. OT?

Please see our previous response on Fisher-Flows without OT on language modelling.

Closing comment

We hope that our responses were sufficient in clarifying all the great questions asked by the reviewer. We thank the reviewer again for their time, and we politely encourage the reviewer to consider updating their score if they deem that our responses in this rebuttal along with the new experiments merit it.

References

[1] Bose et. al (2024). SE(3)-Stochastic Flow Matching for Protein Backbone Generation. arXiv preprint arXiv:2310.02391.

[2] Villani. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008. ISBN 9783540710509.

[3] Tong et. al. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint 2302.00482, 2023.

[4] Klein et. al . "Equivariant flow matching." Advances in Neural Information Processing Systems 36 (2024).

[5] Sahoo et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).

I want to thank the authors for their efforts in adding more experiments.

First regarding the Riemannian OT, thank you for emphasizing the contributions upon previous works. I am now convinced that more credit should be attributed to the authors than I initially acknowledged. Second regarding the significance of change in geometry from simplex to sphere, Table 1 in pdf actually shows no advantage to the sphere and to my understanding Table 2 compares only to discrete diffusion methods (where I would expect comparison like Table 1 to support such a claim.) However, I do find the work in this paper interesting and well presented, I will increase my score.

We thank the reviewer for their time and effort assessing our paper. We now address the key points in the review, while additional experiments are included in the global response and rebuttal PDF.

Circular justifications for the Fisher Rao

We understand that certain aspects of the theory could be explained more clearly. We now explain our reasoning step by step. First, note that a distance function on a Riemannian manifold is fully determined by the choice of metric, and that choosing a different distance function amounts to choosing a different metric.

We require the following key desiderata from a metric to employ flow matching:

• Analytic expressions of manifold operations like the exponential and logarithmic map.
• Easy parameterisation of the tangent space in order to define vector fields.
• Interpretability and naturality, i.e. connection to established statistical and physical theory.
• Numerically stability of the metric over the entire manifold.

The use of the FR metric satisfies both points 1 and 2 above and more and point 4 is achieved via the sphere map. In addition, we argue that the pervasiveness of the use of the forward KL divergence in generative models such as diffusion, flows, and VAE, testifies to its naturalness. As to point 3, the KL divergence is natural due to its strong link to fundamental concepts from Information Theory, being interpretable in terms of entropy. In this sense, KL is perhaps the simplest natural divergence over categorical probabilities.

We now explain why our choice does not lead to circular reasoning using Prop 1. Note that Prop 1 starts off with hypotheses over a generic metric “g” in the $W_{2,g}$ distance from the hypothesis. The thesis of the proposition shows that natural gradient with respect to KL-divergence-based $W_{KL}$ distance produces naturally the $W_{2, g_{FR}}$ gradient flow, with now “$g_{FR}$” being FR metric, which was not assumed in the hypotheses. Thus this is not a circular statement, as it deduces FR as the limit of KL-divergence, at the level of probability spaces.

Are there computational reasons to use this?

The FR metric leads to theoretical benefits as it is also the Wasserstein geodesic over probability paths on the simplex. The Euclidean metric, as explored in Dirichlet Flow Matching, does not enjoy such benefits and arguably is one reason that is empirically less performant than Fisher-Flows.

In Section 3.4, a similar problem arises…

We believe there may be a slight misunderstanding as the logic of our setup is to first pick a metric and then show OT is computationally feasible. Indeed other metrics are possible and will lead to different OT maps; but, crucially, they may not be easy to compute for generative models. Fortunately, for the FR metric the distance is simply the distance on a sphere, which is easy to compute allowing us to efficiently compute the OT cost $c(x_0, x_1)$.

Why do we care about the lengths of paths defined by the Fisher-Rao metric? … velocity fields faster to simulate?

Shorter paths lead to lower variance training: on the sphere and the simplex. The paths do not cross at intermediate points $x_t$, which would have added to noisier training—in line with the literature see [1,2]. We note OT often leads to vector fields that are faster to simulate using numerical solvers as they incur smaller per step error, as studied for $SE(3)$ in [3].

Here do you mean to imply something about this choice of p_t? … Prop. 2 shows that, by using FR geodesics as our flow, we automatically use the geodesic path between $p_0$ and $p_1$ with respect to the natural Wasserstein metric over probability measures. A consequence of this is that the links to kinetic energy and dynamical formulations apply (see Prop. 4). ## Empirical validation There is no experiment that uses a flexible source distribution We invite the reviewer to kindly read our global response which includes new experiments over text that uses a new source distribution which is a convex combination of a specialised mask token and the uniform distribution.

It's unclear why the sphere map helps ….

The metric appears when we compute $x_t$ due to exp and log map on the simplex (see Appendix B). Looking at equation 11, for the exp map, we notice we need to compute the norm which is induced by the Riemannian metric so $<v^2_p, v^2_p>_{FR}$. Similarly, the log map on the simplex in equation 12 explicitly requires the inner product. Exploiting the sphere map, we obtain $x_t$ and avoid numerical instability near the boundary of the manifold. Other metrics beyond the FR are indeed possible on the simplex but it may not have analytic formulae for $x_t$.

Similarly, I think a good ablation to do is just to use the Euclidean metric here …

Please see our global response PDF where we have included the Euclidean metric in Linear Flow Matching on the toy datasets as well as our DNA experiments. ## Clarifying experimental setup In addition to DNA sequences we have now also included more standard generative modelling experiments on molecules and larger scale language modelling that demonstrate the scalability of Fisher-Flows on a host of different domains.

We thank the reviewer again for their valuable feedback. We hope that our rebuttal addresses their questions and concerns, and we kindly ask the reviewer to consider fresher evaluation of our paper if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise. ## References [1] Bose et. al (2024). SE(3)-Stochastic Flow Matching for Protein Backbone Generation. arXiv preprint arXiv:2310.02391.

[2] Tong et. al. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint 2302.00482, 2023.

[3] Klein et. al "Equivariant flow matching." Advances in Neural Information Processing Systems 36 (2024).

We would like to thank the reviewer for the time and effort they spent on reviewing our work. We are appreciative of the fact that the reviewer found our geometric perspective a “novel extension” to flow matching and that this leads to numerically stable training which is supported by better empirical performance. We now address the concerns raised by the reviewer. ## Experiments beyond DNA We acknowledge the reviewer's valid concern. In this rebuttal we included two more standard experimental settings for discrete data in molecular generation (QM9) and language modelling over text (LM1B). We kindly invite the reviewer to read our global response, and the 1pg rebuttal PDF, which gives more detail on these experiments. In summary, we found Fisher-Flows to be scalable to over 800k categories in text and achieve the best upper bound to text PPL amongst discrete diffusion models. For QM9 our results outperform flow matching and are competitive, almost saturating the task, with molecular diffusion models. We hope that our new experiments sufficiently demonstrates the scalability of Fisher-Flows to much higher dimensions and across a greater breadth of domains.

Choice of baselines

As suggested by the reviewer we’ve now included multiple baselines as part of our 1pg rebuttal PDF. Specifically, we added Linear Flow Matching (Euclidean metric) and Multinomial Flows to the synthetic task, where we found the latter to be numerically unstable as we increased the number of categories. We also added Linear Flow Matching to our DNA Promoter and Enhancer experiments where we found Fisher-Flows to be on par. For our new language modelling experiment on LM1B we compared against discrete diffusion models such as SEDD as well D3PM, as well as SOTA concurrent work in masked discrete diffusion in MDLM [1]. We found that Fisher-Flows still outperforms previous methods like D3PM and SEDD and is marginally better than MDLM.

## Reporting of Dirichlet Flow Matching results

“Why the original numbers from the Dirichlet FM paper were not used?

Using the reference code implementation of DFM we noticed severe methodological issues that made it clear that the original numbers in their paper were an inaccurate reflection of true performance. We found the following flaws in their experimental setup: We found for the DNA datasets the provided command used --subset_train_as_val argument which sets the validation set as the training set. This meant that their reported metrics were measured on the training set. We found their released DNA classifier to have only $11$ test accuracy which makes it a poor choice as a representation space to compute FBD. Their synthetic data experiment did not standardise a training and validation split and meant that their training set size was a function of run time—effectively infinite. Operationally, DFM was trained for $10^{10}$ points per epoch which is a tremendous amount for a synthetic task. We standardised the experiment to train all models with a more modest 10^6 points.

After fixing these issues we reran DFM and found that it performed worse than what was initially reported. We are appreciative of the DFM authors for releasing their code that enabled us to find these experimental flaws, some which we outlined in Appendix F.4.2 and F.4.3 but will discuss further in our updated manuscript.

Reporting PPL as a metric

We agree with the reviewer that calling our reported result perplexity on DNA is inaccurate. To compute our metric for DNA sequences we used the continuous time change of variable of the flow along the ODE trajectory over the manifold to compute the log-likelihood of the observed sequence (as done in Riemannian Flow Matching). Since the output space is discrete this likelihood on the manifold is an upper bound to the discrete log likelihood. More accurately this is a valid Evidence Lower Bound (ELBO) which was proven in concurrent masked discrete diffusion models [1, 2]. We also note that for autoregressive models the used PPL metric is a specific factorisation over the joint $\log p(x_{1:N})$ and any valid ELBO is comparable. Nevertheless, the reviewer is correct and we will amend our reported results by indicating this upperbound when reporting PPL for diffusion/flow matching models.

Finally, we have also included a new Gen PPL eval metric for DNA sequences. This is the PPL supplied a larger pre-trained autoregressive model on the outputs samples of a generative model. We found this metric could not distinguish the performance of DFM and Fisher-Flows, but we report it for completeness.

Q1

Please see our global response for new larger scale experiments, including on text.

Q2

Please see above about additional baselines as well as our global response.

Q3

Please see our response above which finds flaws in the empirical setup of DFM.

Q4

We computed an upper bound to the PPL using the continuous time change of variable formula associated with an ODE. This is a valid ELBO see [1-2] for proof.

Q5

OT in Fisher-Flows is not computationally expensive as both the sphere and the simplex have closed form distance functions which allows us to efficiently calculate the OT plan over mini-batches. This adds minimal overhead to training and is not affected by scaling.

Q6

Thank you for pointing this out. We will fix it in the updated paper. ## Conclusion We thank the reviewer again for their time. We believe we have answered to the best of our ability all the great questions raised by the reviewer. We hope our answers allow the reviewer to consider potentially upgrading their score if they see fit. We are also more than happy to answer any further questions.

## References [1] Sahoo et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).

[2] Shi et al. "Simplified and Generalized Masked Diffusion for Discrete Data." arXiv preprint arXiv:2406.04329 (2024).