We introduce Statistical Flow Matching (SFM), a novel and mathematically rigorous flow-matching framework on the manifold of parameterized probability measures inspired by the results from information geometry. We demonstrate the effectiveness of our method on the discrete generation problem by instantiating SFM on the manifold of categorical distributions whose geometric properties remain unexplored in previous discrete generative models. Utilizing the Fisher information metric, we equip the manifold with a Riemannian structure whose intrinsic geometries are effectively leveraged by following the shortest paths of geodesics. We develop an efficient training and sampling algorithm that overcomes numerical stability issues with a diffeomorphism between manifolds. Our distinctive geometric perspective of statistical manifolds allows us to apply optimal transport during training and interpret SFM as following the steepest direction of the natural gradient. Unlike previous models that rely on variational bounds for likelihood estimation, SFM enjoys the exact likelihood calculation for arbitrary probability measures. We manifest that SFM can learn more complex patterns on the statistical manifold where existing models often fail due to strong prior assumptions. Comprehensive experiments on real-world generative tasks ranging from image, text to biological domains further demonstrate that SFM achieves higher sampling quality and likelihood than other discrete diffusion or flow-based models.
Paper
References (68)
Scroll for more · 38 remaining
Similar papers
Peer review
Summary
The paper extends Flow Matching to discrete ("categorical") variables, similarly to a whole flourish of other papers that appeared in the last months, all of which work on the "probability simplex", i.e. they pass from a discrete alphabet A to the finite dimensional probability space P(A), which can be identified with a simplex with vertex set A. Here the emphasis is on the Fisher-Rao metric over this simplex, which the authors claim gives a natural riemannian metric that is well-suited to the flow-matching formalism.
Strengths
The method is well founded and the geometric insights are inspiring and seem "right". It seems that the method performs well, as indicated by experiments. The paper is well written and clear, easy to read execept for a few hiccups (mentioned in "questions" section).
Weaknesses
The Riemannian metric on the subset of the sphere $\{x\in \mathbb R^n: |x|^2=2, x_i\geq 0 \text{ for }1\le i\le n\}$ is not that different (in fact a $C^1$-close deformation of) the one on the actual "flat" simplex. So why would the results be so much different between the two cases? The grounding for a big gap from this setup to the standard one is not clear. It is also not clear why this method would have to perform better than Stark et al. "Dirichlet flow" paper.
Questions
By the way, what's the difference between this paper and https://arxiv.org/abs/2405.14664 ? I understand and follow well the core paper, so I have only a few minor questions: 1) line 94: what does it mean to "condition on the delta measure"? A measure is not a random variable right? I don't fully follow what this means. 2) line 151: "we assume a single data dimension" -- what does that mean? also line 152 "extended to any data dimension" is kind of confusing, but I'm sure this is minor.. please clarify? 3) for section 3.3, line 171 and following, says "the Fisher information metric defines a second-order optimization [...]".. I don't fully understand what this means exactly. A metric is a metric, it doesn't have agency and so it doesn't define anything. I think this could be clarified by adding a few more formulas and by expanding on what the authors actually mean. 4) also for section 3.4 I don't fully follow: what Optimal Transport are the authors considering, how is that set up and how is that justified? At the moment I don't find this section very useful to a reader. 5) about section 3.5: why would one want to calculate NLL? The reason for that is missing, so I think it would make sense to shorten the derivation of the NLL formula to a minimum (move them to the appendix?) and actually explain where this section is coming from, as it's not crucial to actual flow matching objectives, and at the moment seems a bit artificial tbh
Rating
6
Confidence
5
Soundness
3
Presentation
3
Contribution
3
Limitations
see "weaknesses" section.
I thank the authors for the rebuttal, but I am not understanding one of the answers. About the sphere metric, I don't see how taking a simplex and curving the interior (mapping from a flat simplex to a curved one) can prevent overshooting over the boundary. Could you please elaborate? To clarify my doubt, think of the case of 3 categories. We then compare a flat equilateral triangle with straight lines, to a "quadrant" of a sphere, where geodesics on a sphere are so to say, "pieces of equators", or maximal circunferences on the sphere. Then the segment going from a vertex of the straight triangle to the opposite side, makes an angle $\alpha\in[\pi/3, \pi/2]$ with the boundary whereas, in the curved setting, the geodesic from a vertex of the quadrant to the opposite side makes an angle $\alpha=\pi/2$ with the boundary. So why does one trajectory overshoot less than the other? I can see the performance of your experiments, but to me the geometry of the sphere quadrant and of a simplex still seem quite similar. In fact, adding to what I said in the review: the two metric spaces are not just $C^1$-diffeomorphic, but actually have finite distortion one with respect to the other, and I think the distortion constant is smaller than $2$. So, what would warrant such a big difference of performance then? Given the above comparison of agles of geodesics, I don't think the answer is related to overshooting, it must be something else. Emphasizing again, to be clear: this is just an important curiosity for me, because I can see the better performance of the experiments. But still I'm not convinced about the actual principle/reason beyond the better performance at the moment.
We thank your feedback for our rebuttal and we are more than happy to further discuss the potential benefits of considering the curved geometry induced by the Fisher information metric. We will further elaborate on our hypothesis of overshooting for SFM versus LinearFM. In the middle column of Fig.1, we plotted the geodesics between the same source and target pairs under the Euclidean setting (flat simplex) and the Riemannian settings (using the Fisher information metric). Besides the curved geodesics, we also noted that the spacing between adjacent points is also different (though they are linearly spaced with respect to the geodesic distance). It can be seen from the figure that points are **more clustered near the boundary than those near the middle region**. Also note that the ground truth conditional vector fields have constant length for all timestep $t$ for both the flat simplex (always $\mu_1-\mu_0$) and the sphere (the arc length between $x_0,x_1$). In this way, a predicted vector field on the sphere with the same norm near the boundary will instead move the point for a smaller distance on the simplex to avoid overshooting. The Euclidean vector field, on the other hand, remains constant near the boundary. Furthermore, for points near the boundary but not at the vertex (which may occur in our toy example), it can be seen from Fig.1 that the vector fields and the geodesic are **curved to be more parallel to the boundaries**. It can be also demonstrated mathematically (in our response to Reviewer xYE5) by looking at the direction term $\sqrt{\mu\odot\nu}-\langle\sqrt{\mu},\sqrt{\nu}\rangle \mu$. For $\mu$ close to the boundary with component $\mu_k\approx 0$, its corresponding vector field will also have a close to $u_k\approx 0$ component, which is different from linear flow matching's fixed $\nu-\mu$. In this way. the Riemannian vector field avoids further pushing the points outside the boundaries. We further noted that the **geodesic distance in Eq.3 cannot be bounded by the Euclidean distance**. More rigorously, there does not exist a finite constant $C$ such that $d\_\text{cat}(\mu,\nu)\le C\\|\mu-\nu\\|\_2,\forall \mu, \nu$. In other words, there isn't a finite distortion constant due to the singularity of the transform on the boundary. This can be demonstrated with Taylor expansion at $\mu$ as: $$ d_\text{cat}(\mu,\mu+\Delta\mu)=2\arccos\left(\sum_{i=1}^n\sqrt{\mu_i(\mu_i+\Delta\mu_i)}\right) \approx 2\arccos\left(\sum_{i=1}^n\mu_i+\frac{\Delta\mu_i}{2}-\frac{\Delta\mu_i^2}{8\mu_i}\right) =2\arccos\left(1-\sum_{i=1}^n\frac{\Delta\mu_i^2}{8\mu_i}\right) \approx2\sqrt{2}\sqrt{\sum_{i=1}^n\frac{\Delta\mu_i^2}{8\mu_i}}=\sqrt{\sum_{i=1}^n\frac{\Delta\mu_i^2}{\mu_i}} $$ Compared to the Euclidean distance $\\|\Delta\mu\\|=\sqrt{\sum_{i=1}^n\Delta\mu_i^2}$, it is clear that $d_\text{cat}(\mu,\mu+\Delta\mu)$ cannot be bounded when some $\mu_i$ is close to zero. An alternative theoretical benefit of SFM is its connection to natural gradient, as we have mentioned in the common rebuttal. We also noted that similar non-flat assumptions can be found in previous work including DirichletFM which follows the specific path of Dirichlet distributions and achieves better performance than LinearFM. Again, we sincerely thank you for bringing in such an inspirational discussion and we hope our explanation addresses your questions.
Sorry but I'm not convinced yet about the overshooting justification. Do you maybe have some empirical evidence about this? About Fig. 1, I have the suspicion that it deforms the sphere quadrant nontrivially, in order to plot it to a simplex, and that deformation does not preserve angles and gives the wrong visual impression. Precisely: for geodesics starting from the center, in the case of the simplex, the angles geodesics form with the boundary are in the range $[\pi/6, \pi/2]$ and for the sphere quadrant they are in the range $[\pi/4, \pi/2]$. Angles in the second case are higher than in the first case, so geodesics are less parallel to the boundary for the sphere, going against the "overshooting becomes better" argument. Correct me if I got the intervals wrong. They come from the following observation: angles of a "spherical triangle" are of $\pi/2$ so when a geodesic goes from the center to a vertex (which is the geodesic making the smallest angle with the boundary), it makes an angle of half of $\pi/2$, namely $\pi/4$. The same computation for the "flat simplex" (an equilateral triangle) gives angle of $\pi/6$. Another point, maybe this helps for some clarification: I don't follow your computation about "sphere quadrant distance blowing up", I think it doesn't blow up. $2\mathrm{arccos}(\langle x,y\rangle)$ (formula (7) in the paper) is twice the angle between vectors $x,y\in \mathbb S^{d-1}_2$ (sphere of radius $2$) so this says that you are actually working with a sphere with usual geodesic distance (this distance is radius=2, times angle). On the other hand, you wrote $\mathrm{arccos}(\sqrt{\langle x,y\rangle})$ in your above computation. Is there a mismatch? Of course for small angle $\epsilon>0$ we have "distance $\approx$ radius times angle", but if you put the square root, you get the quantity looking like $\sqrt{\epsilon}/ \epsilon$, and that becomes arbitrarily high as $\epsilon\to 0$, but I don't see where the square root comes from. I'm still convinced that you work with a statistical manifold isometric to sphere quadrant, that the curvature of the sphere slightly helps. I think that the distortion between that and a simplex is controlled, a factor between 1 and 2, which still gives a nontrivial increase in performance. Can this be accurate? I'm happy to discuss, if you have time to discuss more on this. It's not a big issue with the paper, but just want to get to the bottom of it. Please give higher priority to other rebuttals if you want.
We thank your prompt reply and we are happy to continue this inspirational discussion with you. We first want to assure you that Fig.1 was **accurately plotted using the exponential and logarithm maps of the Riemannian simplex equipped with Fisher-Rao metric directly** (Eq.22 & 23 in Appendix A.3) with geodesic $\mu_t=\exp_{\mu_0}(t\log_{\mu_0}\mu_1)$ (should be equivalent to Eq.3). In other words, this visualization is not made from deforming a sphere quadrant, although such derivation should result in the same outcome. We note that the Riemannian structure can be defined directly on the simplex, we can call it *Riemannian simplex* (equipped with the Fisher information metric, and geodesic as Eq.3), which is different from a *flat simplex* (equipped with the Euclidean metric, as in LinearFM). The diffeomorphism to the unit sphere was introduced later in Eq.5 to address the numerical stability issue. This diffeomorphism is an isometry (up to a constant scaling factor) between the *Riemannian simplex* and the sphere, but NOT between the flat simplex and the sphere. We notice that the examples you mentioned have presumably involved the comparison of angles in different geometries (e.g., "a flat equilateral triangle with straight lines ... geodesic-boundary angle ranging [$\pi/6,\pi/2$]" assumes Euclidean geometry on Euclidean simplex, whereas the sphere angle ranging [$\pi/4,\pi/2$] assumes Riemannian geometry on sphere), which in our opinion may not be a very meaningful comparison. Instead, we tried to compare Euclidean angles between different simplexes (Riemannian vs Euclidean simplex) which are more comparable. Note **the geodesic on the Riemannian simplex are not straight lines** but are curved as demonstrated in Fig.1, therefore could make an "Euclidean angle" smaller than $\pi/6$ (i.e. more parallel to the boundary compared to "Euclidean angle" of Euclidean geodesic). If we switch to consider Riemannian angles, it should be preserved during the diffeomorphism between the sphere and Riemannian simplex, but it is not perceptually meaningful to visualize Riemannian angles, and again it makes less sense to us to compare Riemannian angles and Euclidean angles. As for relation to overshooting, our preliminary hypothesis is that, since modern ML optimizers and ODE solvers are mostly developed for Euclidean operators, the curved geometry having a vector field more parallel to the boundary under the "Euclidean angle" definition helps numerical stability. We have demonstrated mathematically that the direction term $\sqrt{\mu\odot\nu}-\langle\sqrt{\mu},\sqrt{\nu}\rangle \mu$, of the vector field will have a close to $u_k\approx 0$ component when $\mu_k\approx 0$, which is different from linear flow matching's fixed $\nu-\mu$. Empirically we did observe that SFM is less sensitive to ODE solver specifications in the ablation study. Last but not least, less overshooting is only one of the potential benefits we hypothesized, and we believe that SFM's capability to employ geometry of statistical manifold and the connection to natural gradient also helps it learn efficiently in the generative model training process. We thank you for encouraging this insightful discussion, and we would love to continue exploring the theoretical foundations of SFM and include more such discussions in the paper. ## Regarding Riemannian geodesic not bounded by Euclidean We hope to emphasize that **our derivation in the response is consistent with our paper**, and it is not about distance on the sphere (Eq.7) but rather Riemannian geodesic (Eq.3) vs Euclidean geodesic on the simplex. To avoid confusion, we use $\mu,\nu\in\Delta^{n-1}$ as probability measures in simplex, and $x,y\in S_+^{n-1}$ as points on the sphere. To make a comparison of the different geodesics, we should use a consistent pair of points over the simplex $\mu,\nu$. In other words, we should either compare $d\_\text{cat}(\mu,\nu)$ with $\\|\mu-\nu\\|\_2,\forall \mu, \nu$, or $d\_S(\pi(\mu),\pi(\nu))$ with $\\|\mu-\nu\\|\_2,\forall \mu, \nu$, and NOT compare $d\_S(x,y)$ with $\\|x-y\\|\_2,\forall x,y\in S\_+^{n-1}$. This is the reason why our calculation was based on naive categorical distance in Eq.3 but not Eq.7 on the unit sphere. We have demonstrated mathematically that the two distance metrics are "diverse" in the sense that there does not exist a finite constant to bound the former with the latter. As we mentioned the isometric property in Proposition 1 that $d\_\text{cat}(\mu,\nu)=\frac{1}{2}d_S(\pi(\mu),\pi(\nu))$, the argument can be applied to the spherical distance (up to a constant of 2 that does not affect the bounding result). In short, we have provided mathematical proof that **such a fixed constant does not exist for bounding** the geodesic distance with the Euclidean distance. We are **not** trying to bound the spherical distance $d\_S(x,y)\le C\\|x-y\\|\_2,\forall x,y\in S\_+^{n-1}$, for which a fixed $C$ indeed existing between 1 and 2.
Ok I see now my mistake/misunderstanding, thanks for insisting and making it clear. I agree that $2 d_S(\pi(\mu), \pi(\nu)) = d_{cat}(\mu,\nu)$, and that $\|\mu - \nu\|_2$ is a high distortion of $d_{cat}(\mu,\nu)$, in particular near the boundary. The point is that map $\pi$ mapping flat to curved simplex, is not bi-Lipschitz at the boundary of the domain, and even though the image is a subdomain of the sphere that is low-distortion to a simplex, the actual identification map $\pi$ isn't bounded distortion. Now I understand the overshooting comment better. I think that in order to be clear, **in the caption of figure 1 you should mention that in the upper row middle picture, do not plot the geodesics on the statistical manifold via exponential map, but rather via the parametrization map $\pi$**. In retrospect I now realize that the fact that you do not use the exponential map had to be evident: the exponential map of the sphere (and thus by isometry, the exp map of the statistical manifold with Fisher-Rao distance) would in fact not map it to a straight triangle as the domain, it would give a curvilinear triangle. But I didn't think twice and the "exponential" in the caption kind of made me misunderstand.
We thank your understanding and we are happy that we have reached a consensus in our discussion. We acknowledge that there might be some confusion regarding the Fig.1 caption and we will add more mathematical details to make it clearer regarding the exponential map. Again, we sincerely thank you for bringing up such an inspirational discussion, and we believe our rebuttal and clarifications have fully addressed your concerns in this regard.
The appreciation for the thoughtful discussion is mutual, thank you for it.
Summary
The paper proposes an approach to generative modelling of discrete data based on the Riemannian Flow Matching [1] algorithm where the authors use the Fisher metric to define the geometry of this space. In detail, the authors define the generative modelling from the discrete data as sampling from the empirical distribution on the statistical manifold (the space of categorical distributions). There every point is a categorical distribution on a fixed state space and the generative model maps the noise distribution to the data distribution on the manifold, i.e. it operates with distributions of categorical distributions. In order to apply Riemannian Flow Matching, the authors define the metric tensor, tangent space, geodesics, logarithmic and exponential maps. Furthermore, for computational stability, the authors map the statistical manifold to the sphere and redefine all the differential geometric tools there. Finally, the authors perform an extensive empirical study applying their method for language, discrete image generation, and sequence modelling in bioinformatics. The method demonstrates competitive results with prior works. [1] Chen, Ricky TQ, and Yaron Lipman. "Riemannian flow matching on general geometries." *arXiv preprint arXiv:2302.03660* (2023).
Strengths
The paper presents a complete study, i.e. it approaches an important problem in the field (generative modelling for discrete data) proposes a reasonable approach, presents the proposed approach in a comprehensible way, and provides minimal necessary empirical study of the idea.
Weaknesses
Although the paper does not raise major concerns and overall satisfies the criteria of a NeurIPS publication, the practical motivation of the method is insufficient. Indeed, the problem of generating categorical distributions is interesting and the paper answers this question definitively (to a reasonable extent). However, then the authors quickly jump to the conclusion that a natural way to model discrete data is to sample a categorical distribution rather than to sample *from* a categorical distribution. Indeed, in Algorithm 2, the authors don’t even mention how they sample these categorical distributions (as I understand they consider every sample to define a point mass on the statistical manifold). First, the transition between sampling categorical distribution (in theory) and sampling discrete data (in practice) should be stated and explained clearly. Second, this transition requires a practical motivation (perhaps an empirical study). Indeed, why would one consider the Fisher-Rao distance to be a natural distance between two sentences in a natural language? Why one would need to consider a sentence to be a categorical distribution on the statistical manifold rather than a sample from such distribution? The paper does a poor job of answering these questions and describing how the proposed method works with discrete data. There are some minor concerns regarding the presentation: - The name Statistical Flow Matching does not match the title of the paper. - Line 29. It is not clear what `assumptions that fail to capitulate` mean. - Line 52. It is not clear what `existing models often fail due to impromptu prior assumptions` means. - Line 105 requires citations. - In lines 112-117, please specify that equipping the manifold with the Fisher-Rao metric. - The sampling notation in Eq. (8) is confusing due to the diffeomorphism applied to a probability density. - There is a typo in line 118. The tangent space is called `the target space`. - Section 3.3 reads as a preliminary section on the natural gradient rather than a connection of the proposed method to the optimization literature.
Questions
I suggest adding practical examples where the data is a categorical distribution that is not concentrated in a single point. This would be a significant motivation for the proposed method.
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
The paper lacks a discussion of the applicability of the proposed abstraction to the generation of discrete data (see Weaknesses).
Summary
This work introduces statistical flow matching, an improved method for discrete generation on the simplex by utilizing the geometry induced by the Fisher information metric. By mapping points on the simplex to points on the sphere, flows wrt the Fisher metric can be efficiently computed. SFM is tested on toy examples in the simplex as well as on binarized MNIST, Text8, and a Promotor design application.
Strengths
- Great figures with above average effort and conceptual clarity that make it very easy to understand the method and theory geometrically. - Simple, but seemingly effective addition to flow matching on data supported on a simplex. - Could be quite significant if it is truly competitive to diffusion-based approaches.
Weaknesses
- I’m quite concerned about the “Pseudo log likelihood” metric being compared to log likelihood. While the authors know and say it is incomparable, it is for some reason still compared in Tables 1 and 2 directly to the discrete NLL. I don’t understand why this is done. I might recommend an evaluation similar to MultiFlow for text8, where the evaluation is done on generated samples, and can be used to compare fully discrete, vs. the continuous state approach taken here. Its extremely concerning how much the pseudo-NLL changes based on the choice of a seemingly arbitrary threshold $t_{max}$ and the number of steps (Table 4). To me that means this score should either not be used, or should at most only be used to compare like methods. - Qualitatively, the text8 generations look extremely poor relative to other recent related work. I’m very uncertain that “SFM achieved the second-best performance among all diffusion or flow baselines”. It would be great to include qualitative samples from other established methods for comparison. I think it is essential to establish the performance of SFM on the text8 dataset since this is arguably the only non-toy dataset and most established benchmark. If this text8 baseline can be fixed and these results shown to hold on valid and comparable metrics, I would substantially increase my score. Even if the text8 result is not as exciting, as long as it improves over existing flow-based approaches (Including MultiFlow), I believe this paper could still be of interest to the community, and I would still consider raising my score. MultiFlow: https://arxiv.org/pdf/2402.04997
Questions
Questions: - I don’t understand why the change of measure term is undefined at the boundary. Isn’t the transformation the square or sqrt transformation? Why is this undefined at the boundary? Could the authors elaborate on this point? - I don’t understand the “forward noising process” $q(\tilde{\mu}_1 | \delta)$ as far as I can tell this is not defined in this work. - Table 4,5 “ODE solver” is what solver? Euler is an “ODE solver”. Could the authors specify which solver was used? - Why is LinearFM quite different between 300 steps and “ODE Solver” in Table 5? - Table 5 does not look like an FID score to me, which must be strictly positive. Perhaps this is an NLL? - Why is Multiflow cited but not compared to in the text8 task? Minor comments: - “we define the ODE log-likelihood as” — This is more accurately described as a change in likelihood perhaps? - “Swill roll” —> “Swiss roll” perhaps is meant in multiple places?
Rating
7
Confidence
4
Soundness
3
Presentation
4
Contribution
4
Limitations
Yes
Summary
This paper proposes flow matching on the manifold of discrete distributions, where each point represents a probability mass function (pmf). In essence, this is done by parameterizing a vector of size n for each n-class categorical distribution. Instead of using Euclidean geodesic,s the authors discuss the use of Riemannian geodesics given the Fisher information matrix. To address the constraint that this should represent a probability mass function, the authors suggest using constraint ||x||^2 = 1 (i.e. the normalization is over L2), which allows geodesics to be well-defined on when points lie on the boundary. Furthermore, since this is within a continuous normalizing flow framework, authors claim that it is possible to compute exact log-likelihoods, as opposed to related diffusion model counterparts. Experiments are carried out over a range of datasets where n=2, 26, 4.
Strengths
- Well written. I found the writing clear and to the point. - A natural extension of existing works to flow matching on the manifold of discrete pmfs.
Weaknesses
Overall, I feel the idea is a straightforward adoption of prior work on riemannian FM in terms of novelty, but the paper is nicely presented and packaged together. The topic of generative flows with discrete data is also timely.
Questions
### Main questions / concerns: _Regarding likelihood definition._ In Sec 2.1, you use Radon-Nikodym derivative $p=\frac{d\mu}{d\nu}$ to define the density p, where $\nu$ is used to denote the reference measure. In Sec 3.5, you state the change of measure as if it is a normal density function rather than treat it as a Radon-Nikodym derivative. - How do you chose $\nu$? - Can you say that the Radon-Nikodym derivative is defined for all $\mu$ that you use in practice? - Does this change of measure (Eq 11) correspond to some $\frac{d\mu_t}{d\nu}$, i.e., fixed reference measure? I feel that even though $p$ is defined through a Radon-Nikodym derivative, there isn't actually a formal treatment of $p$. In the text, $p$ seems to just be treated as a density function. This should be okay specifically because $x$ lies in a finite dimensional space, but this special case of the Radon-Nikodym derivative should be stated. _Regarding choice of metric in likelihood computation._ In Tables 1 & 2 involving NLL values, are you sure these values are comparable between methods? For instance, the NLL depends on the choice of metric g (as in Eq 11), so if the $div_g$ is replaced with Euclidean divergence, this will have very different NLL values. I see that in both tables, LinearFM is directly compared to SFM. Even if different g's are used during training, they can still be compared if using the same g for NLL computation. - Did you use the same g for NLL computation? - Which g did you use (Euclidean or Riemannian)? _Regarding choice of metric in desiging conditional u_t._ In Figure 1, it might be a good idea to also visualize the velocity field u_t. I think there's also an interesting property of the Riemannian u_t that isn't discussed, and it's that this u_t is parallel to the boundary. Whereas the Euclidean u_t will have a nonzero inner product between the u_t and the normal direction when evaluated on / close to the boundary. Is this true? _Regarding the connection to natural gradient._ It's unclear what the point of Section 3.3 is. Right now, this section feels pointless as it doesn't say anything about the proposed method. What your point is should be made more explicit. - Why should we care if the Fisher information matrix shows up as the Hessian of KL? What does this imply? - Can you show that the geodesic under this Riemannian metric is following the "natural gradient"? Is this what you wanted to claim here? ### Minor suggestions: - State in the main text that you are using minibatch OT couplings. Sec 3.4 does not mention how OT is used in the actual training algorithm. (I had assumed it is done over the full training set until finding Alg 2.) - The "target" space --> "tangent" in Sec 3.1.
Rating
7
Confidence
4
Soundness
4
Presentation
4
Contribution
3
Limitations
Authors adequately addressed limitations.
Alternative BPC Calculation
We hope our previous rebuttal has adequately addressed your questions. To further support our discussion, we would like to introduce a potentially more comparable BPC formulation inspired by a new concurrent work MDLM [1]. MDLM provided a general form of ELBO for the continuous flow setting with a more formal proof. This ELBO applies to both continuous flow models and flow models that rely on discrete jumps with logits. Specifically, the ELBO has the form: $$ \mathcal{L}=-\mathbb{E}_{\mu_1\sim q(\mu)}\left[\int_0^1\frac{1}{1-t}\log\langle \mu_t,\mu_1\rangle dt\right] $$ where $\mu_t$ follows the predicted inverse trajectory starting from $\mu_1$. We believe BPC calculated with this ELBO is comparable to SEDD/D3PM/Multiflow (as they were compared in [1] as well). The BPC based on this variational bound is provided in the following table. Our BPCs were averaged over the first 1000 sequences of length 256 in the test set. | Model | BPC↓ | | -------------- | ------------- | | SFM w/ OT | 1.412 ± 0.006 | | SFM w/o OT | 1.410 ± 0.004 | | LinearFM | 2.197 ± 0.008 | | D3PM-absorb | 1.45 | | D3PM-uniform | 1.61 | | BFN | 1.41 | | SEDD-absorb | 1.32 | | SEDD-uniform | 1.41 | | Discrete Flow | 1.23 | | Argmax Flow | 1.80 | | MultiFlow η=0 | 1.41 | | Transformer XL | 1.08 | Similar to our evaluation results using GPT-J-6B, **our SFM reached a comparable BPC with MultiFlow, BFN and SEDD-uniform, outperformed many existing discrete diffusion/flow models including D3PM, LinearFM, and Argmax flow**. We will update the results in the revised manuscript as such BPCs are more comparable than pseudo-BPCs. We sincerely hope that our discussion addressed all your concerns and we look forward to further discussion. [1] Sahoo, Subham Sekhar, et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).
I thank the authors for their extensive additional evaluations on text8 and detailed responses. I believe the evaluations are substantially improved and give a much clearer picture of the performance of various methods. I particularly appreciate the inclusion of the new ELBO from the concurrent MDLM work. I believe this is much better than what was suggested in MultiFlow, which has its problems that the authors have pointed out. While SFM is not the best method for this task, and probably not where the community will settle given other concurrent work. I am satisfied with its performance and appreciate the contribution overall. I thank the authors for their hard work and raise my score 3 --> 7.
We thank your recognition of our work and we are happy that our rebuttal has fully addressed your questions. Again, we thank your insightful review regarding the evaluation of the Text8 dataset that helps make our work more concrete and comprehensive. We will make sure to include these new evaluation results of both the GPT-J-6B NLLs and the more solid and comparable BPCs from the MDLM paper in our revised manuscript.
the motivation is still unconvincing but I consider the method to be important
Thank you for the response. The presentation of the proposed method as the right way to generate text remains unconvincing to me. I'm not claiming that this is the wrong approach, I feel that the paper does a poor job of convincing the reader in this. However, I recognize the significance of generative modeling on the manifold of distributions equipped with the Fisher-Rao metric (which I'm well aware of). I agree, that the Fisher-Rao metric is a fundamental concept, and that's why I think it doesn't require further motivation through the discussion of the natural gradient.
We first thank your recognition of the significance of our SFM for generative modeling on the manifold of distributions and for approving the validity of our approach. We would like to further clarify why our choice of modeling discrete data as one-hot probabilities is natural and widely adopted in prior work on discrete flow/diffusion. We will make sure to improve the arrangement of sections to provide more background to the audience earlier in the text. As described in our rebuttal, previous diffusion and flow models also predominantly viewed discrete data as a distribution, and the generative task for these models effectively learns a distribution over distributions. For example, DDSM and Dirichlet FM both assumed the discrete data as probabilities following some Dirichlet distribution (points on the simplex). D3PM, MultiFlow, and SEDD viewed each token as continuous logits (points on the probability simplex after softmax). We provided extended discussions about this setting in Sec.5 in our original paper. Similarly, our SFM views tokens as points on the simplex, differing only by providing a more mathematically meaningful geometry for the probability simplex. Regarding the "right way" to generate text, to the best of our knowledge, **all current discrete diffusion and flow models for natural language modeling were based on the probability path on the simplex or related to the path on the simplex using logits**. In this sense, we believe that working with simplex is more of the "natural" and "standard" way to the diffusion community following existing work. We thank your insightful comment, which reminds us that such an assumption may diverge from that in other text modeling communities. We provided an explanation to the possible gaps and misunderstandings in our rebuttal, and we will add them to our revised manuscript to make it more accessible to general audiences with different backgrounds. We would also appreciate it if you could provide pointers to works that do not view text tokens as one-hot categorical distributions (e.g., do not use cross-entropy loss), and we will include discussions on them. Again, we sincerely thank your review and response and we hope we have fully addressed your questions and comments. We will better organize our paper in the revised manuscript.
Sincere Gratitude from Authors
Dear Reviewers, We sincerely thank you for your insightful and constructive feedback throughout the review process. Your thoughtful comments have helped make our work more concrete and comprehensive. We are glad that our rebuttals have successfully addressed your questions and comments, and we are grateful for your unanimous recognition of our efforts and contributions in our proposed SFM framework. We especially enjoyed the interactive discussions during the rebuttal period, which has been inspirational and fruitful for us. We will make sure to add the additional evaluation metrics on Text8, clarify some mathematical notations and details, and include additional insights inspired by our discussions in our revised manuscript. Again, we thank you for your valuable reviews and for helping us bring this research to fruition. Warm Regards, Authors.
Decision
Accept (poster)