Fisher Flow Matching for Generative Modeling over Discrete Data

Summary

The paper aims to build a Flow Matching [1] based generative model for discrete data. The approach models the discrete data as a categorical distribution that resides on the simplex, thereby translating the problem into continuous flows. Equipped with the Fisher-Rao metric, the simplex is identified as a Riemannian manifold, and the task is reduced to Riemannian Flow Matching [2]. To further simplify the geometry of the space, the sphere map is used, which identifies the simplex with the positive orthant of a hypersphere. Finally, the conditional path between every pair $(x_0, x_1)$ is defined via the geodesic interpolant induced by the FIsher-Rao metric. For marginalization, two joint distributions are considered: 1. $x_0$ and $x_1$ are independent, i.e., $\pi(x_0,x_1) = p_0(x_0)p_1(x_1)$ as in [2]. 2. $\pi(x_0,x_1)$ is the optimal transport (OT) plan between $p_0$ and $p_1$, approximated with minibatch [3].

Strengths

1. The paper is well-written, and the method is rigorously presented. 2. The author successfully frames the proposed method for generative models of discrete data in terms of well-established existing methods, which eases the understanding of the approach and potentially allows for a relatively low cost of implementation. 3. Aside from the computational cost of the OT plan, the method does not introduce additional computational costs compared to previous works.

Weaknesses

1. The use of Riemannian Optimal Transport with a mini batch is of low novelty. This has also been done in [3] and on euclidean space in [6,7]. 2. Two key contributions of the paper seems to be the change in geometry, that is the introduction of the Fisher-Rao metric and the sphere map. However the paper lacks a clear experimental demonstration of the portion these two components contribute to results. Specifically: 2.1. Regarding Figure 2, allegedly, the author claim that Fisher-Flow (FF)-OT on $\mathbb{S}_+^2$ (e) performs best, however it is not clear, one may even say that FF-OT on the $\Delta^2$ (c) perform better. Furthermore, even without FF-OT it is not clear that using the sphere map helps ( (b) vs (d)). 2.2. In Figure 3.a the author compares FF simplex, FF-OT simplex, FF sphere, and FF-OT sphere, that shows that FF-OT sphere preforms the best. It is also apparent that with increasing dimension size the advantage of the use of the OT plan decreases. And Figure 3.b shows that FF-OT sphere out performs [4]. However, comparing Figure 3.a and 3.b it seems that the other three, i.e., FF simplex, FF-OT simplex, FF sphere are actually outperformed by [4] which uses flows on the simplex but no OT plan. To conclude, this experiment suggests that the apparent advantage of FF-OT sphere on [4] may exist only for relatively low dimensions. 2.3. The other two experiments only show FF-OT sphere. 2.4. No ablation that focus on the use of Fisher-Rao metric compared to other potential metrics (though the author does provide a theoretical justification).

Questions

1. Can the author expand on the difference of his use in Riemannian Optimal Transport with a mini batch compared to previous works? 2. The author claims that [5] is restricted in the choice of possible source distributions. Can the author please elaborate on that? 3. Could the author possibly provide more experimental evidence to demonstrate that the proposed change in geometry indeed contributes to a considerable improvement and that the gain in results is not mainly from the OT plan?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

It is suspected that the proposed method might not be as effective in large dimensions (i.e., large vocabulary size ) compared to other methods. [1] Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2023). Flow Matching for Generative Modeling. arXiv preprint arXiv:2210.02747. [2] Chen, R. T. Q., & Lipman, Y. (2023). Flow Matching on General Geometries. arXiv preprint arXiv:2302.03660. [3] Bose, A. J., Akhound-Sadegh, T., Huguet, G., Fatras, K., Rector-Brooks, J., Liu, C.-H., Nica, A. C., Korablyov, M., Bronstein, M., & Tong, A. (2024). SE(3)-Stochastic Flow Matching for Protein Backbone Generation. arXiv preprint arXiv:2310.02391. [4] Stark, Hannes, et al. "Dirichlet flow matching with applications to dna sequence design." arXiv preprint arXiv:2402.05841 (2024). [5] Campbell, Andrew, et al. "Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design." arXiv preprint arXiv:2402.04997 (2024). [6] Tong, Alexander, et al. "Improving and generalizing flow-based generative models with minibatch optimal transport." arXiv preprint arXiv:2302.00482 (2023). [7] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings." arXiv preprint arXiv:2304.14772 (2023).

