Categorical Flow Matching on Statistical Manifolds

Authors: chaoran  cheng
jiahan  li
jian  peng
ge  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-05 20:52:39

Abstract: 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.

Main paper: 
 introduction: 
 Recently, conditional flow matching (CFM) models [37] have achieved remarkable success in various generative domains including image generation [37, 17, 28] , molecule [58, 57, 30] and protein design [67, 10, 35] , and sequence generation [59, 38, 12] . While attempts to generalize CFM and diffusion models to discrete categorical data have been made, they typically exert ad hoc assumptions on the structure of the discrete distribution. One group of work relies on stochastic jumps of Markov chains in either the discrete-time [6, 38, 1] or continuous-time setting [11, 54, 12] that discards the continuous nature of the underlying categorical distributions. Other work directly works with the probability simplex [7, 59] or the corresponding logit space [27, 39, 23] with potentially imperfect assumptions that fail to capture the underlying true geometry of the statistical manifold. Furthermore, likelihood is often approximated by variational bounds in previous discrete generative models due to the lack of tractable exact likelihood.We propose to incorporate the intrinsic geometry of the statistical manifold by viewing categorical data as points on the statistical manifold of categorical distributions. Inspired by the mathematical results from information theory, we utilize the Fisher information metric [50] to naturally equip such a manifold with a Riemannian structure and develop an efficient generative training scheme for learning the vector fields without stability issues. We summarize our contributions as the following:(1) We propose Statistical Flow Matching (SFM), a novel and mathematically rigorous generative framework on the manifold of parameterized probability measures. SFM does not pose any prior assumptions on the statistical manifold but instead deduces its intrinsic geometry via mathematical tools. To tackle the discrete generation problem, we instantiate SFM on the manifold of categorical distributions. We deduce closed-form exponential and logarithm maps and develop an efficient flowmatching training algorithm that avoids numerical issues by leveraging a diffeomorphism between manifolds. SFM effectively leverages the intrinsic geometric properties by following the shortest paths of geodesics between the noise and target distributions on the statistical manifold.(2) Our distinctive geometric perspective of the statistical manifold allows us to further apply optimal transport during training and derive tractable exact likelihood for any given sample of probability measure, both of which are unachievable for most existing methods. We also introduce new theoretical insights by establishing connections among Riemannian flow matching, information geometry, and natural gradient descent, which allows us to interpret SFM as following the steepest descent of the natural gradient from the optimization angle.(3) We demonstrated with a toy example on simplex that SFM can learn more complex patterns on the statistical manifold where existing models often fail due to impromptu prior assumptions. We further conducted extensive experiments on diverse real-world discrete generation tasks involving computer vision, natural language processing, and bioinformatics. SFM consistently outperformed existing diffusion or flow-based models and also achieved comparable results with autoregressive models on character-level generation. 
information geometry: 
 In this work, we are interested in learning a parameterized family of probability measures. It is known from information theory that all probability measures over the sample space form the structure known as statistical manifold. Mathematically, consider probability densities p = dµ dν : X → R defined by the Radon-Nikodym derivative where µ is a probability measure on the sample space X and ν is the reference measure on X . Suppose the statistical manifold P = P(X ) = {p : X dµ = X p(x; θ)dν = 1} is parameterized by θ = (θ 1 , θ 2 , . . . , θ n ) ∈ Θ, this parameterization naturally provides a coordinate system for P on which each point is a probability measure µ with the corresponding probability density function p(x; θ). The Fisher information metric is defined asg jk (θ) = E X ∂ log p(X; θ) ∂θ j ∂ log p(X; θ) ∂θ k = X ∂ log p(x; θ) ∂θ j ∂ log p(x; θ) ∂θ k p(x; θ)dν (1)Rao demonstrated in his seminal paper [50] that statistical manifold can be equipped with the Fisher information metric to obtain a Riemannian structure, the study of which is known as information geometry [5, 4, 8] . This geometric view of statistical manifolds allows us to derive key geometric concepts for our statistical flow matching framework. For example, a geodesic γ(t) : [0, 1] → P defines a "shortest" path (under the Riemannian metric) connecting two probability measures on the statistical manifold. The geodesic distance between two probability measures, also known as the Fisher-Rao distance [50] , measures the similarity between them. The tangent space T µ (P) at a point µ ∈ P can be naturally identified with the affine subspace T µ (P) = {ν| X dν = 0} where each element ν is a signed measure over X . The exponential map exp µ : T µ (P) → P and logarithm map log µ : P → T µ (P) can also be defined on the statistical manifold. While the geodesic for a parameterized family of probability measures can be obtained numerically by solving the geodesic equation when closed-form expression is unknown (see Appendix A.1), it usually requires expensive simulations. Fortunately, closed-form geodesic distances are available for many common distributions including categorical, multinomial, and normal distributions [43] , which motivates our method.The conditional flow matching (CFM) framework [37] provides a simple yet powerful approach to generative modeling by learning a time-dependent vector field that pushes the prior noise distribution to any target data distribution. Such a flow-based model can be viewed as the continuous generalization of the score matching (diffusion) model [55, 56, 25] while allowing for a more flexible design of the denoising process. The Riemannian flow matching framework [14] further extends CFM to general manifolds on which a well-defined distance metric can be efficiently computed.Consider a smooth Riemannian manifold M with the Riemannian metric g, a probability path p t : [0, 1] → P(M) is a curve of probability densities over M. A flow ψ t : [0, 1] × M → M is a timedependent diffeomorphism defined by a time-dependent vector field u t : [0, 1] × M → T M via the ordinary differential equation (ODE): d dt ψ t (x) = u t (ψ t (x)). The flow matching objective directly regresses the vector field u t with a time-dependent neural net v(x t , t) where x t := ψ t (x). However, this objective is generally intractable. Both [37, 14] demonstrated that a tractable objective can be derived by conditioning on the target data x 1 at t = 1 of the probability path. The Riemannian flow matching objective can be formulated as [14] EQUATION)where q is the data distribution, p 0 is the prior distribution, and x t := ψ t (x|x 0 , x 1 ) denotes the conditional flow. [14] further demonstrated that if the exponential map and logarithm map can be evaluated in closed-form, the condition flow can be defined as x t = exp x0 (t log x0 x 1 ), and the corresponding vector field can be calculated asu t (x t |x 0 , x 1 ) = d dt x t = log xt (x 1 )/(1 -t).We also adapt this formulation for our statistical flow, which we will elaborate on in the next section. We note that, since we are working with manifolds of probability measures M = P(X ), we will use p, q for probability densities over the manifold P(X ) and µ, ν for probability measures (or probability masses for discrete distributions) on the manifold P(X ) to avoid confusion. 