Summary

This paper proposes a framework that enables flow matching over a d-dimensional simplex, by instantiating a riemannian flow matching algorithm using the Fisher-Rao metric. Some motivation in connecting the Fisher-Rao metric to natural gradient descent and Riemannian optimal transport is used to justify the choice of Fisher-Rao metric. Experimental validation is carried out on DNA sequences.

Strengths

The proposed method is quite a natural application of riemmanian flow matching to distributions over simplices.

Weaknesses

- I felt the writing is a bit too mathematically dense, and the justifications for the Fisher-Rao metric feel a bit circular (see below for questions on Propositions 1 & 2). - Additional ablation experiments on non-toy data would be good to have to verify some of the claims in the paper. - It's unclear to me (due to lack of expertise in the application area) to judge how well this approach performs.

Questions

### Regarding the mathematical justifications I find that Sections 3.3 and 3.4 are circular in their reasoning. In Section 3.3, it is stated "the choice of Fisher-Rao metric [is] the optimal one on the probability simplex". However, here "optimality" is defined in terms of the expected KL, W_KL. This is an arbitrary choice of optimality. There is obviously a correspondence between a distance function and a Riemannian metric, and any metric can be stated as the optimal one for some distance function. So: - Why is considering the W_KL ball in Eq 8 useful? Are there computational reasons to use this? Are there theoretical justifications over other distance functions? In Section 3.4, a similar problem arises. Here it is stated that the flows (1) lead to shorter global paths and (3) have lower kinetic energy. However, this is circular because the distance of a path (and consequently, the kinetic energy) is an implication of the chosen Riemannian metric g. Here the correspondence between distance function <==> Riemannian metric shows up again. - Why do we care about the lengths of paths defined by the Fisher-Rao metric? - Does the choice of Fisher-Rao metric give velocity fields that are faster to simulate? I'm not completely sure what the point of Proposition 2 is. Equation 10 looks to me like a definition of p_t given a coupling pi. That is, given that we've already solved the optimal transport coupling, we can then take any probability path p_t that transports between the marginals of pi, so Proposition 2 is just stating a definition of a particular choice of p_t. - Here do you mean to imply something about this choice of p_t? E.g., that it allows expressing W_2 in terms of p_t through a dynamical formulation, where the kinetic energy (as defined by the Fisher-Rao metric) shows up? ### Regarding empirical validation The paper makes the following claims: (1) flexibility of source distribution (2) the sphere map has better numerical stability properties (3) the choice of Fisher-Rao metric enables continuous reparameterisation and OT The experiments do not yet completely justify these points, from what I can tell. (1) There is no experiment that uses a flexible source distribution. - Is there a case for using something that isn't a Dirichlet distribution which can be handled by Dirichlet FM? Or at least, a justification over the uniform noise distribution? (2) It's unclear why the sphere map helps. Concretely, there is a division by p in the computation of the inner product (Eq 5). However, this inner product is unneeded for the geodesic computation, since we can easily solve the geodesic on the sphere and then map back to the simplex. So this inner product seems to show up only in Eq 7, the training objective. However, given that this metric only depends on x_t, and that the model conditions on x_t, the optimal solution of CFM (Eq 7) is unaffected. So one can easily use any other metric for CFM without affecting the optimal solution. - Have the authors tried just using another metric for CFM and training on the simplex representation? (3) Similarly, I think a good ablation to do is just to use the Euclidean metric here. One can easily use the interpolants defined by the Fisher-Rao metric instead of linear, but otherwise just use regular Euclidean flow matching. Otherwise, related to the first concern, it is unclear why the connections to kinetic energy (as in Riemannian optimal transport) is important. ### Clarify about experiment setup As someone not familiar with DNA sequences, the experiment section does not provide enough information for me to gauge the usefulness of these experiments. For instance, it should at least be stated what d is (I imagine d=4?). My subjective opinion is that many methods could work reasonable well with such a small d. It would be even better if the authors could comment on the real world implications of their improved metrics (how does PPL affect a scientist who would use this generative model?)

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper discusses a bit the limitations regarding higher d, but does not provide details on why the current approach cannot be applied. I am also curious why the authors did not consider relatively toy datasets (such as variants of MNIST) that would be more understandable / appealing for the general machine learning community. For instance, it is unclear to me if non-DNA experiments would be too hard, or even too easy, compared to the chosen experiments.