method: 
 Different from previous work that treated each distribution separately, we adopt an integral viewpoint toward the manifold of probability distributions. In this section, we focus on the statistical manifold of categorical distributions to demonstrate the application of our method on discrete generation tasks. However, we emphasize that our proposed SFM is applicable to any statistical manifold with a closed-form geodesic distance and can be broadly used in generative tasks targeting probability measures on the statistical manifold.Consider the discrete sample space X = {1, 2, . . . , n}, an n-class categorical distribution over X can be parameterized by n parameters µ 1 , µ 2 , . . . , µ n such that n i=1 µ i = 1, µ i ≥ 0. In this way, the reference measure ν is the counting measure and the probability measure µ can be written as the convex combination of the canonical basis of Dirac measures {δ i } n i=1 over X : µ = n i=1 µ i δ i . Geometrically, this manifold can be visualized as the (n -1)-dimensional simplex ∆ n-1 . The geodesic distance between two categorical distributions is given in [50, 43] asd cat (µ, ν) = 2 arccos n i=1 √ µ i ν i (3)The tangent space at a point µ can be identified with T µ (P) = {u| n i=1 u i = 0} and the corresponding inner product at point µ is defined asEQUATIONwhere P • denotes the interior of the statistical manifold P. Note that the inner product is ill-defined on the boundary, causing numerical issues near the boundary. To circumvent this issue, we introduce the following diffeomorphismπ : P → S n-1 + , µ i → x i = √ µ i (5)which maps the original statistical manifold to the positive orthant of a unit (n -1)-sphere S n-1+ = {x| n i=1 x 2 i = 1, x i ≥ 0} (see Fig.2). Note that we have π(µ) 2 2 = n i=1 √ µ i 2 = n i=1 µ i = 1.The geodesic on S n-1 + follows the great circle, and the following proposition holds between the geodesic distance on S n-1 + and P:EQUATION)A proof is provided in Appendix A.3). This indicates that we can work with the geodesic distance on the unit sphere instead with up to a constant factor:d S (x, y) = arccos( x, y ), x, y ∈ S n-1 + (7)The We provide the analytical form of the exponential and logarithm maps on the statistical manifold of categorical distributions in Appendix A.3. Although it is possible to directly learn the vector field following the loss in Eq.( 2), such a direct parameterization has numerical issues near the boundary.As described in Sec.3.1, we apply the diffeomorphism π in Eq.( 5) to derive a more stable training objective on the spherical manifold:EQUATION)where p 0 is the prior noise distribution over P and q is the data distribution; π * denotes the standard pushforward measure. v :S n-1 + × [0, 1] → T S n-1+ is a learnable time-dependent vector field network that maps the interpolant x t on the unit sphere to a tangent vector in the tangent bundle T S n-1 + . The ground truth conditional vector field u S is calculated using the exponential and logarithms maps on the sphere (Appendix A.2). The overall training and inference scheme is visualized in Fig. 2 and described in Alg.2 and 3 in Appendix C.We further implement a Euclidean flow matching model on the probability simplex with linear interpolation between the noises and the target distributions. Though linear flow matching offers "straight" lines under the Euclidean assumption, it is unaware of the intrinsic geometric properties of the statistical manifold and turns out to trace longer paths in terms of the Riemannian metric. The objective for linear flow matching can be described asL LinearFM = E t∼U [0,1],µ0∼p0(µ),µ1∼q(µ) v(µ t , t) -(µ 1 -µ 0 ) 2 2 (9)In the discussion above, we assume a single data dimension on the statistical manifold. This can be extended to any data dimension by treating them as independent channels of the input. In practice, each probability measure on the simplex is represented by a matrix X ∈ [0, 1] D×n where D is the data dimension and each row sums to 1. The priors are independently sampled from the uniform distribution over the (n -1)-simplex and each data dimension is independently interpolated with the same timestep t ∼ U [0, 1]. The flow model, on the other hand, takes all the data dimensions as input to capture the dependence between different dimensions. During sampling, existing ODE solvers and the simple Euler method are used to integrate the vector field through time to obtain the final categorical distributions. Discrete samples are then drawn from the generated categorical distributions for evaluation. Details regarding model parameterization and sampling are further described in Appendix C.1 and C.2. We further provide an interpretation of our proposed statistical flow matching framework from the perspective of an optimization problem. From the optimization viewpoint, for a generative model on a statistical manifold, we want to minimize the discrepancy between the target distributions and the generated distributions. Naturally, the Kullback-Leibler divergence D KL (µ(θ) µ(θ 1 )) can be used as a measure of statistical distance. We note that the Fisher information metric can be obtained as the Hessian of the KL divergence with respect to the parameter θ. This can be demonstrated by the fact that the KL divergence reaches the global minimum of 0 at θ = θ 1 , so all first-order partial derivatives are zero and the Hessian is positive semi-definite. Therefore, Taylor expansion of KL divergence at θ 1 with ∆θ = θ -θ 1 givesEQUATION)From the geometric viewpoint, the geodesic, by definition, is a (locally) length-minimizing curve with respect to the corresponding Riemannian metric. Therefore, by following the direction of the vector field that decreases the geodesic element ds 2 = dθ 2 g ≈ ∆θ 2 g , we are indeed following the steepest direction that minimizes the local KL divergence. In this sense, the Fisher information metric defines a second-order optimization scheme for the KL divergence by following the "steepest" direction of the Hessian. Indeed, The steepest direction ∆θ that decreases the KL divergence is known as the natural gradient in the existing literature [2, 3] and has been explored in optimization known as natural gradient descent [47, 41, 45] . Instead of optimizing along the normal gradient, the natural gradient is defined as ∇L = F -1 ∇L where F is the Fisher information matrix estimated from a batch of sampled data. Results established in [2, 41] demonstrated that stochastic natural gradient descent is asymptotically "Fisher efficient" and is indeed the steepest descent direction in the manifold of distributions where distance is measured in small local neighborhoods by the KL divergence. Following the geodesic defined by the Fisher information metric, our SFM framework also shares these benefits with an additional advantage of analytical expressions for geodesics, as we focus on the family of categorical distributions instead of general statistical models. Such a theoretical connection may contribute to the better performance of SFM.The geometric viewpoint of the statistical manifold offers a continuous and differentiable formulation of generative modeling and also provides a robust way to measure the distance between two categorical distributions via the well-defined geodesic distance in Eq.(3). In contrast, the Markov chain-based methods cannot establish a robust distance measure due to the stochastic jumps between discrete states. Inspired by the optimal transport formulation in previous work [44, 20, 67, 62] , we naturally extend it to our statistical setting in which the cost is defined by averaging the statistical distances over data dimensions asd cat (X, Y ) = 1 D D k=1 d cat (µ (k) , ν (k) ) for X = {µ (k) } D k=1 , Y = {ν (k) } D k=1 .An optimal assignment of the initial noises to the target data distributions can potentially lead to more efficient training, which we demonstrated empirically in our ablation studies.Unlike diffusion-based models which rely on variational lower bounds for likelihood estimation, our proposed method shares the continuous normalizing flow's capability of exact likelihood calculation. For an arbitrary test sample x ∈ M, using the change of measure formula, the likelihood can be modeled by the continuity equation [14, 42] :EQUATIONwhere div g is the Riemannian divergence and v t (x t ) := v(x t , t) is the time-dependent vector field. In this way, the pushforward probability measures p t (x t ) defined via the learned flow can be obtained as the integral of the Riemannian divergence back through time on the simulated trajectory x t . Following [9, 7] , we define the ODE log-likelihood as the change of the log-likelihood as:EQUATIONwhere the trajectory x t is obtained by solving the differential equation ∂ ∂t x t = v t (x t ) reverse through time with the initial condition x 1 at t = 1. In this way, we have log p(x 1 ) = log p ODE + log p 0 (x 0 ) where log p 0 (x 0 ) is the log-likehood of the prior distribution at data point x 0 . In practice, we follow previous work [7] to use Hutchinson's trace estimator [29] to efficiently obtain an unbiased estimation of the divergence using standard normal random variables. To further account for the transform π from P to S n-1 + and the reverse transform π -1 , additional change of measure formulae need to be applied. Consider the pushforward of the probability measure P over P under the diffeomorphism π defined by π * P (V ) := P (π -1 V ), we have the change of measure identity π * P (π(µ))| det dπ(µ)| = P (x) for x = π(µ). Therefore, by adding the two log-determinant of the pushforward measures, the log-likelihood can be formulated asEQUATIONThe above formula is well-defined for all interior points of P, but the change of measure terms are undefined on the boundary. This can be understood as the fact that discrete likelihoods can be arbitrarily high [61] . Following [7] , we derive a variational lower bound for the likelihood as the marginalized probability over the small neighborhood of a Dirac measure µ 1 = δ at the boundary:EQUATIONwhere q(μ 1 |δ) can be viewed as the forward noising probability at μ1 which is close to δ with a closed-form likelihood. p(δ|μ 1 ) is the categorical likelihood (cross-entropy) and p 1 (μ 1 ) is the model likelihood in Eq.( 13). The overall workflow of calculating NLL is demonstrated in Fig. 2 , and more details regarding likelihood computation can be found in Appendix B.1. It is worth mentioning that the continuous likelihood calculated here is defined with respect to the probability distribution over the statistical manifold P. In contrast, the bits-per-dimension score for autoregressive models is usually defined with respect to a specific categorical distribution on P and therefore not comparable, as was also noted in [7] . 
experiments: 
 Our proposed SFM framework can be leveraged to approximate arbitrary distribution over P, i.e., any distribution over the parameterized family of probability measures. We first demonstrate this with the toy example of a Swiss roll distribution on the probability simplex. We then conduct extensive experiments on real-world datasets for discrete generation across various domains including computer vision, natural language processing, and bioinformatics to demonstrate the effectiveness of our proposed model. We train and evaluate two different versions of the SFM framework: SFM w/ OT for our proposed model with optimal transport during training and SFM w/o OT for the model without optimal transport. A naive approach that directly works with the exponential and logarithm maps on the statistical manifold without the diffeomorphism (SFM w/o π) is also compared in the toy example. We also implement a linear flow matching model on the probability simplex using the loss in Eq.( 9) as an additional baseline, for which we dub LinearFM. For discrete data, we always use Eq.( 14) to obtain finite negative log-likelihood (NLL). More details regarding model architectures and the experimental setup can be found in Appendix C.We project the Swiss roll dataset onto the 2-simplex with some additional margins to make sure no point resides on the boundary. We used a time-dependent MLP to model the vector field for all models. The generative points on the simplex and NLL are calculated using the Dopri5 ODE solver [19] and are shown in Fig. 3 . We note that most existing models assume the target data to be Dirac measures (one-hot distributions) and cannot be applied to this task. Both DDSM [7] and DirichletFM [59] imposed strong prior assumptions on the data distribution as Dirichlet distributions, which failed to capture the complex geometric shape of the Swiss roll data points as demonstrated in the generated samples. On the contrary, directly built upon the geometry of the statistical manifold, our SFM model successfully captured the detailed data geometry. As no data resides on the boundary, exact NLL can be obtained using Eq.( 13) and was averaged over all points in the dataset. Though linear flow matching could also capture the geometry of the data, our SFM outperformed it in terms of NLL.We extend our SFM to model discrete generation tasks in computer vision. The binarized MNIST dataset [52] is the binarized version of the original MNIST dataset [33] by thresholding the original continuous value to be either 0 or 1, thus can be viewed as a 2-class generation task with a data dimension of 28 2 = 784. We also compared D3PM [6] with an additional absorbing state. We used a convolutional neural net adopted from [56] with additional additive Fourier-based time embeddings. The quantitative evaluation of NLL and Fréchet inception distance (FID) are shown in Tab.1. All the generated results and NLL calculations were based on the Dopri5 ODE solver. Additional ablation results can be found in Appendix D.2 and generated images can be found in Appendix D.3.The NLL for diffusion-based models (D3PM and DDSM) was also calculated based on the variational bounds derived in their papers. Our proposed model consistently outperformed both the linear flow matching and other diffusion baselines in terms of FID and NLL, achieving higher sample quality and likelihood.The Text8 dataset [40] is a medium-size character-level corpus consisting of a small vocabulary of 27, which includes the 26 lowercase letters and the whitespace token. We followed the convention in previous work [6, 23] to use random chunks of length 256 in both training and evaluation without any preprocessing. We used a 12-layer diffusion transformer (DiT) [48] based predictor, similar to the one used in [38] . As our exact likelihood is not directly comparable to bits-per-character (BPC) reported in previous work which relies on an alternative variational bound for the likelihood, we additionally calculate such BPC inspired by [51] . See Appendix B.2 for additional details. All generated samples and NLL were estimated using the Dopri5 ODE solver. Quantitative results of BPC are provided in Tab.2 and generated texts are provided in Appendix D.3. The results for the autoregressive language models are also provided as a reference (Discrete Flow [63] is based on autoregressive normalizing flow).Note that, unlike all of the other baselines, SFM and LinearFM do not directly optimize such a BPC as a training objective, so such an evaluation metric is unfavorable for our model (see Appendix B.2). Despite such an unfavorable evaluation, our proposed SFM still achieved the second-best performance among other diffusion and flow baselines that were directly trained with cross-entropy losses.We also noted a significant performance gap between SFM and the naive linear flow matching on simplex, highlighting the importance of capturing the intrinsic geometric properties of the underlying statistical manifold. Optimal transport does not significantly affect the final performance here, which we presume might be due to the long training stage in which vector fields were well-explored. Additional evaluation following [12] is provided in Appendix D.1, in which SFM achieved the best generation entropy as evaluated by the pre-trained GPT-J-6B model [66] .We further apply SFM to the practical task of promoter DNA sequence design in the bioinformatics realm. The promoter is a key element in gene transcription and regulation, and the generation of desired promoters can better help us understand the intricate interactions between human genes. [7] proposed a human promoter sequence dataset containing 100k promoter sequences with the corresponding transcription initiation signal profiles. Each promoter sequence is 1024 base pairs long and is centered at the annotated transcription start site position [26] . Therefore, we can model promoter design as a generation task with 4 categories (the base pair ATGC) conditioned on the given transcription signals. We follow [7] to concatenate the signal as additional input to the vector field predictor built upon 20 stacks of 1-dimensional convolutional layers on the input sequence.Optimal transport can also be applied for conditional generation, as we can pair the conditions with the input to make sure that the conditional arguments align with the target one-hot distributions. 1.47 * BFN [23] 1.41 * SEDD-absorb [38] 1.32 * MultiFlow [12] 1.41 * Argmax Flow [27] 1.80 * Discrete Flow [63] 1.23 * Transformer [6] 1.23 * Transformer XL [16] 1.08 * [15] 0.0395 * Bit-Diffusion (bit) [15] 0.0414 * Language Model 0.0333 † DirichletFM 0.0269 † To provide a quantitative evaluation of the generated promoter sequences, we follow [7] to apply the pre-trained deep learning sequence model Sei [13] to predict active promoters based on predictions of the chromatin mark H3K4me3. As the dataset spans the whole range of human promoter activity levels from highly expressed promoters to those with very low expression, the evaluation is based on comparing the scores between the generated sequences and the ground truth human genome promoter sequences on the test chromosomes. The metric SP-MSE is the mean squared error between the predicted promoter activity of the generated sequences and human genome sequences (lower is better) and is demonstrated in Tab.3. Our SFM was able to achieve the lowest SP-MSE score compared with other baselines. It is worth noting, though, that optimal transport produced slightly sub-optimal results. We hypothesize that this is because, in conditional generative tasks, the final generation should rely heavily on the conditions. Simply matching inputs with the targets using distance measures may discard important information in the conditional arguments and may not be the best practice. 