Summary

The paper proposed a novel flow-based generative model called *Fisher-Flow* for discrete data. The model uses the Fisher metric to deduce a Riemmanian geometric structure of the statistical manifold. The authors also demonstrated the connections to natural gradient descent and optimal transport. Experiments on DNA datasets demonstrated better performance.

Strengths

1. The geometric perspective of the manifold of categorical distributions, to the best of my knowledge, is a novel extension of flow matching models for discrete data. The use of the Fisher-Rao metric can induce a Riemannian structure for continuous parameterization. 2. The authors proposed a diffeomorphism to the unit sphere to avoid numerical instability during training and demonstrated better performance in experiments. 3. The authors provided connections between the proposed model and natural gradient descent and Riemannian optimal transport with mathematical proof. The proposed model can potentially share such theoretical benefits. Ablation studies also supported such claims.

Weaknesses

1. Although being claimed as a general discrete generation model, **the proposed model was only tested on the specific task of DNA generation without further justification of its performance on more complex data such as image and text generation**. The two DNA datasets used in the paper have a small cardinality of 4 categories. Existing diffusion- or flow-based discrete generation models (e.g. BFN [1], D3PM [2], SEDD [3]) demonstrated good results on image generation (discretized pixel values or VQVAE tokens) and/or text generation (tokens) tasks which have higher cardinalities to show their scalability. It remains unclear whether the proposed model can be adapted for more general real-world tasks with a larger number of categories. 2. The choice of baselines is **too limited to be convincing enough** to demonstrate the proposed model's superior performance. 1. For the density estimation task, **no baseline was compared** (results were more of an ablation study). For the toy data in higher dimensions, **only Dirichlet FM was compared**. For these two tasks on simplex, models like multinomial flow [4] should be tested to provide a comparison. 2. For the enhancer DNA design task, **only Dirichlet FM was compared as a flow-/diffusion-based baseline**, despite that many other discrete diffusion models have been proposed, e.g., BFN [1], D3PM [2], SEDD [3]. Some of them were included in the related work section but none was tested as baselines in this task. 3. Moreover, one simple but effective and commonly used baseline is missing across the board, which is linear flow matching on simplex (which was used in DirichletFM [5] and other multimodal flow matching papers as well). 3. Another critical concern lies in the **inconsistency of the reported performance of the same baseline models in this paper and the original paper**. Numbers for the Dirichlet FM reported in this paper for the promoter design task are **significantly worse** than those in the original paper. The original paper reported an MSE of 0.0269 for their best model (Table 1 in [5]) which is far better than the number indicated in this paper (0.034, Table 1 in this paper). Moreover, the linear flow matching model that achieved 0.0281 in Table 1 of [5] (outperforming the MSE of FisherFlow) was dropped in this paper. Similar large discrepancies in the baseline model's performance can be also noted in the FBD scores for enhancer design. Comparing Table 2 in [5] and Table 3 in this paper, the performance for the baseline Dirichlet FM was also significantly worse. As the numbers for other baselines were directly copied without change, it is unclear why the authors reported DirichletFM's performance differently. **Using the original MSEs or scores, the results indicated that the proposed model was not as good as Dirichlet FM or even Linear FM**. 4. In terms of evaluation metrics, **it is unclear whether perplexity is a valid metric** to compare flow/diffusion based models against language models. Perplexity was initially proposed for autoregressive language models, in which it was calculated directly over conditional probabilities of discrete tokens: $$ PPL(X)=\exp\left(-\frac{1}{N}\sum_{k=1}^N\log p_\theta(x_k|x_{<k})\right) $$ However, in the continuous flow/diffusion setting, such probabilities may **not be well-defined** as they are neither autoregressive nor defined over discrete space. Therefore, their PPL cannot be calculated in this way. If the authors instead tried to calculate the joint distribution $\log p_\theta(x_{1:N})$, it remains unclear how to leverage the unconditional flow model to calculate this log-likelihood for arbitrary given input sequence $x_{1:N}$, as the flow is defined over the continuous simplex space with infinite possible initial probability samples $p_0$. The authors did not explain how they were able to calculate the PPL in the paper, and it is unclear whether the derived PPL is comparable to the ones computed for autoregressive models. [1] Graves, Alex, et al. "Bayesian flow networks." *arXiv preprint arXiv:2308.07037* (2023). [2] Austin, Jacob, et al. "Structured denoising diffusion models in discrete state-spaces." *Advances in Neural Information Processing Systems* 34 (2021): 17981-17993. [3] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023). [4] Hoogeboom, Emiel, et al. "Argmax flows and multinomial diffusion: Learning categorical distributions." *Advances in Neural Information Processing Systems* 34 (2021): 12454-12465. [5] Stark, Hannes, et al. "Dirichlet flow matching with applications to dna sequence design." *arXiv preprint arXiv:2402.05841* (2024). directly