related work: 
 As we put a special interest in discrete generation, we start with the existing work on discrete generation. We first list a few traditional autoregressive models [16, 49, 63] and masked language models [18, 53] but will skip them as we mainly focus on the diffusion and flow matching models. Existing diffusion and flow-based discrete generative models can be categorized into two main groups. The first group relies on stochastic jumps of Markov chains by treating discrete classes as Markov states. By choosing a proper transition matrix, the target one-hot distribution can be gradually noised into the stationary distribution. Noticeably, this approach can also adopt an additional absorbing state to mimic the mask token in masked language modeling [6] . D3PM [6] adapted the discrete-time Markov chain with the diffusion framework and SEDD [38] proposed a more efficient training scheme by learning the score entropy between different states. [11, 54] extended it to the continuous-time Markov chain by using the infinitesimal generator (rate matrix) instead, and [12] further applied the flow matching framework. Although the transition matrix or the rate matrix determines the entire diffusion trajectory, there is no canonical way of choosing it for different generation tasks to control the mixing rate. Also, due to the presence of discrete jumps of the Markov chain, exact likelihood calculation is no longer feasible. Only variational bounds were derived in [6, 12] . The other group directly works with the probability simplex or the logit space with certain assumptions of the underlying geometric structure [35] . As an example, our linear flow matching assumes a Euclidean geometry on the probability simplex and is often used as a baseline in previous work [59] . The multinomial diffusion [27] assumed a Euclidean geometry on the logits space so interpolation became multiplication back on the probability simplex. Dirichlet diffusion (DDSM) [7] and Dirichlet flow matching (DirichletFM) [59] exerted a strong prior assumption on the probability path as the Jacobi diffusion process. We provide a more detailed analysis of these models in Appendix A. 4 . For models that directly operate on the logit space, the Bayesian flow network (BFN) assumed Gaussian distributions on the logit space with Bayesian update rules. [39] also assumed a Euclidean geometry in the logit space with targets as soft one-hot logits. The assumptions made in these models may not always hold, e.g., for the Swiss roll on simplex dataset in Fig. 3 . These assumptions also did not necessarily capture the true geometry of the underlying statistical manifold. In contrast, our proposed SFM framework does not exert any strong prior assumptions and is aware of the intrinsic geometry by following the geodesics.We also briefly review the related work on statistical manifolds and information geometry. The field of information geometry, the study of geometric properties of statistical manifolds, was first established in Rao's seminal paper in 1945 [50] discussing the Fisher information metric. Most previous work utilizing information geometry focuses on optimization [2, 3] with a few exceptions on representation learning. [64] utilized the geodesic distance between two univariate Gaussians for function shape matching for retrosynthesis of gold nanoparticles. [34] treated point cloud data as probability measures over the 3D Euclidean space and considered the pullback metric on the latent space to obtain optimal latent coordinates for the autoencoder. [46] further applied such a method on single-cell RNA sequence trajectories. These models demonstrated the advantage of following the proper geometry compared with the vanilla Euclidean setting, but the pullback metric needed to be evaluated with expensive Jacobian-vector products during training. In contrast, our proposed SFM, for the first time, leverages mathematical tools from information geometry to generative modeling. SFM directly operates on the statistical manifold with closed-form geodesics, providing a simulationfree approach for efficient generative training. 
discussion: 
 In this paper, we proposed statistical flow matching (SFM) as a general generative framework for generative modeling on the statistical manifold of probability measures. By leveraging results from information geometry, our SFM effectively captures the underlying intrinsic geometric properties of the statistical manifold. Specifically, we focused on the statistical manifold of categorical distributions in this work and derived optimal transport and exact likelihood formulae. We applied SFM to diverse downstream discrete generation tasks across different domains to demonstrate our framework's effectiveness over the baselines. We noted that SFM can be further extended to non-discrete generative tasks whose targets are probability distributions, which we will leave as future work.We are also aware of the limitations of our SFM framework. As a special case of the flow matching model, the generation is an iterative process of refinement that cannot modify the size of the initial input. This may pose limitations to generation compared with autoregressive models. Furthermore, we adopted the assumption of independence between classes so that the canonical Riemannian structure can be induced by the Fisher metric. However, discretized data like CIFAR-10 [31] (256 ordinal pixel values) do not follow this assumption, and results on such data might be suboptimal.In this section, we give a more detailed mathematical background on Riemannian manifold and information geometry related to this work. A more comprehensive background on the Riemannian manifold can be found in [21] . For information geometry, [8, 4] provide more comprehensive analyses and rigorous formulations.A Riemannian manifold M is a real, smooth manifold equipped with a positive definite inner product g on the tangent space T x (M) at each point x ∈ M. Let T M = x∈M T x (M) be the tangent bundle of the manifold M, a time-dependent vector field on M is a mappingu t : [0, 1]×M → T M where u t (x) ∈ T x (M).A geodesic is a locally distance-minimizing curve on the manifold. The existence and the uniqueness of the geodesic state that for any point x ∈ M and for any tangent vector u ∈ T x (M), there exists a unique geodesic γ : [0, 1] → M such that γ(0) = x and γ ′ (0) = u. The exponential map exp : M × T M → M is uniquely defined to be exp x (u) := γ(1).The logarithm map log : M × M → T M is defined as the inverse mapping of the exponential map such that exp x (log x (y)) ≡ y, ∀x, y ∈ M * . With the exponential map and logarithm map, the time-dependent flow can be compactly written as time interpolation along the geodesic:EQUATIONIt can be demonstrated that the above flow indeed traces the geodesic between x 0 , x 1 with linearly decreasing geodesic distance d g (x t , x 1 ) = (1 -t)d g (x 0 , x 1 ) [14] . For an n-dimensional Riemannian manifold, the geodesic can be obtained by solving the geodesic equation (using the Einstein summation convention):EQUATIONwhere x i are local coordinates of the geodesic and Γ k ij are the Christoffel symbols defined byEQUATIONwhere g ij is the inverse metric such that g ij g jk = δ i k . The closed-form expressions for the geodesic are generally not available. [43] curated a wide range of common statistical manifolds for parameterized families of both the discrete and the continuous distributions. We will focus on the statistical manifold of categorical distributions and also the spherical manifold, as the latter is closely related to the former via the diffeomorphism π in Eq.(5).The positive orthant of the unit (n -1)-sphere S n-1 + is a (n -1)-dimensional manifold. The sphere can be embedded into the ambient Euclidean manifold R n so it inherits the canonical inner product asEQUATIONThe tangent space T x (S n-1 + ) = {u| u, x = 0} is a (n -1)-dimensional hyperplane perpendicular to the vector x. The geodesic on the sphere follows the great circle between two points, and the geodesic distance can be calculated in Eq. (7) . The corresponding exponential map can be calculated as:EQUATION) where sinc(θ) = sin θ/θ is the unnormalized sinc function. The logarithm map can be calculated as:EQUATIONwhere P x (w) = w -x, w x is the projection of vector w onto the tangent space T x (S n-1 + ).The pullback of the Fisher metric g defines a metric on M. For u, v ∈ T θ (M), we haveEQUATIONwhich gives the more common definition of the Fisher information matrix in Eq.( 1). This relation also generally holds for a continuous sample space X but requires more careful handling with measure theory tools. We refer interested readers to mathematical references, e.g., Chapter 3.1 in [8] .The statistical manifold of categorical distributions provides an intuitive way of modeling discrete targets. However, we noted that other parameterized distributions can also be used to derive alternative statistical manifolds with different geometries. We will briefly discuss some manifolds that have been used in previous generative models (though not from the information geometry perspective as we do).The Dirichlet distribution Dir(α) is a multivariate generalization of the beta distribution. It is defined on the (n -1)-dimensional simplex ∆ n-1 with n concentration parameters α = {α i } n i=1 . The probability density function is defined asEQUATION)whereB(α) = n i=1 Γ(α i )/Γ( n i=1 α i )is the multivariate beta function. The closed-form expressions for the geodesics between Dirichlet distributions (and for the marginal beta distributions, as were used in [7] ) are unknown. It is known that, however, this statistical manifold has non-constant negative sectional curvatures everywhere [32] . [7] proposed to follow the specific path of Jacobi diffusion processes on the marginal beta distributions, for which expensive modified Jacobi polynomials needed to be precomputed. [59] further applied the flow matching framework by considering the probability path of Dir(1 + te k ) for some fixed maximum timestep. Still, expensive calculations of the regularized incomplete beta functions were required. Compared with these two generative models, our proposed SFM on categorical distributions has the following advantages:• DDSM and DirichletFM have strong prior assumptions that the categorical distributions follow some specific forward noising process which always terminates at (or near) one-hot distributions. This assumption generally holds for discrete generation, but cannot be generalized to more complex geometry on the probability simplex, as we have demonstrated in our toy example of the Swiss roll on simplex dataset in Fig. 3 .• The forward noising paths are not the shortest geodesics in the sense of the Riemannian metric. In our SFM, the vector fields always follow the geodesic, thus following the sharpest direction of decreasing KL divergence (see Sec.3.3).• Our SFM framework does not require the computation of expensive polynomials and is more mathematically concise.• The final one-hot distributions are unreachable in [59] , as it would require an infinite timestep. Therefore, a variational bound needs to be derived for likelihood estimation, which was not done in the original paper.Multinomial Distribution Consider m i.i.d. experiments that follow a categorical distribution with n possible outcomes and probability µ = {µ i } n i=1 , a multinomial distribution gives the proba-bility of getting x i times the i-th outcome with n i=1 x i = m. The geodesic distance for multinomial distributions is identical to that of the categorical distributions up to a scaling constant [43] :d mul (µ, ν) = √ nd cat (µ, ν) = 2 √ n arccos n i=1 √ µ i ν i (27)Therefore, our model can also be interpreted as multinomial flow matching on the corresponding statistical manifold. The previous work of multinomial diffusion (argmax flow) [27] transformed the one-hot distribution to ne k -1 in the logit space (soft one-hot logits) and assumed the common Euclidean structure of the logit space for constructing the diffusion process. This assumption did not consider the intrinsic geometry of the underlying statistical manifold and also had the inaccessibility issue of the Dirac measure as DirichletFM. We have demonstrated that SFM consistently outperformed the naive multinomial diffusion on the Text8 datasets, indicating the effectiveness of tracing the true geodesics.In this section, we provide more concrete mathematical details regarding the exact likelihood calculation in the main text. We first make a distinction between negative log-likelihood (NLL) and bits-per-dimension (BPD, including bits-per-character, BPC) in this work. As we are dealing with probability distributions over the statistical manifold of probability measures, it is important for us to keep the difference between the probability spaces P(P(X )) and P(X ). In our notation, we use µ, ν to denote elements in P(X ), or categorical distributions, and use p, q to denote elements in P(P(X )), or distributions over categorical distributions. We further use NLL to refer to -log p(µ) for a specific µ and BPD for categorical likelihood -log 2 p(δ|µ) where δ is the ground truth Dirac measure (one-hot distribution). Note this definition of BPD (BPC) follows the convention in autoregressive language models where the conditional probabilities for the next token are calculated and averaged [16, 22] . In contrast, most previous work drew equivalence between these two concepts.Our SFM framework enjoys the exact negative log-likelihood calculation as a continuous normalizing flow model [37, 14] . It is worth noting that negative NLLs are not uncommon [14] . As a probability density, the NLL can indeed go to negative infinity at the boundary. Consider the arcsine distribution defined on the finite support [0, 1] with the probability density functionEQUATIONIt is easy to see the density approaches infinity as θ approaches 0 or 1. Note that the arcsine distribution is a special case of beta distribution Beta( 1 2 , 1 2 ), which is in turn a special case of Dirichlet distribution Dir( 12 , 1 2 ). Each θ can be naturally viewed as the parameter for the Bernoulli distribution (2-class categorical distribution), so the arcsine distribution indeed defines a distribution over Bernoulli distributions with negative NLLs near the boundary. In fact, for discrete data, the target data always come in as Dirac measures, which indicates that the target distribution must be a convex combination of Dirac measures with infinite likelihood: q = n i=1 θ i δ i where δ i := δ δ i is the Dirac measure over the underlying Dirac measure δ i on the discrete class i.We followed [7] to derive a variational bound as marginalized over a small neighborhood of the Dirac measure in Eq.( 14) to obtain finite NLLs. The variational bound can be derived via the standard variational approach aslog p(δ) = E q(µ|δ) [log p(δ)] = E q(µ|δ) log p(δ, µ) p(µ|δ) = E q(µ|δ) log p(δ, µ) q(µ|δ) q(δ, µ) p(µ|δ)= E q(µ|δ) log p(δ, µ) q(µ|δ) + D KL (q(µ|δ) p(µ|δ))≥ E q(µ|δ) log p(δ, µ) q(µ|δ) = E q(µ|δ) [log p(δ|µ) + log p(µ) -log q(µ|δ)] (29)The posterior probability q(µ|δ) serves as the forward diffusion process in [6, 59] with a closed-form expression that depends on the timestep. In our Riemannian flow matching formulation, however, such a forward process is implicitly defined via the time-interpolation along the geodesics. Nonetheless, we do know that the conditional distribution should approach the Dirac measure when t → 1 as the geodesic distance approaches 0: p 1 (µ|δ) ≈ δ. Therefore, as we are marginalized over a small neighborhood, we can define a time-dependent posterior distribution q t (µ|δ) that approaches δ as t → 1. In principle, any forward probability that approaches the Dirac measure at t = 1 with an analytical likelihood can be used to derive the lower bound. For example, we can follow [7] to use the Jacobi diffusion process. To avoid the need for expensive Jacobi polynomial calculation, we instead use simple indicator measures over the simplex asEQUATIONwhere p 0 is the base probability of the uniform distribution over the simplex, µ i = θ i is the ith component of the categorical distribution of µ. In other words, q t (µ|δ k ) is the time-interpolated probability between the Dirac measure and the uniform distribution:q t (µ|δ k ) = tδ k +(1-t)U ∆ n-1 .We have log q t (µ|δ k ) = log p 0 -(n -1) log(1 -t) (31) In practice, we use t = 0.995 and provide more results in the ablation studies in Appendix D.2. The categorical likelihood p(δ k |µ) is the standard cross-entropy loss, also referred to as the reconstruction loss in some previous work [23] . It can be calculated asEQUATION)Prior Likelihood The prior likelihood accounts for the base probability when we sample from the prior noise distribution. During both the training and the inference stages, we uniformly sample from the simplex ∆ n-1 . This can be efficiently done by normalizing n i.i.d random variables from the exponential distribution:EQUATIONThe uniform distribution over simplex is exactly the Dirichlet distribution Dir(1, . . . , 1). Therefore, the log-likelihood can be calculated asEQUATION)where Γ is the Gamma function. Note that the prior likelihood is independent of the samples.The diffeomorphism π in Eq.( 5) induces a pushforward measure on S n-1 + which can be described by the change of measure identityEQUATION)The change of measure term can be calculated using the change of volume formula with the canonical coordinates {x i } n-1 i=1 and {ϕ i } n-1 i=1 asEQUATION)where 0 ≤ ϕ i ≤ π/2. The diffeomorphism x = π(µ) is given byEQUATION)The change of measure term isEQUATION)Similarly, for the inverse mapping, the term isEQUATION)Model Likelihood The model likelihood can be viewed as the pushforward measure of the learned flow:p t = (ψ t ) * p 0 (40)We demonstrated in Sec.3.5 that this term can be calculated as the integral of divergence in Eq.( 12) plus the base probability in Eq. (34) . More specifically, the NLL can be obtained as the solution to the following ODE system back through time with the initial condition of x 1 and log p ODE 1 = 0:EQUATIONFor manifolds that can be embedded in the ambient Euclidean space (e.g., simplex and sphere), the Riemannian divergence can be calculated as the normal Euclidean divergence div g = div [14] . We further use Hutchinson's trace estimator [29] to efficiently estimate the divergence asEQUATIONThe expectation can be efficiently computed using the vector-Jacobian product. For data with a unitary dimension D = 1 (e.g., Swiss roll on simplex), we also calculate the exact divergence using the equality div v = Tr(J(v)), where J(v) is the Jacobian of the vector field v. The computational cost for the exact divergence calculation scales quadratically with the data dimension D, as the interdependence between data dimensions makes it computationally expensive to calculate the Jacobian. Therefore, we only demonstrated the NLL with exact divergence calculation for the Swiss roll dataset in Fig. 3 , in which Hutchinson's trace estimator could make a decent estimation of the exact divergence.Combining Eq.( 31), ( 32), ( 34), (38) , and (39) with Eq.( 14) and ( 13), we can efficiently calculate the NLL given arbitrary Dirac measures as input. The NLL calculation scheme is visualized in Fig. 2 and described in Alg.1.We have demonstrated that the NLL can reach negative infinity. In contrast, the common definition of BPC with the cross-entropy loss -log 2 p(δ|µ) is always non-negative and consistent with previous Algorithm 1 NLL Calculation for Discrete Data 1: Sample μ1 ∼ q t (µ|δ) in Eq.( 30) and calculate -log q t (μ 1 |δ) and log p(δ|μ 1 ).2: Apply the diffeomorphism in Eq.( 5) to obtain x1 = π(μ 1 ) and calculate log | det dπ -1 (x 1 )|.3: Solve the ODE system in Eq.( 41) to obtain x 0 and log p ODE . 4: Apply π -1 to obtain µ 0 = π -1 (x 0 ) and calculate log | det dπ(µ 0 )|. 5: Calculate the base log probability log p 0 (µ 0 ). 6: return NLL as in Eq.( 14).autoregressive language models [61] . However, as our flow-based model is not autoregressive, we instead follow the variational formulation in [51] to calculate the alternative NLL and BPC for comparison. The ELBO in [51] can be calculated asEQUATION)where µ t = exp µ0 (t log µ0 µ 1 ), µ 0 ∼ p 0 (µ) denotes the interpolation along the geodesic and μ1 (µ t ) denotes the learned model's prediction for t = 1 given the current noise data µ t . This can be efficiently done by one-step prediction with the learned vector field v asEQUATION)The standard Euclidean inner product is used in Eq.( 43) so it can be understood as a weighted crossentropy loss. Note that the integral in Eq.( 43) has a singularity at t = 1, making it numerically unstable to estimate using ODE solvers. Instead, we follow [51] to apply the change of variable s = -log(1 -t) to reformulate the ELBO asEQUATION)where t = 1 -e -s . Note that for large s, the corresponding t is very close to 1, so the integrand is very close to zero. Indeed, s = 10 corresponds to 1 -t < 5 × 10 -5 , so we simply set the upper limit to 10 and used the Dopri5 solver to numerically estimate the ELBO. BPC can be then calculated as L ELBO / log 2. This formulation also shares a similar form as the other ELBOs derived in [6, 12] , and is thus comparable to most existing models.Eq.( 43) was directly optimized in Markov chain-based methods like D3PM [6] , MultiFlow [12] , and SEDD [38] , and also autoregressive language models. In contrast, our SFM (and LinearFM) follows the continuous normalizing flow setup in [37] , which enables the exact likelihood calculation instead of ELBOs. Therefore, trained on the alternative objective of minimizing the Riemannian vector field norm in Eq.( 8), SFM did not directly try to minimize such an ELBO. Despite such an unfavorable evaluation compared to other baselines, SFM was still able to achieve the second-best BPC, as demonstrated in Tab.2.GPT-J-6B [66] is a transformer-based large language model with 6B trainable parameters. We follow the pipeline in [12] to first tokenize the generated text using the provided tokenizer with GPT-J-6B.The GPT-J-6B NLL is then calculated based on the predicted logits on the token level using the pre-trained model as:EQUATIONwhere w are tokens and K is the number of tokens. Similarly, the entropy is calculated on the empirical distribution of the tokens based on a large number of generated texts. We noted that, as GPT-J-6B was not pre-trained on Text8, its NLL does not necessarily reflect the true data distribution in Text8, as we will demonstrate more concretely in the evaluation plot in Appendix D.1 and generated samples in Appendix D.3.In this section, we further describe the experimental setup, model architecture, and datasets.Our SFM architecture is encoder-agnostic and can be applied to arbitrary discrete generative tasks.Here, we describe the common setting of the flow model across different datasets. We use the sinusoidal timestep embedding described in [65] as Emb : [0, 1] → R H . We follow [14] to manually project the predicted vector field onto the corresponding tangent space. For the spherical manifold, the projection can be described asEQUATION)For linear flow matching on the simplex, the projection can be described asEQUATION)The projection guarantees that the final prediction lies on the corresponding tangent space, and the data points will stay on the manifold during Euler sampling. We will always assume the projection is performed in the following context of using v t . The training stage of SFM is described in Alg.2.The time complexity of our SFM framework is dominated by the underlying vector field predictor with little overhead. Each model for binarized MNIST and promoter design was trained on a single 80GB NVIDIA A100 GPU for 6-10 hours. Each model for Text8 was trained on four 80GB NVIDIA A100 GPUs for about 7 days. We will further describe the encoders for each generation task in the following subsections. Sample noise distribution µ 0 ∼ p 0 (µ) and target distribution µ 1 ∼ q(µ).if optimal transport then 4:Do batch OT assignments of µ 0 and µ 1 according to the average statistical distances.end if6:Apply the diffeomorphism in Eq.( 5) to obtain x 0 = π(µ 0 ), x 1 = π(µ 1 ).Sample t ∼ U [0, 1] and interpolate x t = exp x0 (t log x0 x 1 ) using Eq.( 19) and (20) .8:Calculate the conditional vector field u S t (x t |x 0 , x 1 ) = d dt x t = log xt (x 1 )/(1 -t).Predict the vector field using v(x t , t) and optimize the SFM loss in Eq.( 8). 10: end whileThe sampling process from the trained model can be described as solving the differential equation ∂ ∂t x t = v t (x t ) from t = 0 to 1 with the initial conditional x 0 sampled from the prior noise distribution. Alternatively, we can write the solution as the integral of the learned time-dependent vector through time asEQUATIONWe always use the Dopri5 ODE solver [19] for our ODE sampling and NLL calculation. For Euler method with N discrete steps, the timesteps 0, 1/N, 2/N, . . . , (N -1)/N are used with a step size of 1/N . The sampling stage is described in Alg.3.The Swiss roll on the 2-simplex is generated by normalizing the span of the 2-dimensional Swiss roll to [0 + ε, 1 -ε] where ε is a small margin to make sure no point lies on the boundary. The third dimensional is automatically obtained as µ 3 = 1 -µ 1 -µ 2 . We fixed a random seed and generated 1000 samples for a full batch training. The vector field predictor is based on simple multi-layer perceptions (MLPs) and can be described asEQUATIONwhere denotes concatenation and we used a hidden dimension of H = 128. Each model was trained for 2000 epochs using full batch gradient descent with an initial learning rate of 10 -3 . 1000Algorithm 3 Sampling from SFM 1: Sample noise distribution µ 0 ∼ p 0 (µ).2: Apply the diffeomorphism in Eq.( 5) to obtain x 0 = π(µ 0 ). for t ← 0, 1/N, 2/N, . . . , (N -1)/N do 7:x t+1/N = exp xt (v(x t , t)/N ) 8:end for 9: end if 10: return µ 1 = π -1 (x 1 ) samples were sampled using the Dopri5 ODE solver, and NLL (based on Hutchinson's trace estimator) was calculated on the whole training data with 20 repeats. The DirichletFM model was originally trained with cross-entropy loss which could not be used for this task. We instead used the binary cross-entropy loss during training and kept all the other flow-based part the same as described in the original paper [59] .We used the preprocessed binarized MNIST dataset from [52] which has a split of 50k/10k/10k. We adopted the CNN-based vector field predictor from [56] with additional additive time embeddings at each convolution layer's output. The model has 4 residual layers [24] and 4 refinement layers [36] with a total number of 29.8M trainable parameters. All models were trained for 100k iterations with a batch size of 256 (approximately 510 epochs) with an initial learning rate of 3 × 10 -4 .For the quantitative evaluation of the generated samples, we calculated the FID score for 1000 generated samples for each model. The statistics for the binarized MNIST dataset were calculated on the whole training set via the pre-trained InceptionV3 [60] model. The NLLs reported in the main text were calculated on a fixed subset with 1000 samples of the test set using the Dopri5 ODE solver. See Appendix D.3 for generated images.We followed previous work [23, 6] to use a fixed split of 90M/5M/5M with a fixed sequence length of 256. We used a 12-layer diffusion transformer (DiT) [48] based predictor, similar to the one used in [38] . The resultant model has 92.4M trainable parameters. Noticeably, our model size is smaller than [23] which used a full 24-layer model, but is similar to [38] . The models were trained for a total number of 3M iterations with a batch size of 512 per GPU (approximately 16 epochs), an initial learning rate of 10 -4 , and an exponential moving average (EMA) decay rate of 0.9999. We further used 1000 iterations as the linear warmup of the learning rate and used a decay rate of 0.8 with a patience of 3 and a validation interval of 2000 training iterations. The model snapshot with the lowest validation loss was saved for evaluation. BPCs were calculated on the first 4k sequences of length 256 of the test set (about 1M characters) with the Dopri5 solver as described in Appendix B.2, and generated texts were also obtained using the Dopri5 ODE solver. GPT-J-6B NLL and the generation token entropy were calculated based on 4k generations of length 256 (about 1M characters). See Appendix D.3 for generated texts.We used the splits from the dataset paper [7] that assign Chromosome 10 to the validation set, Chromosomes 8 and 9 to the test set, and all the other 21 human chromosomes to the training set. This assignment can avoid potential data leakage on the same chromosome. The transcription signals are provided as a float number for each base pair position. We also followed [7] to use a total number of 100k sequences with a context window of 1024 base pairs and added a random offset of 10 to the sequence position during training. The vector field predictor we used was identical to that in [7] , with 20 stacks of 1D convolutional layers and a total number of 13.3M trainable parameters. Our models were trained for 200k iterations with a batch size of 256 and an initial learning rate of 5 × 10 -4 . The model snapshot with the lowest validation SP-MSE on the validation set was saved for evaluation. The test SP-MSE was evaluated on the generated samples on the full test set's transcription signals with 300 Euler steps of generation.In this section, we provide additional results of the evaluation on Text8, additional ablation studies, and generated samples for each task.Figure 4 : GPT-J-6B NLL versus sample entropy. For MultiFlow, D3PM, and autoregressive, the curve represents different logit temperatures from 0.5 to 1. Baseline data are from [12] .We further follow [12] (MultiFlow) to calculate the NLL using GPT-J-6B [66] , a pre-trained autoregressive large language model. As such an NLL can be easily fooled by repeating high-frequency words, we also follow [12] to report additional token entropy, for which a closer entropy to data is preferred. The GPT-J-6B NLL and the token entropy are plotted in Fig. 4 , with the raw baseline data and the permission to reproduce provided by the authors of [12] . An additional random baseline that generates each character independently based on the letter frequency was provided. The ground truth data NLL and entropy are also included as the dotted horizontal and vertical lines. They demonstrate the best results any model can possibly achieve. Our SFM tended to better capture the diversity of the text corpus with the best entropy. Generated texts are provided in Appendix D.3.As GPT-J-6B was not trained on Text8, we noted a caveat that its NLL may not necessarily reflect the Text8 distribution fitness. From the MultiFlow results, such an NLL can be made artificially low by duplicating high-frequency words, e.g., repeated numbers with little semantic meaning as in the low-temperature MultiFlow generations (see Appendix D.3 for examples). We also noted that such an NLL can be easily fooled by randomly generated strings based on letter frequency. Additionally, low-temperature MultiFlow variants achieved lower NLLs than the ground truth data, making this metric less credible. We noted the huge impact of temperature on MultiFlow as it generated more repeated numbers with lower temperatures. In contrast, SFM achieved perceptually similar or even better results with more diversity.We provide ablation studies of the effect of different sampling methods, different sampling steps for the Euler methods, and different t max ≈ 1 for NLL calculation. The results are provided in Tab.4. The effect of the variational timestep t max matched the results from [7] , which stated that a closer timestep to the target data would decrease the NLL. The reasons we choose t = 0.995 instead of higher values are: 1) we noticed numerical stability issues at a timestep very close to the target, 2) we want a fair comparison with [7] which used an equivalent timestep of 0.9975, and 3) as long as we use a same timestep, the results are comparable within. We also found that NLL increased with the number of Euler sampling steps. We hypothesize that it is due to the large divergence change at t ≈ 1 and naive Euler steps tend to overestimate its contribution. Nonetheless, the Euler method with more than 300 steps gave a similar NLL as the ODE solver. We further provide additional results of the NLLs and FID scores on different sample methods in Tab.5 in which we used t = 0.995 and 300 Euler steps. The final performance was quite similar between these two sampling methods. The ground truth images and generated samples of the binarized MNIST dataset using various models and different sampling settings are demonstrated in Fig. 5 . The lower row demonstrates generated samples with different numbers of Euler steps.For the Text8 dataset, generated texts, NLL, entropy, and per sample NLL evaluated by GPT-J-6B using different models and sampling methods are demonstrated in Tab.6, with one sample with an NLL chosen around the average, one above the average, and one below the average for each model to provide a more comprehensive demonstration of the generation. The MultiFlow with different sampling temperatures used the model checkpoint in [12] . As SFM directly samples categorical probabilities instead of making multiple discrete jumps with logits, it is incompatible with temperature-tuning. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We describe the potential societal impacts here. Our proposed model is a general discrete generative framework that can be applied to various domains in machine learning including computer vision, natural language processing, and bioinformatics. Therefore, our proposed SFM can have potential positive impacts on various downstream applications if the machine learning model is properly utilized. For example, the design of promoter DNA sequence can better help us understand the complex interactions of gene activities. However, we also noted the potential negative impacts if the model is misused, especially in bioinformatics realms. We will work closely with both the machine learning community and the science community to ensure the proper usage of our model for various discrete generation tasks.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: Our work did not curate or propose new data or pre-trained models.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?Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: Our work did not involve crowdsourcing or research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.* Note that the formulation here is slightly different from Sec.2.1 where we fix a point µ on P. 
 
 Appendix:  


Decision: The paper extends flow matching to discrete data by working with the associated simplex equipped with the usual Fisher-Rao metric. This is a natural, yet novel, usage of Riemannian flow matching to statistical manifolds. Working with generative flows for discrete data is timely and the paper makes a sound contribution to the field.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: Authors adequately addressed limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: Yes
paper_summary: paper_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).
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: The paper lacks a discussion of the applicability of the proposed abstraction to the generation of discrete data (see Weaknesses).
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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
limitations: limitations: see "weaknesses" section.

Comments: title: Sincere Gratitude from Authors
the motivation is still unconvincing but I consider the method to be important
Alternative BPC Calculation
comment: 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.
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.
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 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.
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.
The appreciation for the thoughtful discussion is mutual, thank you for it.
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.
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 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.
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).
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 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.
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.