Questions

1. Previous existing diffusion- or flow-based discrete generative models have achieved considerable success in image and text domains (see Weakness 1). As a general discrete generation model, will the proposed model also achieve comparable performance on these image or text generation tasks? Will it have scalability issues regarding datasets with a large number of categories? More experiments on other generative domains will be needed to demonstrate the model's effectiveness and scalability. 2. Can you provide more baselines in the visualization of the density estimation toy dataset and some quantitative evaluation metrics for this task to better demonstrate the generation quality? With the wide range of existing discrete diffusion or flow based models, more should also be compared as baselines for both toy examples and the enhancer generation task. Specifically, will the proposed model outperform simple linear flow matching? See Weakness 2 for some existing diffusion or flow baselines. 3. Why the original numbers from the Dirichlet FM paper were not used? If you reran the experiments in the Dirichlet FM paper, can you provide justifications for why you only reran the ones for Dirichlet FM but not for other baselines like linear flow matching? It is also unclear why the results of linear flow matching in the original Dirichlet FM paper were omitted -- it has better performance than FisherFlow according to the Dirichlet FM paper. The reported numbers in this paper were **significantly worse** than those in the original paper. If original numbers were used, the proposed model cannot outperform the Dirichlet FM model. See Weakness 3 for details. 4. How was perplexity calculated for flow-based models? Is it calculated in an autoregressive fashion as the language models? (See Weakness 4) Flow-based model is not an autoregressive model. During generation, all tokens are simultaneously denoised into meaningful sequences. If you instead modeled the joint distribution $\log p_\theta(x_{1:N})$, how was this marginal log-likelihood calculated for arbitrary input sequence $x_{1:N}$? How can you deal with the randomness in the initial samples $p_0$ if a deterministic cross-entropy loss needs to be computed? Furthermore, how was perplexity calculated for random sequences? The log probability $\log p(x)$ for one-hot random sequences should either be 0 or infinite. 5. OT seems to lead to better performance. Will it significantly add to the computation time during training? 6. There is a notation inconsistency in the paper about whether to refer to the Dirichlet flow matching model as *Dirichlet FM* (e.g., Table 1 & 2) or *DFM* (main text).

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Limitations were mentioned in the conclusion section, and the potential negative societal impact was adequately addressed.

Kindly awaiting more feedback

We thank you again for your time and feedback that allowed us to strengthen the paper with new experiments. As the end of the rebuttal period is fast approaching we were wondering if our answers in the rebuttal were sufficient enough to address the important concerns raised regarding 1.) the justification of using the Fisher-Rao metric and 2.) Empirical validation. We note that for point 2.) we included a host of new experiments, including a large-scale experiment on text on the LM1B dataset where we achieved SOTA among discrete diffusion and flow matching models. We would be happy to engage in any further discussion that the reviewer finds pertinent, please let us know! Finally, we are very appreciative of your time and effort in this rebuttal period and hope our answers are detailed enough to convince you to potentially upgrade your score if you believe it's merited.

Re: Decision to raise my score

We thank the reviewer for taking additional time to investigate DFM on DNA Promoter. We will most definitely open-source our code and include a detailed discussion regarding the discrepancy between our DFM reproduction and the published table results. We appreciate that the reviewer views our new experimental findings on LM1B as promising and that the reviewer finds using that applying the Fisher-Rao metric to be novel and interesting to the generative community. Finally, we thank the reviewer again for their effort in helping us strengthen our paper and we will incorporate all the great suggestions made by the reviewer in our updated draft.