Categorical Flow Matching on Statistical Manifolds
Authors: Chaoran Cheng
Jiahan Li
Jian Peng
Ge Liu
Venue: NeurIPS 2024 poster (2024)
Keywords: Generative Model
Flow Matching
Statistical Manifold
Primary area: generative_models
TL;DR: We proposed Statistical Flow Matching (SFM), a novel flow matching framework on the manifold of parameterized probability measures for discrete generation.

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.

Categorical Flow Matching on Statistical Manifolds
Chaoran Cheng"
Jiahan Li' Peking University 1i lijiahanypcopku. edu.
University of Illinois Urbana-Champaiin
1it.//iiiiiiiieeuu edu
Jian Peng
Ge Liu
University of Illinois Urbmna-Champaign jianpeng@1 11 winois. edu
University of Illinois Urbana- Trbana-Champaigee geliu@illinois edu
Abstract
We introduce Statistical Flow Matching (SFM), novel and mathematically rig- orous How-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 instantiat- ing 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 Riemannian structure whose intrinsic geometries are effectively leveraged by following the shortest paths of geodesics. We develop an efficient training and sampling algorithm that over- comes numerical stability issues with diffeomorphism between manifolds. Our distinctive geometric perspective statistical manifolds allows us to apply op- timal transport during training and interpret SFM as following the steepest di- rection 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 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. Our code is available https: //ts:////ttbt//////// www./oommmccieeeeeeee cheng/ statistical- www.waalllllow--oooooo fww.wlowww-oo-o--- wwwowwttciiing
Introduction
Recently, conditional flow matching (CFM) models [38] have achieved remarkable success in vari- ous generative domains including image generation [38, 17, 29], molecule [59, 58, 31] and protein design [68, 10, 36], and sequence generation [60, 39, 12]. While attempts to generalize CFM and diffusion models discrete categorical data have been made, they typically exert ad hoc assump- tions on the structure of the discrete distribution. One group of work relies on stochastic jumps of Markov chains in either the discrete-time [6, 39, or continuous-time setting [11, 55, 12] that dis- cards the continuous nature of the underlying categorical distributions. Other work directly works with the probability simplex [7, 60] or the corresponding logit space [28, 40, 24] with potentially imperfect assumptions that fail to capture the underlying true geometry of the statistical manifold.
Furthermore, likelihood is often approximated by variational bounds previous discrete generative models due to the lack of tractable exact likelihood.
'qqual contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 [51] 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:
(I) 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 flow- matching training algorithm that avoids numerical issues by leveraging diffcomorppism 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 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 theoret- ical 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 toy example 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.
2 Preliminary
2.1 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 struc- ture known as staristical manifold. Mathematically, consider probability densities p dy X R defined by the Mikon-NNkodyym derivative where μ is a probability measure on the sample space X and is the reference measure X. Suppose the statistical manifold P P(X) {p fx du Jx p(x; 8) dv 1} is parameterized by 0 (81,02..... 0,,) O, this parameterization naturally provides a coordinate system for P on which each point is probability measure μe with the corre- sponding probability density function p(x; 0). The Fisher information metric is defined as
T log p(X;0) log ((g)((X;)
a Olo p(z; f) O p(x; 0) pl::; 0) (1)
(0) = Ex
00,
00p
00,
00,
Rao demonstrated in his seminal paper [51] that statistical manifold can be equipped with the Fisher information metric to obtain 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 "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-Ree Rao distance [5]]. measures the similarity between them. The tangent space (P) at a point P can be naturally identified with the affine subspace Tpp(P) {vl fx du 0} where each element v is signed measure over X. The exponential map exp,, TT (P) + P and logarithm map log,," P + (P) can also be defiedd on the statistical manifold. While the geodesic for a parame- terized family of probability measures can be obtained numerically by solving the geodesic equation when closed-form expression is unknown (see Appendix A. it usually requires expensive simu- lations. Fortunately, closed-form geodesic distances are available for many common distributions including categorical, multinomial, and normal distributions [44], which motivates our method.

---

2.2 Conditional Flow Matching on Riemannian Manifold
The conditional flow matching (CFM) framework [38] 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 flow-based model can be viewed as the continuous generaliza- tion of the score matching (diffusion) model [56, 57, 26] while allowing for 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 Pr [0, + P(M) is a curve of probability densities over M. A flow v, |0, 1] x M + M is time- dependent diffeomorphism defined by a time-dependent vector field u, : 1] x M - TMM via the ordinary differential equation (ODE): the (a) = we (e(vv(())) The flow matching objective directly regresses the vector field with a time-dependent neural net v(r,, f) where == (x()). However, this objective is generally intractable. Both [38, 14] demonstrated that tractable objective can be derived by conditioning on the target data Iı at t = of the probability path. The Riemannian flow matching objective can be formulated as [14]
= http://0...lommpppoaamaaaaaaa [1]v ux x =1]1
(2)
where 9 is the data distribution, po is the prior distribution, and 2; the (2|22, 21) 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 5t expso (1 logz, 31), and the corresponding vector field can be calculated as ue ( x1) = de Ft = log., (21)/(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. for probability densities over the manifold ((()) and μt, for probability measures (or probability masses for discrete distributions) on the manifold (()) to avoid confusion.
3 Method
Different from previous work that treated each distribution separately, we adopt an integral view- point 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 gener- ation tasks. However, we emphasize that our proposed SFM applicable to any statistical manifold with closed- geodesic distance and can be broadly used in generative tasks targeting probabil- ity measures on the statistical manifold.
3.1 Statistical Manifold of Categorical Distributions
Consider the discrete sample space X' {,,,2.... n.. an n-class categorical distribution over X' can be parameterized by n parameters P1, .2..... ......,,, Hy, such that Σx= P; = 1. p4 2 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 {o'} } over X: μ En= u.d". Geometrically, this manifold can be visualized as the (n 1)-dimensional simplex An The geodesic distance between two categorical distributions is given in [51, 44] as
dea v) = 2 arccos Vvnn
The tangent space at a point can be identified with (P) = {ul Σm =0 and the corre- sponding inner product at point y defined as
(u, v) Σ 42u p4
MEP+, (T(P)
where P+ denotes the statistical manifold of positive categorical distributions. 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 ST-1, #4 1H+ I, Vp..
(5)
3

---

which maps the original statistical manifold to the positive orthant of unit (n - 1) -pphere St
{ar E oे = æ; 0} (see Fig.2). Note that we have [| ll = E v/4² Ei M4 1.
The geodesic on - 5 follows the great circle, and the following proposition holds between the geodesic distance on 5 and P:
Proposition
ds (s(((1) de) 1, E.
(6)
proof is provided Appendix A.3. This indicates that we can work with the geodesic distance on the unit sphere instead with up to constant factor:
ds(2, v) arccos v)), Ir E
The geodesic distance dy and the inner product (.+) are well-defined for the boundary, and found this transform led to the practical stabilized training of the flow model. Visualizations of the Riemannian geometry the statistical manifold of 3-class categorical distributions and the corresponding Euclidean geometry on the simplex are provided in Fig. I for comparison. The straight lines under the Euclidean assumptions fail capture the true curved geometry of the statistical manifold.
Geodesic Distance
Exponential Map
Logaritim Map
Statistical
Manifold
Euclidcan Manifeld
Figure/ The Riemannian geometry of the statistical manifold for categorical distributions in comparison to Euclidean geometry on the simplex. Left: Contours for the geodesic distances to 14 (1/3, 1/3, 1/3). Middle: Exponential maps (geodesics) from μo to different points near the boundary. Right: Logarithm maps (vector fields) to μco
3.2 Statistical Flow Matching
We provide the analytical form of the exponential and logarithm maps on the statistical manifold of categorical distributions in Appendix A.3. Although possible directly learn the vector field following the loss in (2), such direct parameterization has numerical issues near the boundary. As described 3, we apply the diffeomorphism æ (5. to derive more stable training objective the spherical manifold:
LSEM (ot. (g(r))
[1],(,,, t) uff (are 21)
(8)
where Po is the prior noise distribution over P and g is the data distribution; n, denotes the standard pushforward measure. -T-1 [o, 1] + TST is a learnable time-dependent vector field net- work that maps the interpolant o, on the unit sphere to a tangent vector the tangent bundle T The ground truth conditional vector field 5 calculated using the exponential and logarithms maps the sphere (Appendix A.2). The overall training and inference scheme is visualized in Fig. and described in Alg.2 and 3 in Appendix C.
We further implement Euclidean flow matching model on the probability simplex with linear in- terpolation between the noises 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 as
LineuriM U/ [u, 11 (p ((g(()) [Ilo(e,)) t) (11 mo)1
In the discussion above, assume single data dimension the statistical manifold. This can be extended to any data dimension by treating them as independent channels of the input. practice, each probability measure the simplex is represented by matrix X [0, 1] where D is the data dimension and each row sums to 1. The priors are independently sampled from the uniform distribution over the (n -immmmplex and each data dimension is independently interpolated with the same timestep ~ Uio, 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 then drawn from the generated categorical distributions for evaluation. Details regarding model parameterization and sampling are further described in Appendix C.I and C.2.
(a)
(b)
dd
log det do (ax))
p1
π
(x 21)
log af (
u1
π
L1 Log p(p (Pa)
(c) Na, 21)
log oDB (2:)
140
30
140 log (Mn)
6
220
π
π log dr(pp))
π
6'
Figure 2: Statistical flow matching, (SFM) framework. (a) During training (Sec.3.2), probability measures on P are mapped to 54-11 via diffeomorphism to compute the dime-dependent vector field (in red). During inference, the learned vector field generates the trajectory on S and we map the outcome of ODE back to P (in blue). (b) In the NLL calculation for one-hot examples (Sec.3.5), the probability density is marginalized over a small neighborhood of some Dirac measure avoid undefined behaviors the boundary (in green).
3.3 Optimization View of Statistical Flow Matching
We further provide an interpretation of our proposed statistical flow matching framework from the perspective of optimization problem. From the optimization viewpoint, for generative model on a statistical manifold, we want to minimize the discrepancy between the target distributions and the generated distributions. Naturally, the Kullback-Libler divergence DxL (#()))l((0)) can be used as a measure statistical distance. We note that the Fisher information metric can be obtained as the Hessian of the KL divergence with respect the parameter 6. This can be demonstrated by the fact that the KL divergence reaches the global minimum 0 at f = 01, so all first-order partial derivatives are zero, and the Hessian is positive semi-definite. Therefore, Taylor expansion of KL
divergence at 0 with AO = 0 - 0, gives
82 00,00k
D (=(((())))())))) ~
Σ
A0, A0
D (xa()))l))((e)))
2
jk
18=8;
(10)
1 A ABkg,K (01)
- 1A011
2
From the geometric viewpoint, the geodesic, by definition, is (locally) length- -inimiziingiggeeei curve with respect to the corresponding Riemannian metric. Therefore, by following the direction of the vector field that decreases the geodesic element ds2 |0 || | we are indeed following the steepest direction that minimizes the local KI. divergence. In this sense, the Fisher information metric defines a second-order optimization scheme for the KI. divergence by following the "steepest" direction of the Hessian. Indeed, The steepest direction A that decreases the KL divergence is
5

---

known as the natural gradient in the existing literature [2, 3] and has been explored in optimization known as natural gradient descent[ 48, 42, 46]. Instead of optimizing along the normal gradient, the natural gradient is defined as VL = F-IVL where F the Fisher information matrix estimated from batch of sampled data. Results established in [2, 42] demonstrated that stochastic natural gradient descent 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 instcad of general statistical models. Such a theoretical connection may contribute to the better performance of SFM.
3.4 Optimal Transport on Statistical Manifold
The geometric viewpoint of the statistical manifold offers a continuous and differentiable formu- lation of generative modeling and also provides a robust way to measure the distance between two categorical distributions via the well-defined geodesic distance Eq.(3). In contrast, the Markov chain- methods cannot establish a robust distance measure due to the stochastic jumps between discrete states. Inspired by the optimal transport formulation in previous work [45, 20, 68, 63], we naturally extend it to our statistical setting in which the cost is defined by averaging the statistical distances over data dimensions (X, for X { Y { 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.
DE==dduu(u" v/())
3.5 Exact Likelihood Calculation
Unlike diffusion- -baed 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 I E M, using the change of measure formula, the likelihood can be modeled by the continuity equation 43];
i a
log (()) + div (v)))))) 0,
(1))
where div, is the Riemannian divergence and (=)))) == v(z,, t) is the time-dependent vector field. In this way, the pushforward probability measures (xx) defined via the learned flow can be ob- tained as the integral of the Riemannian divergence back through time on the simulated trajectory xt. Following [9, 7], we define the ODE log-likelihood as the change of the log- log-likelioood as:
= J," (vx))(()) di,
log Dogponi (x1) ODE (
(12)
3 = (x(2) reverse where the trajectory Ie is obtained by solving the differential equation a through time with the initial condition T2 at L = 1. In this way. we have (ogp(z,) = logpode log Po(zo) where log (z() the log-likehood of the prior distribution at data point zp. In prac- tice, we follow previous work [7] to use Hutchinson's 's trace estimator [30] to efficiently obtain an unbiased estimation of the divergence using standard normal random variables. To further account for the transform from P S and the reverse transform -- 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) := PV=--V), we have the change of measure identity n, (4)) det dr (r( = P(x) for x = π (1)). Therefore, by adding the two log-determinant of the pushforward measures, the ligeliiid can be formulated as
log pi (11) log det (xx)) ) + log logpoDig (xI) + log det do(()))) + log po()))
(13)
The 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 [62]. Following [7], we derive variational lower bound for the likelihood as the marginalized probability over the small neighborhood of a Dirac measure ja = 5 at the boundary:
log p()) 2 Eg( log (8) + log (ogp())ā)) + log p1 (ii)),
(14)

---

where ((]1))) can be viewed as the forward noising probability at A1 which is close to 8 with closed-form likelihood. p()) is the categorical likelihood (cross-entropy) and pi (pr) ) is the model likelihood in (1.)) 'he overall workflow of calculating NLL is demonstrated in Fig.2, and more details regarding likelihood computation can be found 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 Swiss roll distribution the probability simplex. We then conduct tensive 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 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 log- arithm maps the statistical manifold without the diffeomorphism (SFM w/o πt) is also compared in the toy example. We also implement a linear flow matching model on the probability simplex using the loss (9) as an additional baseline, for which dub LinearFM. For discrete data, we always use Eq. (14) to obtain finite negative log-likelihoee (NLL), More details regarding model architectures and the experimental setup can be found in Appendix C.
4.1 Toy Example: Swiss Roll Simplex
We project the Swiss roll dataset onto the 2-simplex with some additional margins to make sure no point resides the boundary. We used dinendependent 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.
SFM OT -2.24 = 0.07 (-2.28)
SFM w/o OT -2.27 = 0.07 (-2.29)
SFM w/o x -1.71 = 0.06 (-1.73)
Ground Truth
LinearFM
DDSM
DirichletFM NA
1 (-1.93)
10.78
Figure 3: Generated samples of the Swiss roll simplex dataset and NLL (lower is better). The NLLs are estimated using Hutchinson's trace estimator, whereas those in the parenthesis are exact.
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 [60] 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 (q3( 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.
7

---

4.2
Binarized MNIST
We extend our SFM to model discrete generation tasks in computer vision. The binarized MNIST dataset [53] is the binarized version of the original MNIST dataset [34] by thresholding the original continuous value to be either or 1, thus can viewed as a 2-class generation task with a data dimension of 28
784. We also compared D3PM [6] and DFM [22] as examples of masked diffusion models for discrete generation. We used convolutional neural net adopted from [57] with additional additive Fourier-based time embeddings. The quantitative evaluation of NLL and Fréchet inception distance (FID) are shown in T.
Table 1: NLL and FID of different discrete models on binarized MNIST. The NLLs in the parenthesis are discrete NLLs; therefore, they are not directly comparable. * is from [7].
Model
SFM w/ OT
SFM w/o OT
SFM w/o
LinearFM
NLL
-1.687 0.020
-1.631 = 0.027
0.928 + 0.035 8.2073
0.622 0.022
FID
4.62
5.15
5.91
Model
DirichletFM
DDSM
D3PM
DFM
NLL
NA
0.100 + 0.001
(0.141 + 0.021)
(0.101 = 0.017)
FID
77.35
7.79
67.36
34.42
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 varia- tional 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.
4.3
Text8
The Text8 dataset [41] 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. 24] to use random chunks of length 256 in both training and evaluation without any preprocessing. We used a 12-layer diffusion transformer (DiT) [49] based predictor, similar to the one used in [39]. As our exact likelihood is not directly comparable to bits- bits-per-characterr (BPC) reported in previous work which relies on an alternative variational bound for the likelihood, we additionally calculate such BPC inspired by [52]. 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 reference (Discrete Flow [64] based on autoregressive normalizing flow).
Note that, unlike all of the other basclines. SFM and LinearFM do not directly optimize such a BPC as 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 perfor- mance among other diffusion and flow baselines that were directly trained with cross-entropy losses. We also noted significant performance gap between SFM and the naive linear flow matching on simplex, highlighting the importance of capturing the intrinsic geometric properties of the underly- ing 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] provided in Appendix D.1, in which SFM achieved the best generation entropy as evaluated by the pre-trained GPT-J-6B model [67].
4.4 Promoter Design
We further apply SFM to the practical task 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 human promoter sequence dataset containing 100k promoter sequences with the
8

---

Table 2: BPC on Text8. Results marked
Table 3: SP-MSE (as evaluated by Sei [13]) on the generated promoter DNA sequences. Results marked are from [7] and results marked 1 are from [60].
taken from the corresponding papers. Model
BPC|
SFM w/ OT
1.399 = 0.020
Model
SP-MSE
SFM w/o OT
1,386 0.033
LincarFM
1.651 = 0.027
SFM w/ OT
0.0279
D3PM-absorb(6]
1.47"
SFM w/o OT
0.0258
BFN[24]
1.41"
LinearFM
0.0282
SEDD-absorb(]9]
1.32"
DDSM
0,0334
MultiFlow[12]
1.41"
D3PM-unifone
0.0375"
Argmax [28]
1.80*
Bit-Diffusion (one-hot) [15]
0.0395*
Discrete Flow[64]
1.23
Bit-Diffusion (bit) [15]
0.0414*
Transformer(6) []
1.23*
Language Model
0.0333
Transformer [I
1.08"
DirichletFM
0.0269
corresponding transcription initiation signal profiles. Each promoter sequence 1024 base pairs long and is centered at the annotated transcription start site position [27]. Therefore, we can model promoter design as 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 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.
To provide quantitative evaluation of the generated promoter sequences, we follow [7] to apply the pre-trained deep learning sequence model Sei [13] predict active promoters based 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 pro- moter 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 better) and is demonstrated in Tab.3. Our SFM was able 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.
5 Related Work
As we put special interest in discrete generation, start with the existing work discrete gen- eration. We first list a few traditional autoregressive models [16, 50, 64] and masked language models [18, 54] but will skip them as mainly focus on the diffusion and flow matching mod- els. Existing diffusion and flow-based discrete generative models can be categorized into two main groups. The first group relies 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 grad- ually 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 [39] proposed a more ef-
ficient training scheme by learning the score entropy between different states. [11, 55] extended
the continuous- time Markov chain by using the infinitesimal generator (rate matrix) instead, and [12, 22] further applied the flow matching framework. Although the transition matrix or the rate matrix determines the entire diffusion trajectory, there no canonical way of choosing for different generation tasks to control the mixing rate. Also, due to the presence of discrete jumps of the Markov chain, exact likelihood calculation 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 [36]. As an example, our linear flow matching assumes Euclidean geometry on the probability simplex and is often used as a baseline in previous work [60]. The multinomial diffusion [28] 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) [60] exerted 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 the logit space, the Bayesian flow net- work (BFN) assumed Gaussian distributions on the logit space with Bayesian update rules. [40] also assumed Euclidean geometry in the logit space with targets soft one-hot logits. The assump- tions 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 [51] discussing the Fisher information metric. Most previous work utilizing information geometry focuses optimization [2, 3] with few exceptions on representation learning. [65] utilized the geodesic distance between two univariate Gaussians for function shape matching for retrosynthesis of gold nanoparticles. [35] treated point cloud data as probability measures over the 3D Euclidean space and considered the pullback metric on the latent space obtain optimal latent coordinates for the autoencoder. [47] further applied such 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 generative modeling. SFM directly operates on the statistical manifold with closed-form geodesics, providing a simulation- free approach for efficient generative training.
6 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 the statistical manifold of categorical distribu- tions in this work and derived optimal transport and exact likelihood formulae. We applied SFM to diverse downstream discrete generation tasks across different domains demonstrate work's 's effectiveness over the baselines. We noted that SFM can be further extended to non-discred generative tasks whose targets are probability distributions, which we will leave as future work.
frame-
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 struc- ture can be induced by the Fisher metric. However, discretized data like CIFAR-10 [32] (256 ordinal pixel values) do not follow this assumption, and results on such data might be suboptimal.
References
[1] S. Alamdari, N. Thakkar, R. van den Berg, A. X. Lu, N. Fusi, A. P. Amini, and K. K. Yang. Protein generation with evolutionary diffusion: sequence is all you need. bioRxiv, pages 2023- 09, 2023.
[2] S.-I. Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
[3] S.-i. Amari and A. Cichocki. Adaptive blind signal processing- neural network approaches. Proceedings of the IEEE, 86(10):2026-2) 2048, 1998.
[4] S.-i. Amari and H. Nagaoka, Methods of information geometry, volume 191. American Math- ematical Soc., 2000.
10

---

[5] C. Atkinson and A. F. Mitchell. Rao's distance measure. Sankhya: The Indian Journal of Statistics, Series A, pages 345-365, 1981.
[6] Austin, Johnson, Ho, D. Tarlow, and Van Den Berg. Structured denoising dif- fusion models in discrete state-spaces, Advances in Neural Information Processing Systems, 34:17981-17993, 2021.
[7] Avdeyev, C. Shi, Y. Tan, K. Dudnyk, J.d Zhou. Dirichlet diffusion score model for biologi- cal sequence generation. In International Conference on Machine Learning, pages 1276-1301. PMLR, 2023.
[8] N. Ay, Jost, H. Vân Lè, and Schwachhöfer. Information geometry, volume 64. Springer, 2017.
[9] H. Ben-Hame S. Cohen, J. Bose, B. Amos, A. Grover, M. Nickel, R. T. Chen, and Y. Lip- man. Matching normalizing flows and probability paths on manifolds. arXiv preprint vrXiv:2207.04777 2022.
[10] A. J. Bose, T. Akhound-Sadegh, K. Fatras, G. Huguet, J. Rector-Brooks, C.-H. Liu, A. C. Nica, M. Korablyov, M. Bronstein, and A. Tong. Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint arXiv:2310.0233// 2023.
[11] A. Campbell, Benton, V. De Bortoli, T. Rainforth, G. Deligiannidis, and A. Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-288 28279, 2022.
[12] A. Campbell, Yim, R. Barzilay, T. Rainforth, and T. Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint 3xXi:22402.0499,, 2024.
[13] K. M. Chen, A. K. Wong. Troyanskaya, and J. Zhou. A sequence-based global map of regulatory activity for deciphering human genetics. Nature genetics, 54(7):940 9994, 2022.
[14] R. T. Chen and Y. Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.0m 03660, 2023.
[15] Chen, R. Zhang, and G. Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint arXiv:2208. 04202, 2022.
[16] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Atten- tive language models beyond a fixed-length context. arXiv preprint 11XIv::901.20000) 2019.
[17] Q. Dao. H. Phung, B. Nguyen, and A. Tran. Flow matching in latent space. arXiv preprint arXiv:2307.08699, 2023.
[18] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 181000080 2018.
[19] J. Dormand and P. Prince. A family of embedded runge-kutta formulae. Journal of Computa- tional and Applied Mathematics, 6(1):19-26, 1980.
[20] K. Fatras, Y. Zine, S. Majewski, R. Flamary, R. Gribonval, and N. Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792. 2021.
[21] S. Gallot, D. Hulin, and J. Lafontaine. Riemannian Geometry. Universitext. Springer Berlin Heidelberg, 2004.
[22] I. Gat, T. Remez, N. Shaul, Kreuk, R. T. Chen, G. G. Synnaeve, Y. Adi, and Y. Lipman. Discrete flow matching. arXiv preprint arXiv:2407.15595, 2024.
[23] A. Graves. Generating sequences with recurrent neural networks. arXiv preprint
M.Xivv/308. 0850, 2013.
[24] A. Graves, R. Srivastava, T. Atkinson, and Gomez. Bayesian flow networks. arXiv preprint WxXiv:2308.07033 2023.

---

[25] K. He, X. Zhang. S. Ren, and J. Sun. Deep residual learning for image recognition. In Pro- ceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778. 2016.
[26] Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems,
2020.
[27] C.-C. Hon, J. A. Ramilowski, J. Harshbarger, N. Bertin, O. J. Rackham, J. Gough, E. Denisenko, S. Schmeier, T. M. Poulsen, J. Severin, et al. An atlas of human long non-coding mas with accurate 5' ends. Nature, 543(7644):199-004 2017.
[28] E. Hoogeboom, D. Nielsen, P. Jaini, P. Forré, and M. Welling, Argmax flows and multino- mial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34:1 2212-21245555 2021.
[29] V. T. Hu, W. Zhang, M. Tang, P. Mettes, D. Zhao, and C. Snoek. Latent space editing in transformer-based flow matching. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 38, pages 2247-2255, 2024.
[30] M. F. Hutchinson. A stochastic estimator the trace of the influence matrix for laplacian smoothing splines. Communications in Paristics-Simulation and Computation, 18(3):1059- 1076, 1989.
[31] L. Klein, A. Krämer, and F. Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36, 2024.
[32] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[33] A. Le Brigant, S. C. Preston, and Puechmorel. Fisher-rao geometry of dirichlet distributions. Differrnniial Geometry and its Applications, 74:101702, 2021.
[34] Y. LeCun, C. Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: htlp://yann.le http://yammleecn.coom
2, 2010.
[35] Y. Lee, S. Kim, J. Choi, and F. Park. A statistical manifold framework for point cloud data. In International Conference on Machine Learning, pages 12378-12402. PMLR, 2022.
[36] J. Li, C. Cheng, Z. Wu, R. Guo, S. Luo, Z. Ren, J. Peng. and J. Ma. Full-atom peptide design based on multi-modal flow matching. arXiv preprint arXiv:2406.00735, 2024.
[37] G. Lin, A. Milan, C. Shen, and I. Reid. Refinenet: Multi-path refinement networks for high- resolution semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1925 1934, 2017.
[38] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[39] A. Lou, C. Meng, and S. Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310. 16834, 2023.
[40] R. K. Mahabadi, Tae, H. Ivison, J. Henderson, Beltagy, M. E. Peters, and Cohan, Tess: Text-to-text conficondit simplex diffusion. arXiv preprint arXiv:2305.08377 2023.
[41] M. Mahoney. Large compression benchmark, 2011.
[42] J. Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
[43] E. Mathicu and M. Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 53:2503-2515, 2020.
[44] H. K. Miyamoto, C. Meneghetti, and Costa, closed-form expressions for the fisher- distance. arXiv preprint arXiv:2304.14488 14888, 2023.
12

---

[45] K. Nguyen, D. Nguyen, T. Pham, N. Ho, et al. Improving mini muni-batch optimal transport via partial transportation. In International Conference on Machine Learning, pages 16656-16690. PMLR, 2022.
[46] L. Nurbekyan, W. Lei, and Y. Yang. Efficient natural gradient descent methods for large- scale pde-based optimization problems. SIAM Journal on Scientific Computing. 45(4):A1621- A A1655, 2023.
[47] A. Palma, S. Rybakov, L. Hetzel, and J. Theis. Modelling single-cell ma-seq trajectories on a flat statistical manifold. NeurlPS 2023 Al for Science Workshop, 2023.
[48] H. Park, S.-I. Amari, and K. Fukumizu. Adaptive natural gradient learning algorithms for various stochastic models. Neural Networks, 13(7):755-764, 2000.
[49] W. Peebles and S. Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
[50] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
[51] C. R. Rao. Information and the accuracy attainable in the estimation of statistical parameters. In Breakthroughs in Statistics: Foundations and basic theory, pages 235 247. Springer, 1992.
[52] S. S. Sahoo, M. Arriola, Y. Schiff, A. Gokaslan, E. Marroquin,
Chiu, A. Rush, and V. Kuleshov. Simple and effective masked diffusion language models. arXiv preprint arXiv:22406.011 077524 2024.
[53] R. Salakhutdinov and I. Murray. On the quantitative analysis of deep belief networks. In Proceedings of the 25th international conference on Machine learning, pages 872-879, ACM, 2008.
[54] Salazar, D. Liang, Q. Nguyen, and Kirchhoff. Masked language model scoring. arXiv preprint arXiv: 14659, 2019.
[55] J. E. Santos, 7. R. Fox, N. Lubbers, and Y. T. Lin. Blackout diffusion: generative diffusion models in discrete-state spaces. In International Conference on Machine Learning, pages 9054-9059. PMLR, 2023.
[56] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[57] Y. Song and S. Ermon. Improved techniques for training score-based generative models. Ad- vances in neural information processing systems, 33: 1511281-22448 2020.
(58) Y. Song, J. Gong, M. Xu, Z. Cao, Y. Lan, S. Ermon, H. Zhou, and W.-Y. Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2024.
[59] H. Stark, B. Jing, R. Barzilay, and T. Jaakkola. Harmonic prior self-conditioned flow matching for multi-ligand docking and binding site design. In NeurlPS 2023 Al for Science Workshop, 2023.
[60] H. Stark, B. Jing, C. Wang. G. Corso, B. Berger, R. Barzilay, and T. Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.05841 2024.
[61] C. Szegedy, V. Vanhoucke, S. loffe, J. Shlens, and Z. Wojna. Rethinking the inception archi- tecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1-2268., 2016.
[62] L. Theis, v. d. Oord, and M. Bethge. A note on the evaluation of generative models. arXiv preprint arXiv:1511.018.0 2015.
[63] A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.00442. 2023.
13

---

[64] D. Tran, K. Vafa, K. Agrawal, L. Dinh, and B. Poole. Discrete flows: Invertible generative models of discrete data, Advances Neural Information Processing Systems, 32, 2019.
[65] K. Vaddi, H. Chiang, and L. D. Pozzo. Autonomous retrosynthesis of gold nanoparticles spectral shape matching. Digital Discovery, 1(4):502-510, 2022.
[66] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polo- sukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[67] B. Wang and A. Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github:
www.coo/////oo.oo wttpl/.coo/oiooieie wiwgooolllo//ema transformer- jax, May 2021.
[68] J. Yim, A. Campbell, A. Y. Foong, M. Gastegger, J. Jiménez-Luna, S. Lewis, V. G. Satorras, B. S. Veeling, R. Barzilay, T. Jaakkola, et al. Fast protein backbone generation with se (3) flow matching. arXiv preprint 2rX1n:2310000529,, 2023.
14

---

Supplementary Material
A Riemannian Manifold and Information Geometry
In this section, we give a more detailed mathematical background on Riemannian manifold and in- formation 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 anal- yses and rigorous formulations.
A.1 Riemannian Manifold
A Riemannian manifold .M is real, smooth manifold equipped with a positive definite inner product g on the tangent space (M) at each point E M. Let TM UneM Tx be the tangent bundle of the manifold .,, time-dependent vector field on M is mapping 14 [0, X + TM where ue (2) E (MM.. A geodevic is a locally restance-minimzzing 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 E (MM), there exists a unique geodesic
[0. M such that 7(0) = x and y((0) u. The exponential map exp M T'M - M is uniquely defined be exp, (2) (1).
The logarithm map log .M x M + TM is defined as the inverse mapping of the exponential map such that exp, (log., (v))
y, Vz, y E M. Note that the formulation here is slightly different from Sec.2.1 where we fix a point on P. With the exponential map and logarithm map, the time- dependent flow can be compactly written as time interpolation along the geodesic:
(15) can be demonstrated that the above flow indeed traces the geodesic between zo. 21 with linearly decreasing geodesic distance dg (2t, = (1 t)dg 21) [14]. For an n-dimensional Rieman- nian manifold, the geodesic can be obtained by solving the geodesic equation (using the Einstein summation convention):
Tt == the (z+ ro, 21) exPa, (t (oo, f1), t E 10.
d2in
+ 1 dx dæ
di dt = 0, =
(16)
d/2
where xf are local coordinates of the geodesic and r are the Christoffel symbols defined by dgmi
(Qgmi
0gis
1 29 km
i,j, k 1,,1,,2 .., n2
(17)
Ox
drm
where n is the inverse metric such that gjz = 81. The closed-form expressions for the geodesic are generally not available. [44] curated a wide range of common statistical manifolds for parame- terized families 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 the former via the diffeomorphism π in (5).
A.2 Spherical Manifold
The positive orthant of the unit (n - 1))--erp ST a (n 1)-dimensional manifold. The sphere can be embedded into the ambient Euclidean manifold Rn such that it inherits the canonical inner product as
(u, v)s = (u. = Σ U.Y. U. E (S4
(18)
1=1
The tangent space Tu (ST--)) = (ul(u, .²) 00 is (n 1 dimmenniioaal hyperplane perpendicular to the vector The geodesic the sphere follows the great circle between points, and the geodesic distance can be calculated in Eq.( The corresponding exponential map can be calculated as:
exp (w) = x cos ||u]|2 + u sinc ||]]22 where sinc( (0) = sin 0/0 is the unnormalized sinc function. The logarithm map can be calculated as:
(19)
log, (y) arccos v)) Pa z)
(20)
PL
x21
where (ww)) w (x, w) is the projection of vector w onto the tangent space (ST-))).
15

---

A.3 Statistical Manifold of Categorical Distributions
By calculating the Fisher information matrix for categorical distribution, we can obtain the explicit form of the metric tensor as [44]
6gk 1 1
g))(m)
(21)
Pi
where o is the Kronecker delta. Substituting this into (u, v) leads the Riemannian inner product in (4.( The statistical manifold of categorical distributions ( closely related to S by the diffeomorphism in Eq)( In fact, и is an isometry between the Fisher information metric on P and the standard induced metric on 5 up to a constant scaling factor of 1/4. To see this, we have the following proposition:
Proposition 2.
(T(() θT(H) Ou Ov
1 (u, v)as u, (P)).
1
(22)
Proof.
(T(()) OT(H))
(
d
+ tu)
d
t + tv)
u4
U/
1 (u, v)...
Ou Ov dt
0 dt
1=
2
2/14
4
i=1
Note that the infinitesimal are length on the geodesic can be expressed ds² || dx
Eju (z) dæy where ||- ll, is the canonical Riemannian norm induced by the inner product. easily arrive the result in Proposition 1. Based this, the
Integrating the geodesic, exponential map and logarithm map on this statistical manifold can be written as
2
(xp,,uu=
Cicos cos
u
sinc
(23)
2/F4
2/F
2//4
aarccos( (VTH, VD))
log (V)
(V/ (Vn, VD,u)
(24)
V VI- (VH,V)
where O denotes the pairwise multiplication. We mentioned the main text that directly using the exponential and logarithm maps the statistical manifold has numerical issues near the boundary. Instead, we used the mapping π to work with the spherical manifold. Nonetheless, Eq. and (24) provide useful tools for visualization of the statistical manifold, demonstrated in Fig.
We specifically note the property that the Riemannian structure induced by the Fisher information metric leads to vector fields more parallel to the boundaries. This can also be derived mathematically from the logarithm map (24). Consider the direction term /p @ vー (VI,, Vv)u. For μ close to the boundary with μıg R 0, its corresponding vector field will also have a close to u4 R O component, which different from linear flow matching's fixed v μ. We hypothesize that one potential benefit of such curved geometry over the naive Euclidean geometry is that the former helps prevent overshooting across the boundaries. Specifically, consider a target point on the boundary. The Euclidean vector field will continue to push the points outside the manifold, whereas the Riemannian vector field tends to travel parallel to that boundary to prevent going across the boundary. Once overshooting happens during sampling, the model may exhibit undefined behavior it was never trained on the points outside the manifold.
We also noted the relation between the Fisher information metric defined in Eq.(1) and the canonical inner product defined in (4) for general statistical manifolds. Let M be an n-dimensional differ- entiable manifold and consider an embedding defined by M + P(M), 0 + p(0) Σiex (0)0''
16

---

The pullback of the Fisher metric g defines a metric on M. For u, To(M), we have
go (u, v) p" (g)o(u,)) v)
Op Op
9p(0)
Ou 0v
Σ
1 dp:
(0)
0pi (0)
(25)
Di (i) du
Zu
0lo blog Pi
a log Ps
Σ (()
(0)
(0)
du
dc
which gives the more common definition of the Fisher information matrix in (q.(1). This relation also generally holds for continuous sample space X' but requires more careful handling with mea- sure theory tools. We refer interested readers mathematical references, e.g., Chapter 3.1 in [8].
A.4 Other Statistical Manifolds
The statistical manifold categorical distributions provides an intuitive way of modeling discrete targets. However, we noted that other parameterized distributions can also be used to derive alterna- tive 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).
Dirichlet Distribution The Dirichlet distribution Dir(a: is a multivariate generalization of the beta distribution. It is defined on the (n -immeennion simplex A with concentration parameters a = {(a)) The probability density function is defined as
p(z;a)
II
2 DEA-11 t=1
(26)
B(a)
where B(a) -
05) is the multivariate beta function. The closed-form expres- sions 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 [33]. [7] proposed to follow the specific path of Jacobi diffusion processes on the marginal beta distributions, for which expensive modified Jacobi polyno- mials needed be precomputed. [60] further applied the flow matching framework by considering the probability path of Dir (1 + tex. 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 specific forward noising process which always terminates (or near) one-hot distributions. This assumption generally holds for discrete generation, but cannot be gener- alized 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 not the shortest geodesics in the sense of the Riemannian met- ric. 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 more mathematically concise.
The final one-hot distributions are unreachable [60]. as it would require an infinite timestep. Therefore, variational bound needs to be derived for likelihood estimation, which was done the original paper.
Multinomial Distribution Consider m i.i.d. experiments that follow categorical distribution with possible outcomes and probability μ {u multinomial distribution gives the proba-
17

---

bility of getting 2; times the i-th outcome with Ei= = m. The geodesic distance for multino-
mial distributions is identical to that of the categorical distributions up to scaling constant [44]:
denal (14, v)= /ndea (14, v) 2/n arccos
Σ Piy
(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) [28] transformed the one-hot distribution to neş 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 inacces- sibility issue of the Dirac measure DirichletFM. We have demonstrated that SFM consistently outperformed the naive multinomial diffusion on the Text8 datasets, indicating the effectiveness of tracing the true geodesics.
B Likelihood Calculation
In this section, provide more concrete mathematical details regarding the exact likelihood cal- culation 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 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,
u, V to denote elements in ((()). or categorical distributions, and use q to denote elements in ((P())). or distributions over categorical distributions. We further use NLL refer to log p(u) for a specific μl and BPD for categorical likelihood - log, p(0(p) where 8 is the ground truth Dirac measure (one-hot distribution). Note this definition of BPD (BPC) follows the convention in autore- gressive language models where the conditional probabilities for the next token are calculated and averaged [16, 23]. In contrast, most previous work drew equivalence between these two concepts.
B.1 Negative Likelioelhhood (NLL) Calculation
Our SFM framework enjoys the exact negative log-likelihood calculation as a continuous normal- izing flow model [38, 14]. It is worth noting that negative NLLs are not uncommon [14]. As 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 function
p(0)
0 (0. 1).
(28)
TV / 0)
is easy the density approaches infinity approaches 1. Note that the aresine distri- bution is a special case of beta distribution Beta( ( z), which is in turn a special case of Dirichlet distribution Dir( 2: 2)- Each I can be naturally viewed as the parameter for the Bernoulli distribu- tion (2-class categorical distribution), so the arcsine distribution indeed defines 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: Σm= 0,8; where d; := 6zz is the Dirac measure over the underlying Dirac measure 5' on the discrete class i.
1
Variational Bound We followed [7] to derive a variational bound as marginalized over small neighborhood of the Dirac measure in (14) to obtain finite NLLs. The variational bound can be
18

---

derived via the standard variational approach as
log p(ò) = By(u)d) [log p(6)
p(ó, 14) p(a(5)
]
log
p(6, 14) q(6, (1)
a(m18) p(ul)) )
log
(29)
p(ó, pi)
= (9(
log
+ (Kc(((uul0))<((1))
q(al8)
p(6, p4)
2 (9(H))
log
g(p(8)
= -E [log + 1ogg
log
The posterior probability ((4(6) serves as the forward diffusion process in [6, 60] with closed-form expression that depends on the timestep. In our Riemannian flow matching formulation, however, such a forward process implicitly defined via the time-interpolation along the geodesics. Nonethe- less, we do know that the conditional distribution should approach the Dirac measure when / + 1 as the geodesic distance approaches P1 (1466) R 8. Therefore, as we are marginalized over a small neighborhood,
can define time-dependent posterior distribution (1(8) that approaches 8 as t + 1. In principle, any forward probability that approaches the Dirac measure at 1 with an analytical likelihood can be used to derive the lower bound. For example, 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
Po (I-t)-T, 0,
S Sttisiinin n
(30)
otherwise
where Po is the base probability of the uniform distribution over the simplex, μe 0, is the i- th component the categorical distribution μ. In other words, qe (x(6+) the -nterreerpol probability between the Dirac measure and the uniform distribution: (:(88*)) 18, + +1-t)(A) -- We have
log qe (u16*) log (n 1) log(1 -). t).
(31)
In practice, t 0.995 and provide more results in the ablation studies in Appendix D.2. The categorical likelihood p(5* (μ)) is the standard cross- cross-entropy loss, also referred to as the reconstruc- tion loss in some previous work [24]. It can be calculated
log (68 (4) log pkk
(32)
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 A--1 This can be efficiently done by normalizing n i.id random variables from the exponential distribution:
0 ~ Exp(1), =0 0,/
Oz. i= 1,2,.....
(33)
The uniform distribution over simplex is exactly the Dirichlet distribution Dir(1..... Therefore, the log-likelihood can be calculated
log Pa log [(n)
(34)
where I' is the Gamma function. Note that the prior likelihood is independent of the samples.
Change of Measure The diffcomorphism æ in Eq.(5) induces a pushforward measure on S s which can be described by the change of measure identity
TTT((()))) det (x)) P(x), x x()).
(35)
19

---

The change of measure term can be calculated using the change of volume formula with the canoni- cal coordinates {=: and {p: as
μi 31
31 COS 4p1
1/2 22
sin (1 008 P2
(36)
Pn-1 In-1,
In-1 = sin pin ....... sin Pa-2 cos Pn-1
where 0 s: S 7/2. The diffeomorphism - π (14) is given by
8 = cos² P
x2 = sin 2 01 cos? 62
(37)
-n-1 sin 41 .9..ssin'
sin" Pn
cos Pn-1-
The change measure term is
log det 11 (x) = -n log Σ
log Zi.
(38)
i=
Similarly, for the inverse mapping, the term
1
log det (T)) (n- 1))- log 2
log μ4-
(39)
2 1=1
Model Likelihood flow:
The model likelihood can viewed as the pushforward measure the learned
Pr (vt))pp-
(40)
We demonstrated in Sec.3.5 that this plus the base probability in (34) More specifically, the NLL can be obtained as the solution to the following ODE system back through time with the initial condition of # and log ODEE 0:
be calculated the integral of divergence in Eq.(12)
d dt
- ]-[- vx(xr) (
log ODE p.
(41)
For manifolds that can be embedded the ambient Euclidean space (c.g., simplex and sphere), the Riemannian divergence can be calculated as the normal Euclidean divergence div, div [14]. We further use Hutchinson's 's trace estimator [30] to efficiently estimate the divergence as
div (x) = EE= (v((r)))
(42)
The expectation can be efficiently computed using the vector-Jacobian product. For data with uni- tary dimension = (e.g., Swiss roll on simplex), we also calculate the exact divergence using the equality div u (r)), where J(w) is the Jacobian of the vector field The computa- tional cost for the exact divergence calculation scales quadratically with the data dimension D, as the interdependence between data dimensions makes it computationally expensive 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 decent estimation of the exact divergence.
Combining (31(),), (32), (34), (38), and (39) with Eq.(14) and (13), efficiently calculate the NLL given arbitrary Dirac measures as input. The NLL calculation scheme is visualized in Fig.2 and described in Alg.1.
B.2
(BPC) Calculation
We have demonstrated that the NI.L can reach negative infinity. contrast, the common definition of BPC with the cross-entropy loss logg p(0)p) is always non-negative and consistent with previous
20

---

Algorithm NLL Calculation for Discrete Data
1: Sample Ai ~ (plo) in (3.(3 and calculate - log qr (مx 6) and log p()lū))
2: Apply the diffeomorphism in Eq.(5) obtaín क = π [ and calculate log] det dr (x)))
3: Solve the ODE system in (41( to obtain and log log /
4: Apply -- to obtain 1A0 = x ---1 (xp) and calculate log det dr (po))-
5: Calculate the base log probability log (Ho)..
6: return NLL as in Eq.(14).
autoregressive language models [62]. However, as our flow-based model not autoregressive, we instead follow the variational formulation in [52] to calculate the alternative NLL and BPC for comparison. The ELBO in [52] can be calculated as
LELBO = (uer) Vo - log(Ri (Pt), m1) df
1
(43)
where exP (o M1); denotes the learned model's prediction for 1 given the current noise data pa-- This can be efficiently done by one-step prediction with the learned vector field as
Po(H) denotes the interpolation along the geodesic and ji (Ht)
(1 = exp,r. ((1 ())))
(44)
The standard Euclidean inner product used in Eq.(43) so can be understood as a weighted cross- entropy loss. Note that the integral in (q. has a singularity at t = 1, making it numerically unstable estimate using ODE solvers. Instead, we follow [52] apply the change of variable = log(1 - t) to reformulate the ELBO as
[% (og( m1) ds
LELBO = - Buy ((()
(45)
where t = 1 e Note that for large s, the corresponding t is very close to 1, so the integrand is t < x 10-5, so we simply set the upper limit
very close to zero. Indeed, s = 10 corresponds to 1- to 10 and used the Dopri5 solver to numerically estimate the ELBO. BPC be then calculated as CELBO/ log 2. This formulation also shares a similar form as the other ELBOs derived in [6, 12], and is thus comparable most existing models.
(43) was directly optimized in Markov chain-based methods like D3PM [6], MultiFlow [12]. and SEDD [39], and also autoregressive language models. In contrast, our SFM (and LinearFM) follows the continuous normalizing flow setup in [38], which enables the exact likelihood calculation instead of ELBOs. Therefore, trained on the alternative objective of minimizing the Riemannian vector field norm in (8), SFM did not directly try to minimize such an ELBO. Despite such an unfavorable evaluation compared to other baselines, SFM was still able achieve the second-best BPC, as demonstrated in Tab.2.
B.3 GPT-J-6B Likelihood
GPT-J-6B [67] is transformer-base large language model with 6B trainable parameters. We follow the pipeline in [12] 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:
LGPT
1 K
Σ log (wx/wix-1),
(46)
where w are tokens and K is the number of tokens. Similarly, the entropy is calculated on the empirical distribution of the tokens based on large number of generated texts. We noted that, as GPT-J-6B was not pre-trained on Text8, its NI.L does not necessarily reflect the true data distribu- tion in Text8, generated samples in Appendix D.3.
will demonstrate more concretely in the evaluation plot in Appendix D.I and
C Experimental Setup
In this section, we further describe the experimental setup, model architecture, and datascts.
21

---

C.1 Model Parameterization
Our SFM architecture is encoder-agnostic and can be applied 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 [66] as Emb [0. 1] R" We follow [14] manually project the predicted vector field onto the corresponding tangent space. For the spherical manifold, the projection can be described as
(244) = üt (zx) - (Iee (z,))))
(47)
For linear flow matching on the simplex, the projection can be described as
1 Σ (1)))
(x)))) = (x))))
(48)
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 et. 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 80GB NVIDIA A100 GPU for 6- 10 hours. Each model for Text8 was trained on four 80GB NVIDIA A100 GPUs for about days. We will further describe the encoders for each generation task in the following subsections.
trained on single
Algorithm 2 Training SFM
1: while not converged do
2: Sample noise distribution (A0 ~ (p( and target distribution μ41 ~ q(4)..
3: if optimal transport then
4:
Do batch OT assignments of po and fu according to the average statistical distances. 5: end if
6:
Apply the diffeomorphism in (5) to obtain an π (po) = ((M).
7: Sample / U00 1] and interpolate H, exPx, (1 (oog,, .n) using (19 and (20).
8:
Calculate the conditional vector field uf (zz|=0, x1) = f = log (x) ) - t).
9: Predict the vector field using v(z,, f) and optimize the SFM loss in (8).
10:
end while
C.2
Model Sampling
The sampling process from the trained model can described solving the differential equation -e vz(2,) from 0 to 1 with the initial conditional zo sampled from the prior noise distri- bution. Alternatively, we can write the solution as the integral of the learned time-dependent vector through time as
21 = = L (z( dt.
(49)
We always use the Dopri5 ODE solver [19] for our ODE sampling and NLL calculation. For Euler method with N discrete steps, the timesteps 0, 2/N, (N 1)/N are used with step size of 1/N. The sampling stage is described in Alg.3.
C.3
Swiss Roll
The Swiss roll on the 2-simplex is generated by normalizing the span of the 2-dimensional Swiss roll to [0 + e, - where s a small margin to make sure no point lies on the boundary. The third dimensional is automatically obtained μ43 1000 samples for a full batch training. The vector field predictor is based on simple multi-layer perceptions (MLPs) and can be described as
#1 - U2. We fixed random seed and generated
vz t) MLP(MMP(::)))MPPPP
(50)
where || denotes concatenation and we used hidden dimension of H 128. Fach model was
trained for 2000 epochs using full batch gradient descent with initial learning rate of 10-3, 1000
22

---

Algorithm 3 Sampling from SFM
1: Sample noise distribution Mo ~ (p()).
2: Apply the diffeomorphism in (5. obtain 20 = T(p)).
3: if ODE sampling then
4:
Solve 0t (1x) using Dopris ODE solver with initial condition Zo-
5: else
Euler method
6:
for + 0, 0,1/N,22/N, ....., (N- 1)/N do
7:
1/41/N expz, (0(x,,t)/N)
8: end for
9:
end if
10:
return frn = T
samples were sampled using the DopriS ODE solver, and NLL (based on Hutchinson's 's trace es- timator) 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 [60].
C.4 Binarized MNIST
We used the preprocessed binarized MNIST dataset from [53] which has a split of 50k/10k/10k We adopted the CNN- -iae- vector field predictor from [57] with additional additive time embeddings at each convolution layer's output. The model has 4 residual layers [25] and 4 refinement layers [37] with total number of 8M trainable parameters. All models were trained for 100k iterations with a batch size of 256 (approximately 510 epochs) with initial learning rate of 3 x 10
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 Inception V3 [61] 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.
C.5 Text8
We followed previous work [24, 6] to use fixed split of 90M/5M//M with a fixed sequence length of 256. We used 12-layer diffusion transformer (DiT) [49] based predictor, similar to the one used in [39]. The resultant model has 92.4M trainable parameters. Noticeably, our model size is smaller than [24] which used full 24-layer model, but is similar to [39]. 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 decay rate of 0.8 with a patience of and 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 IM characters) with the Dopris solver as described in Appendix B.2, and generated texts were also obtained using the DopriS ODE solver. GPT-J-6B NLL and the generation token entropy were calculated based on 4k generations length 256 (about IM characters). See Appendix D.3 for generated texts.
C.6 Promoter Design
We used the splits from the dataset paper [7] that assign Chromosome 10 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 total number 100k sequences with a context window of 1024 base pairs and added 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 13.3M trainable parameters. Our models were trained for 200k iterations with batch size of 256 and an initial learning rate
23

---

of 5 X 10 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 Fuler steps of generation.
D Additional Result
In this section, we provide additional results of the evaluation on Text8, additional ablation studies, and generated samples for each task.
D.1
Additional Result on Text8
GBT-J-6B NLL vs. Entropy
7.0
MultiFlow
SFM w/ OT
w/o OT
MultiFlow r
SEDD mask
6.5
D3PM
SEDD uniform
Autoregressive
6.0
Data Entropy
X random
GPT:j-6B
5.0
4.5
Data GPT NLL
4.0
5.0
6.0
6.5
7.5
Entropy
Figure 4: GPT-J-6B NLL versus sample entropy. For MultiFlow, D3PM, and autoregressive, the curve represents different logit temperatures from 0.5 Baseline data are from [12].
We further follow [12] (MultiFlow) to calculate the NLL using GPT-J-6B [67], a pre-trained autore- gressive 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 preferred. The GPT-J-6B NLL and the token entropy are plotted in 4ig.4, with the raw baseline data and the permission to reproduce provided by the authors of [12]. An additional random base- line that generates each character independently based 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, noted 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 MultiFlow it generated more repeated numbers with lower temperatures. In contrast, SFM achieved perceptually similar or even better results with more diversity.
D.2 Ablation Study
We provide ablation studies of the effect of different sampling methods, different sampling steps for the Euler methods, and different tmax ~ for NLL calculation. The results are provided in Tab.4. The effect of the variational timestep mmax matched the results from [7], which stated that a closer timestep to the target data would decrease the NI.L. The reasons we choose = 0.995 instead of higher values are: 1) we noticed numerical stability issues at a timestep very close to the target, 2) want 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 due the large divergence change at t R and naive Fuler steps tend to overestimate its contribution. Nonetheless, the Fuler method with more than 300 steps gave similar NL.L 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 i = 0.995 and 300 Euler steps. The final performance was quite similar between these two sampling methods.
Table NLL for different sampling meth- ods, sampling steps, and maxx
Table 5: NLL and FID for different sampling meth- ods with fmax = 0.995. For the Euler method, we used 300 Euler steps.
Method
#step
tmax
NLL
0,9995
-2.545 + 0.029
Method Model
NLL
FID.
0.999
-2,442 = 0.018
0.995
-1.689 = 0.031
SFM w/ OT -1.687 + 0.020
4.62
Fuler
0.99
-1,019 0.032
ODE
SFM w/o OT
-1.631 + 0.027
5.15
LinearFM
0.622 + 0.022
5.91
100
-1.77
0.020
500
0.995
11677 = 0.022
SFM w/ OT
-1.689 = 0.031
5.00
1000
1..44 = 0.024
Euler
SFM w/o OT
-1.653 + 0.028
4.86
LinearFM
0.499 0.022
6.47
ODE
0.995
-1.687 = 0.020
D.3
Generated Samples
C
6
a C
0
2
6
B
0
4 8
/
3
V
P
7
6
5
)
P
7 o
4
#
5
9
4
6
5 1
9
6
6
1
6
1/50
8
9
0
3
0
0
4
b
0
8
9
2
C
2
3
2
A
7 0
5
9
1
0
#
A 3
9
0
2
D
Ground Truth
SFM 9
OT,
ODE
SFM w/o OT,
ODE
SFM w/o R, ODE
4
4
2
5
5
7
G
8
2
0
9
0 0
4
7
8
4
5
9
6
9
7
7
D
2
8
y 5
9
5
7
9
0
9
5
7
f #
6
2
8
4
9
8
4
7
7
7 2
6
4
4
5
9
1
4 O 4 G
24
4
2
9
9
O
LincarFM. ODE
DDSM
D3PM
DFM
7
8
5
3
4
2
0
6
0
3
9
8
8
0
5
3
9
2
2
6
9
Q
6
0
6
6
0
/
5
7
7
0
4
O
5
3
2
C
3
3
3
2
5
2
1 0
2
4
2
9
2
9
4
7
4
0
O
7
SFM w/ OT, Euler 100
SFM w/ OT, Euler 300
SFM w/ OT, Euler 500
SFM w/ OT, Euler 1000
Figure 5: Generated samples of the binarized MNIST dataset from various models and different sampling settings.
The ground truth images and generated samples of the binarized MNIST dataset using various mod- els and different sampling settings are demonstrated Fig.5. The lower row demonstrates generated samples with different numbers of Euler steps.
25

---

For the Text8 dataset, generated texts, NLL, entropy, and per sample NLL evaluated by GPT-J-6B us- ing different models and sampling methods are demonstrated in 6,b.6, with one sample with an NLL chosen around the average, one above the average, and one below the average for each model to pro- vide 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, is incompatible with temperature-tuning.
Table 6: Generated samples and GPT-J-6B NLLs. For each model, the average NLL and entropy calculated on 4k generations are also provided. The NLL marked for each generation the sample NLL, as evaluated by GPT-J-6B.
SFM w/ OT, ODE, NLL: 6.762, Entropy: 7.340
zero zaro ZATO more, an well an. the needed all, of charch, the comntry a. higner, upconing. bank, the, country. comment, on quebec, a dite www.oooccccccc.oiooo the account of wil.dd__oo___.... htrpo___lo______oeeeeeee www.pcccaaccccccooo coce, tain, three, 2ero seeeee__
NLL: 6.336
wtt.eee_ettoniiiineeeeeeeeeseee wwt.:iinetterrr.oomm .ff soothh_ot, of ttttttt___uue cout_eott__ooo peott___oo__oeete pooddddaa___ta the
such, as in, muscargeeeecc of, coincinacion with nous, auch, an, adler rartie, the hilly, pats, exedhan, of norowls a, night, of blond, the trween
NLL: 6.805
to___ro_roore______._________e lett__ittett____..c__ this__tt_________ lase__lbb___ee pe_i__i_nnn_bo bol.enolntttsss wwwppppppptnnnoeeeee taisxro t_ot_ wwwenneerggg.e...... wwwlge.iilloeeeeeooo an_ulaaaarraaa.oaaaaaaaaaaaa__ee www.eertoin__teee wtapt_haa_______
toorstor_ twenni mothherr, toury. toury.ghout, phoutt_forrt for, his, biittry _ooo blaster, the, three, atallran, faailyy aourow,, including. fhe_file A_ro for_biii_
NLL: 7.522 bloebogeebbookee..oo books,, moy, woy_e_sthiite stttbiite, e_nodit
nance cine, author www.phigggyy____e lacaded the cecoocddd_dd haroour open_ nource for whichhhhhoorcr orrie, changed, the blseboge,
SFM w/o OT, ODE, NLL: 6.811, Entropy: 7.387
becane mmoon thtthh__ee wwwppehhhccaalaaaeeeeeeeeeeeeeee www.ey_to..oooeeoooooeeoooooooot toltcyytt_theeooomo hewren areeeeeeeddddddeeeeeeee in_ stc_nidsic__d ot__iio_________t_eee_eeeeeeeee_ thetth___ee Jechneit, ttw.ccaaaattttt Lime Ltt.___o_oe_________ jeno_.oj.nne_... wattta.valleyooea... one___ inochules, tettttteeaaa yaas veriiied by vy_xany_com charberrlyy...c... brollary, applicatioon. including, thore, eynoohooo with, orbate somo, of, tett
NLL: 6.407
cable sos_nast___aae badtaittdeeeee_eeeeeeee___e to porro____ooo eix eight, nine rade bofavor, parrr_.eer_peee prr.atrtteerr, of disruppptonn this, hitspplaooor, would, forfariigg_to to. www__thooe__t__eeeeeeeeeeeeeeeee
NLL: 6.819
the surrent_ccuntaee bus, this ieid__ for_sou__aaaaeeeeeee fww_ssuuddellaaaeaaaaaaaaa deeela_even, wweeeiiii four one, three chrroebit__otii citt_aof__celeree soo___o_ aatter__,a,_u a_ wtw.ttttootttrrroooooo
nomic, ancestor, mohn berg hiaret rrr_rrdd__rreee rthonstrak, utter technoloyy baddendin rodels, on, bendrays, hypothesser, anti, ser dynami bave, Deeeree_be___eeeeeeeeeeeeeeeeeee int.llggeettttt.... to wo_.evelop..eee european antoppannnttroooiicc astronoxic, comifice in, the tho.ppoduction.o.... ten, conificos, of derddddolpppao and prine
NLL: 7.479
LinearEM, ODE, NLL: 6.935, Entropy: 7.356
is seslelllddddee ga.gaaiik, college, in, thw__thteeeeeeee htw____iv___sssco_aaaaaaae seasoo_of, feason, as, wiwww the atnospheze, is sa__aseed_ aftrr the list_ called bia, before, iun
NL.:: 6.466
s_collloe__aa aaree_______aaa at__stuprord, universisy_of
owwddsyooooooooooooooooooooooeee london, also___ssaaa_ castbrsgdddddee..oo... burrough, burrougha.Jaariann hearians_cooliaaaaaeeet cow.loooootttlaaeee shthhhhhhhhhhhheee_eee_aa__eeeee ye_ apwllo__ooeencoooeeeooeeeeeeeeee colloge_na
cao_ www__eight_eee eight_ parhapp that, one__nn_______ s_strees ronan, travuapared, the practices of, wel.telllatttt.o apeazerahhp settled and arayaaayooo _ff, the.
NL.L: 6.935
two_set_of wtw.iooeoeeeiooeeeeeeeee
www.datiennnoooooooooooooooooooe otettt__eeeeeeeeeeeeeee batbe_coofeeerrtte honovater_ wwwpduugtteeetrreeeeeeeee ctw_ttt___vi__ar_________ or dan, brom_roo siinn_john__inee
devel_nortang_oe alwpuaarrrdyaaa bicckl_____b_uutee bet__vadde___eeeoooooooaeeeeeeee wwwpolve_eeeeeeeeee the of auctiosal vittl_daag_abee aueuttffouuu four_ pet_jekks_ttteeeee there_vas_a_aeeeeeee woiiaad_probbenn sopla,ste
NLL: 7.454
hat_aoooooiioeaeeeeeee recni sigond_pp griaf of, cultura of the siddle bone, laland, for nore, dww.gaaaiiiiiimmmmmmm kickk trade, handa, and rangers, has
MultiFlow, T 1, NLL: 6.728, Entropy 7.387
goosecat, productions, nine, n, farous, veran, stanlay gtadday.ggoddann jarderww__rrrrrrreereee ncmbbb, out, zoon moonnnovia,, barrbbaaabbbrr laroy, barbara, levix, and brw.ddr purcesn away attaar,,aaa aka,
snap opera, by andrew
sea nonaters,
late, nine
NLL: 5.906
she hill othalarmeach sochrans, wileasn munthar, cwar lorba nhannns bardare nho word tore, lan, 1im zu phaile, gael, cangellauigg lan gall, iamonta, tilbonns guilechtt shuechtaneeerrr weaeccatteettgglllleeeee waiillaieegustnnn cuntmole_asthhllee aathache, cuichent, seiihhhttcccttt
NLL: 6.648
thuis_ woilllith leois wlwi..gloonggillll
Ida, itx replacerent, and or, not, rist para, decabeod, ran, and dreat, n. ak, or ubangoatrndes, dislogie or comos cainns ax fallows nake, calawy_comanne _eaaaaa_____e__________ _ge.t_ www.oaaaaaatt att_statii_pppeeee probllan reducw, itttoiitttttttooooooeeeoo percppptontttonnncco for wewelllattttann man, juvish, viewa, the
NLL: 7.426
MultiFlow, T = 0.9, NLL: 6.249, Entropy: 7.240
inisr, wwwreeoorerreeeee www.rrttttteeeeeoee Beciaaee__e________e sev_jork,_..g v_g_ wow.ad_oldddd..mm one. ww__iine__nnoo. t90, sonen, twxteexxxeeeeeeeeeee hertggeeedbbbttrrrrr.. waayeeaddadddd_____e Lyan, lucky poncy_pecher, cnn, cnne_iine mine
NLL: 5.277
esight_eight_tte moderndddeneee
lyceun_pew__oore harpprrrrrrrr mead,
nine, pew.eeeeeeeeeee aereen_eegt__eeee seaptttaarro__uoe gewrggooggggooo one, mine cnts, one nodern, hww.hhtttrrr.. _ex
tight_it_deplete__e stt_t_ttoteeeeeeeeteeee of wt__th_etarters_..t the hig.high_aa_s land_ of wtt.ttooo__lllleeeeeeeeee twwpppapp__oo___p bon, mostootteeeee pucaluucassnn qnd_annttia__eeeeeeeeo aand__dd________________________ asd_ens_ deiddd_jaa de_asaa hencees___eeee______ aea soo_iandoue_pp peteooeeeppecee____ee_e tte_torch_cmm crach,crc_reee iitoo stt_eette__pcceeee pplaces_to to, to_eato hence to seed out, their, blood_ fron, thea, und_ thee_ttt____. tooha
NLL: 6.352
1on_ct, of ot__jstt_zolsceeeooooooo_ wwweeoooe___otooeoooooee___eeeee bob www.eprrrgiieddd.coo elurs_botil bow.se.ooiing.iiiii. tonny_trate wtt.ttrttteeeeeeeee feed_ and chaccig.rroceee rocked doos, oft off_r_oo_ r.foe_ both_eides__tee .an, uon_uneanted, and violent, violence and the, ggined, a recene, pecen the, trial, van, the. clach john, d, d.d the firet, fttpsftttttttlll__ttttpt stich
fonny.
NLL: 7.074
MultiFlow, T' 0.8, NLL: 5,552, Entropy: 6.840
eix_ tour alexander_ Jerodd_rueeiarr, priber, d_one _ninennco coee_oss_eeeee
ctte_oosee_e______eeeeeeeeeeeeee cae.eooooeeeeeee wight__tsssss seven, cantel chase, american, author di one, minoonnneee.e thre
NI. 4.090
e_two_ose www__eigh__seeeeeeeeeeeesessssse sec___seeeee_eeeee wwwehttiggoiiiieeeeeeeeeeeee ses_s__ao__________e_e lurodaaaraaee posiiaaasttttnnnnmmm cosposer __one_sine cue__sin_soi__feeee flu_foour_oo ooe, e-gätt_eight_tae two rebertebrt___ v.kissir wtwsstoooossto.ssrreeaa.o axericaar_pub
usadlushab, wlddines sahira j __singh____a wewe_ive_eeeeeeeeeeeeeeeeeeeeeee five__eren_nin nine, jajaaaaaiaaaee....a _لاgن_ jjjandiaara Saandiarajaghar jughar, rohaaj batsssittaauue singh_ one, sive_iiee t_zero t_ sardoyaz_ teachings of parraaataaarraaaomm hatteeebbrrrrrnnnn_oooe sacee__io_____eeeeeeee eight ons pranajaahr, Jaharkars_ jogyar paaprrggrnooooooom of joar _at_oo__tttte ano_five_nin
NI.L: 5.223
o__po_o__isppiieeeeeeeee y daughter god_ pod_originn originates, to, as.oo_ss__.pr__e of inspiration, all_aas_dd of whes, thor_sball_uiteeee witeess, the totortttotionn you and will your and
NL..: 6.840
neighhour bour and reupskyour_ andd lov_ wiseeieeeeeeeeeee www.e_ooop_cccceoooooooo_o speppedeedddd____.....opo thy, beautifut
26

---

Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have accurately outlined major contributions and application scope in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations in Sec.6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairnsss.
While the authors might fear that complete honesty about limitations might used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
27

---

Justification:
We have provided theoretical background and assumptions in Sec.2.1 and
Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have provided detailed experimental setup information in Appendix C for reproducing all the experiments in this work. Our code is available at https: //github. com/ taaeeensstaaissial-flloo..
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or the contribution it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation,
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
28

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
The datasets used in this work are publicly available. Detailed experimental setup information provided in Appendix C for reproducing all experiments, and code is available at httpss// thub. cwwccc/cccceeeceee wwwr--ccnn/stttiiieeeeeeeeeseeee -al.siiiccclooo-mmciiaaaa
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines futtps://nipe cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so No is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips cttpssccccccccccccccpccccccccccc cooollccggides.e / ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We have included the detailed experimental settings in Appendix C. For data splits and hyperparameter choosing, we mostly followed previous work without hyperpa- rameter search.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yss]
Justification: For most results of our own models, provided standard deviations of the scores and metrics. Other collective metrics were averaged over large number of samples to eliminate the impact of randomness. More details were provided in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
29

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of eπrоrs is not verified.
For asymmetric distributions, the authors should careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have provided GPU specifications in Appendix C.I.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well 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 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https /ttp///eerrpp.ee/ee//ee cc/ htcppuulic/EEhtteeee ctwpe/thhhiuuieeeeeeeeeeeaeea
Answer: [Yes]
Justification:
Our work conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
30

---

Justification: We describe the potential societal impacts here. Our proposed model is gen- eral discrete generative framework that can be applied to various domains in machine learn- ing 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimesss considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific 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 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 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 unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion 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).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 allow for controlled use of the model, for example by re- quiring 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 best faith effort.
12. Licenses for existing assets
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:
31

---

Answer: [Yes]
Justification: All code, papers, and assets properly credited and cited in work.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswi paperswithcode thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: Our work did not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
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.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
32

---

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 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 equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and 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.
33


Reviews:
[{"review_id": "BqitaGkols", "reviewer": "Reviewer_xYE5", "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.", "questions": "### Main questions / concerns:\n\n_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.\n- How do you chose $\\nu$? \n- Can you say that the Radon-Nikodym derivative is defined for all $\\mu$ that you use in practice?\n- Does this change of measure (Eq 11) correspond to some $\\frac{d\\mu_t}{d\\nu}$, i.e., fixed reference measure?\n\nI 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.\n\n_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. \n- Did you use the same g for NLL computation? \n- Which g did you use (Euclidean or Riemannian)?\n\n_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?\n\n_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.\n- Why should we care if the Fisher information matrix shows up as the Hessian of KL? What does this imply?\n- Can you show that the geodesic under this Riemannian metric is following the \"natural gradient\"? Is this what you wanted to claim here?\n\n\n### Minor suggestions:\n\n- 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.)\n- The \"target\" space --> \"tangent\" in Sec 3.1.", "limitations": "Authors adequately addressed limitations.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Well written. I found the writing clear and to the point.\n- A natural extension of existing works to flow matching on the manifold of discrete pmfs.", "weaknesses": "Overall, I feel the idea is a straightforward adoption of prior work on riemannian FM in terms of novelty, but the paper is nicely presented and packaged together. The topic of generative flows with discrete data is also timely.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appreciation of our work's clear presentation as a natural and timely extension of Riemannian FM on discrete data. We thank your thoughtful suggestions and will fix typos and clarify minibatch-OT in the revision. We'd like to address your comments as follows.\n\n## Q1 Regarding likelihood definition\n\nWe mentioned the Radon-Nikodym derivative in 2.1 as the definition of probability density functions on general manifolds. In our likelihood derivation (Sec.3.5), the density $\\frac{dP}{dV}$ can be well-defined over the statistical manifold which is finite-dimensional ($n-1$ dimensional for $n$ categories) using the Riemannian volume as the reference measure. This is in contrast to works like functional FM where data points are continuous functions lying on the infinite-dimensional Hilbert space. We will clarify this as you suggested.\n\nIn terms of the choice of measure, we make sure to use a consistent reference measure for all change-of-measures (COM) in Sec 3.5. Note that the Riemannian volume on the sphere coincides with the Euclidean one (see response to Q2), and only the COM for the two transformations $\\pi$,$\\pi^{-1}$ will be impacted by the choice of measure. To make the NLL comparable to Euclidean flows, we have kept the reference measure consistent as the Euclidean volume form in the equations in Appendix B.1. We will make these mathematical details clearer and explicitly state the choice of volume in the revised manuscript.\n\nFor categorical distributions $\\mu$, we used the canonical counting measure over the discrete sample space $\\mathcal{X}=\\{0,1,\\dots,n\\}$.\n\n\n## Q2 Regarding likelihood computation\n\nWe thank your insightful question regarding the divergence in NLL calculation. You are completely right that even though vector field is learned with Riemannian geometry, we can still compute likelihood with Euclidean divergence. As a matter of fact, we have taken comparability into consideration in our derivation of likelihood. Note that for manifolds that can be embedded in the ambient Euclidean space (e.g., simplex and sphere), the Riemannian divergence is the same as Euclidean divergence (for embeddable manifolds, we have $\\log|\\det g|\\equiv0$). By using the Euclidean volume form as the reference measure in all COMs and computing prior likelihood on Euclidean simplex, **we effectively used Euclidean divergence on the simplex** for SFM and LinearFM to guarantee that the results are comparable.\n\nIt is possible, though, to calculate the Riemannian likelihood for SFM. To do this, we need to make several adaptations. The prior likelihood should be adjusted to $p_0=\\Gamma(n)/\\sqrt{|\\det g|}$ such that the integral $\\int_\\Delta p_0dV=\\int_\\Omega p_0\\sqrt{|\\det g|}d\\mu_1\\cdots d\\mu_{n-1}=1$. The change of measure term in Eq.37 & 38 also needs to be adjusted as $\\log|\\det d\\pi|=(n-1)\\log 2$ and $\\log|\\det d\\pi^{-1}|=-(n-1)\\log 2$. \n\n\n## Q3 Regarding choice of metric in designing conditional $u_t$\n\nThank you for mentioning the nice property of Riemannian $u_t$ being parallel to the boundary. We'd like to mention that our visualization of the logarithm map can be viewed as a demonstration of $u_t$'s direction  (up to a constant scaling factor that depends on $t$), as the vector field can be calculated in terms of the logarithm map as $u_t(\\mu_t|\\mu_0,\\mu_1)=\\log_{\\mu_t}(\\mu_1)/(1-t)$ (assuming a linearly decreasing geodesic distance). We indeed noticed that the Riemannian structure induced by the Fisher information metric leads to vector fields more parallel to the boundaries. This can also be demonstrated mathematically by looking at the logarithm map in Eq.23. Consider the direction term $\\sqrt{\\mu\\odot\\nu}-\\langle\\sqrt{\\mu},\\sqrt{\\nu}\\rangle \\mu$. For $\\mu$ close to the boundary with $\\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$. \n\nWe hypothesize that one potential benefit of such a curved geometry over the naive Euclidean geometry is that the former helps **prevent overshooting across the boundaries**. Specifically, consider a target point on the boundary. The Euclidean vector field will continue to push the points outside the manifold, whereas the Riemannian vector field tends to travel parallel to that boundary to prevent going across the manifold. The sphere exponential map also naturally ensures the point (after transformation back) lies on the simplex. Once overshooting happens during sampling, the model may exhibit undefined behavior as it was never trained on points outside the manifold. We will emphasize this benefit more clearly in our future revision.\n\n\n## Q4 Regarding the connection to natural gradient\n\nPlease see the common rebuttal for details.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "83bW730PJl", "reviewer": "Reviewer_SAto", "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.", "questions": "Questions:\n\n- 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?\n- 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.\n- Table  4,5 “ODE solver” is what solver? Euler is an “ODE solver”. Could the authors specify which solver was used?\n- Why is LinearFM quite different between 300 steps and “ODE Solver” in Table 5?\n- Table 5 does not look like an FID score to me, which must be strictly positive. Perhaps this is an NLL?\n- Why is Multiflow cited but not compared to in the text8 task?\n\nMinor comments:\n\n- “we define the ODE log-likelihood as” — This is more accurately described as a change in likelihood perhaps?\n- “Swill roll” —> “Swiss roll” perhaps is meant in multiple places?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Great figures with above average effort and conceptual clarity that make it very easy to understand the method and theory geometrically.\n- Simple, but seemingly effective addition to flow matching on data supported on a simplex.\n- Could be quite significant if it is truly competitive to diffusion-based approaches.", "weaknesses": "- I’m quite concerned about the “Pseudo log likelihood” metric being compared to log likelihood. While the authors know and say it is incomparable, it is for some reason still compared in Tables 1 and 2 directly to the discrete NLL. I don’t understand why this is done. I might recommend an evaluation similar to MultiFlow for text8, where the evaluation is done on generated samples, and can be used to compare fully discrete, vs. the continuous state approach taken here. Its extremely concerning how much the pseudo-NLL changes based on the choice of a seemingly arbitrary threshold $t_{max}$ and the number of steps (Table 4). To me that means this score should either not be used, or should at most only be used to compare like methods.\n- 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.\n\nIf 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. \n\nMultiFlow: https://arxiv.org/pdf/2402.04997", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your recognition of our work's clear presentation and novel geometric perspective for discrete generation. We'd like to clarify that Text8 results of MultiFlow were added on Jun 5 after the NeurIPS ddl, thus we were unable to compare them although it was cited.\n\nWe wish to emphasize that **our major contributions are beyond the Text8 experiment**. SFM is a general generative framework for measures of categorical distributions, and we have demonstrated **the effectiveness of SFM across diverse domains** including computer vision (binarized MNIST), bioinformatics (DNA design), and NLP(Text8). We also demonstrated performance improvements **using evaluation metrics independent of likelihood calculation** (FID & Sei MSE). We further established the underlying connection of SFM to natural gradient and proposed likelihood formulation with comparability in mind.\n\nWe respectfully disagree that Text8 is the only non-toy dataset. DNA design has a **tangible societal impact**, and DDSM & DirichletFM all predominantly focused on this task. The data dimension of MNIST (784) and DNA (1024) is also larger than Text8 (256). Nonetheless, we understand the importance of text8 evaluation and now include additional results in the common rebuttal following your suggestions. Below we provide more details and justifications in response to your comments.\n\n## Q1 Choice of evaluation metrics\n\nWe're afraid that you might have misinterpreted our evaluation setup. For Swiss roll and MNIST, **all baselines are diffusion or flow models with continuous parameterization**. We never compared NLLs with *discrete-space likelihood* from discrete models (e.g. autoregressive) in these tasks. Secondly, **all NLLs used in discrete tasks (Tab.1,2,4,5) are NOT pseudo-likelihoods**. They are based on the variational bond described in Eq.14 and are comparable to ELBOs from the diffusion models we compared. Only BPC for Text8 (Tab.2) was calculated based on pseudo-likelihoods as described in Appendix B.2. According to the original PLL paper, it provides a reasonable evaluation metric between non-autoregressive models.\n\nWe used the **identically variational formulation** in the DDSM paper to be comparable. We also provided **FID scores that do not rely on NLL calculation**, for which our models still outperformed other baselines by a margin. More discussion on the validity of such a variational likelihood can be found in Sec.3.5 of DDSM.\n\nWe justified our choice of $t_\\text{max}=0.995$ in the calculation of the variational bound for NLL in Appendix D.1. We carefully chose this value to make sure it was comparable to DDSM settings. Our ablation study on the choice of $t_\\text{max}$ agreed with that in DDSM, where $t$ closer to data gives a tighter estimated bound (it can't get arbitrarily low). We believe it is reasonable to compare DDSM and LinearFM using similar $t_\\text{max}$. \n\n## Q2 Additional results on Text8\n\nWe provide GPT-NLL vs token entropy, and representative samples for SFM and MultiFlow with $T=0.8,0.9,1$ for visual quality check. SFM has closer to data token entropy and comparable GPT-NLL with other methods in MultiFlow's result, although slightly worse on the latter. We respectfully disagree that our samples \"look extremely poor\", especially when compared with MultiFlow results with matching entropy ($T=1$). The perceptual quality of SFM and MultiFlow are very close, with noticeable misspellings in both. We also found GPT-NLL very unreliable and can be easily fooled by random strings. As an example, the following sample from MultiFlow at $T=1$ *looks* even worse despite a GPT NLL of 6.648.\n```\nshe hill obhalarnnach eochrans eileann munthar cearlomha mhaonna  tardare mho mord tore lian linn mu phaile gael cangallauig laothuis guilleith leois glion guildh lara gall innonte tilbonne guilecht shuachtansert guillaste guatnaoic asthache cuichant conai\n```\n **We will include these results on Text8 for a more comprehensive evaluation as you suggested, and tone down on our original claims. We sincerely hope that our contributions besides text8 can be valued properly.**\n## Q3 Change of measure term\nThe change of measure term describes the change of the log-likelihood between different manifolds as $\\log p_S(x)=\\log p_\\Delta(\\mu)+\\log |\\det d\\pi(\\mu)|$ where $x=\\sqrt{\\mu}$. Although the density is always well-defined, the log-likelihood is undefined at the boundary, as it involves taking the logarithm of a zero probability density (see Eq.37 & 38).\n\n## Q4 Forward probability $q(\\mu|\\delta)$\n\nThe detailed definition of the forward probability $q(\\mu|\\delta)$ was provided in Eq.29 & 30 in Appendix B.1. The forward diffusion process $q_t$ defines a small neighborhood for variational estimation. In DDSM, the authors followed the fixed probability path of the Jacobi diffusion process with known forward probability. In our flow-based setting, we can use simpler indicator measures as the linear interpolation between the delta measure and the uniform distribution on the simplex.\n\n\n\n## Q5 ODE solver\n\nAs described in Sec.4.2, all generated results and NLL calculations are based on the Dopri5 solver unless otherwise specified. The Dopri5 solver is an adaptive solver with a good balance between accuracy and efficiency.\n\n\n## Q6 LinearFM Sensitivity to solvers\n\nWe noted that LinearFM tends to have very negative divergence around the one-hot distribution. Therefore, Euler step may overestimate the contribution of divergence in this case while adaptive solvers can provide a more accurate result. In contrast, as the SFM vector field is defined on the sphere, it is less sensible to solvers and results in a more stable divergence.\n\n\n## Q7 Typos & minor issues\n\nThe results in Tab.5 are NLLs. The ODE likelihood defined in Eq.12 is indeed the change of log-likelihood: $\\log p^\\text{ODE}:=\\log p(\\mu_1)-\\log p(\\mu_0)$. We thank you for pointing out these issues, and we will fix them in the revised manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dbl20oe3Ff", "reviewer": "Reviewer_hmQs", "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.\n\nIn 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.\n\nFinally, 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.\n\n[1] Chen, Ricky TQ, and Yaron Lipman. \"Riemannian flow matching on general geometries.\" *arXiv preprint arXiv:2302.03660* (2023).", "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": "The paper lacks a discussion of the applicability of the proposed abstraction to the generation of discrete data (see Weaknesses).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a complete study, i.e. it approaches an important problem in the field (generative modelling for discrete data) proposes a reasonable approach, presents the proposed approach in a comprehensible way, and provides minimal necessary empirical study of the idea.", "weaknesses": "Although the paper does not raise major concerns and overall satisfies the criteria of a NeurIPS publication, the practical motivation of the method is insufficient. Indeed, the problem of generating categorical distributions is interesting and the paper answers this question definitively (to a reasonable extent). However, then the authors quickly jump to the conclusion that a natural way to model discrete data is to sample a categorical distribution rather than to sample *from* a categorical distribution. Indeed, in Algorithm 2, the authors don’t even mention how they sample these categorical distributions (as I understand they consider every sample to define a point mass on the statistical manifold).\n\nFirst, 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.\n\nThere are some minor concerns regarding the presentation:\n\n- The name Statistical Flow Matching does not match the title of the paper.\n- Line 29. It is not clear what `assumptions that fail to capitulate` mean.\n- Line 52. It is not clear what `existing models often fail due to impromptu prior assumptions` means.\n- Line 105 requires citations.\n- In lines 112-117, please specify that equipping the manifold with the Fisher-Rao metric.\n- The sampling notation in Eq. (8) is confusing due to the diffeomorphism applied to a probability density.\n- There is a typo in line 118. The tangent space is called `the target space`.\n- Section 3.3 reads as a preliminary section on the natural gradient rather than a connection of the proposed method to the optimization literature.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your high recognition of our work's significance in the realm of discrete generative modeling and for raising interesting questions and suggestions. We will fix the typo and citation in the revision. Our responses to your questions are as follows:\n\n## Q1 Motivations\n\n### Sampling *over* categorical distributions instead of *from* a single categorical distribution\n\nWe agree that both formulations can be effectively used on discrete sampling, but their compatibility with different generative models may vary. It is important to note that SFM (and the majority of diffusion/flow models, VAEs, and GANs) are not autoregressive (AR). Unlike LM which predicts logits for the next single token deterministically given context, diffusion and flow models **operate directly on the joint distributions of all tokens**. Therefore, it is inefficient to learn a single set of deterministic categorical distributions to capture diverse joint distributions. Alternatively, the majority of non-AR methods seek to transform from prior noise distribution (e.g. multivariant Gaussians) to target joint distribution, effectively resulting in a two-step formulation where stochastic samples over joint distribution space are drawn first (i.e., logits) and then decoded with a second sampling step (argmax or softmax). Secondly, diffusion/flow models are naturally defined over continuous space. To leverage such a framework, existing discrete diffusion/flow models utilize relaxation to operate on logits which are then projected to probability simplex, similar to the statistical manifold we defined but without a properly defined geometry. In other words, the two-step sampling in SFM is not a niche design choice but rather **universal for discrete diffusion and flow models** (e.g., D3PM, SEDD, BFN, LinearFlow). \n\nOur distinction to prior works is the proposal of a better objective for training the flow model over the categorical distributions (first step), where we equip the manifold with a Riemannian metric. For the second step, we always used argmax when decoding discrete samples.\n\n### Fisher-Rao metric\n\nWe would first like to clarify a potential misunderstanding. During training, we did not use the Fisher-Rao metric for calculating the distance between two sequences (i.e., discrete samples). Instead, a metric is **necessary for defining probability paths along the geodesic** from the prior distribution (usually uniform over the simplex) towards the target distribution (concentrated around the vertex). The role of the Fisher metric is to equip the manifold of categorical distributions with a Riemannian structure on which vector fields can be calculated as the learning target. The Fisher information metric is the natural canonical Riemannian metric on the statistical manifold. It also enjoys the benefit of following the \"steepest\" direction of the natural gradient of decreasing KL divergence, as we have mentioned in Sec.3.3. \n\nNevertheless, we respectfully disagree that treating discrete data as onehot distributions and using Fisher-Rao distance over them is unnatural. In fact, **discrete data have been commonly treated as distributions in classic discrete generative model training**, when using cross-entropy loss. Fisher-Rao distance is yet another metric between distributions, which indicates the steepest direction for KL divergence minimization. We will make this clear in the revision to better motivate this setup as suggested by the reviewer.\n\n\n## Q2 Name of SFM & title\n\nAs mentioned in the previous text and Sec.3, our proposed SFM is a **general generative framework for modeling measures (distributions) over the statistical manifold** of a family of parameterized distributions. In this work, we presented the realization of our framework on categorical distributions which have wide applications in various discrete generation domains. Therefore, we wish to keep the general naming of *statistical* FM and the *categorical* in the title to inform the audience about the application. We will think of a better title to avoid confusion.\n\n\n## Q3 Impromptu assumptions in prior work\n\nFor Line 29, we want to convey that the naive Euclidean assumption of the probability simplex or the logit space does not capture the true geometry of the statistical manifold. Similarly, in Line 52, we want to emphasize that the Euclidean assumption in previous existing models does not have solid mathematical grounds, which may lead to worse performance on statistical manifolds. \n\n\n## Q4 Fisher-Rao metric\n\nThe Fisher information matrix for categorical distributions is $g_{jk}(\\mu)=\\frac{\\delta_{jk}}{\\mu_j}+\\frac{1}{\\mu_n}$ for $1\\le i,j,\\le n-1$, where $\\delta_{jk}$ is the Kronecker delta. Substituting this into $\\langle u,v\\rangle_\\mu$ leads to the Riemannian inner product in Eq.4.\n\n\n## Q5 Sampling in Eq.8\n\nWe will rewrite the sampling as $x_0\\sim\\pi_*(p_0(\\mu))$ where $\\pi_*(p_0)$ denotes the standard pushforward measure induced by the diffeomorphism $\\pi$. This is equivalent to first sampling $\\mu_0\\sim p_0(\\mu)$ and then taking $x_0=\\pi(\\mu_0)$.\n\n\n## Q6 Connection to natural gradient\n\nPlease see the common rebuttal for details.\n\n## Q7 Non one-hot data\n\nWe really appreciate your recognition of the potential usage of SFM on general categorical distributions besides one-hot. In the paper, we have provided the Swiss Roll example as a demonstration, where other baselines with strong priors may fail. We believe that SFM could be of interest to Bayesian inference, distributional RL, and other tasks that involve sampling over distributions. We defer these use cases to our future work to extend SFM due to limited time and resources.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kSkuvSdPlN", "reviewer": "Reviewer_WPWv", "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.", "questions": "By the way, what's the difference between this paper and https://arxiv.org/abs/2405.14664 ?\n\nI understand and follow well the core paper, so I have only a few minor questions:\n1) 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.\n2) 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?\n3) 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.\n4) 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.\n5) 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": "see \"weaknesses\" section.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The method is well founded and the geometric insights are inspiring and seem \"right\".\n\nIt seems that the method performs well, as indicated by experiments.\n\nThe paper is well written and clear, easy to read execept for a few hiccups (mentioned in \"questions\" section).", "weaknesses": "The Riemannian metric on the subset of the sphere $\\{x\\in \\mathbb R^n: |x|^2=2, x_i\\geq 0 \\text{ for }1\\le i\\le n\\}$ is not that different (in fact a $C^1$-close deformation of) the one on the actual \"flat\" simplex. So why would the results be so much different between the two cases? The grounding for a big gap from this setup to the standard one is not clear. \n\nIt is also not clear why this method would have to perform better than Stark et al. \"Dirichlet flow\" paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your high recognition of our work's novel geometric insight and clear delivery. We will address your questions and concerns as follows.\n\n## Q1 Advantage of Riemannian structure\n\nUnlike LinearFM which assumes a flat simplex and straight flows, SFM considers the Riemannian structure induced by the Fisher information metric for categorical distributions. The **naive Euclidean assumption may fail to capture the true geometry** of the statistical manifold. Furthermore, we demonstrated in Sec.3.3 that **following the vector field induced by the Fisher information metric coincides with the steepest direction of the natural gradient that decreases the KL divergence**, which may also contribute to a better performance from the optimization point of view. Another possible advantage comes from the curved geodesic in Fig.1. Unlike straight geodesic under the Euclidean setting, the Riemannian vector fields are curved towards the boundary (more parallel to the boundary), making it hard to *overshoot* across the boundary.\n\n\n## Q2 Comparison to DirichletFM\n\nWe have included comparisons to DirichletFM in Appendix A.4 to provide some insights into SFM's better performance. DirichletFM considered the specific probability path in which **the target distributions are always assumed to be Dirichlet**. Such an assumption is not applicable to more complex target distributions and fails the Swiss roll toy example in Fig.3. Additionally, its probability path does not reflect the shortest geodesic distance on the statistical manifold, which may lead to suboptimal flows and vector fields. In contrast, our SFM framework constructs flows **based on the geodesic on the statistical manifold**, where vector fields follow the steepest direction of decreasing KL divergence. Our method can applied to any source and target distribution, making it more flexible and efficient.\n\n\n## Q3 FisherFlow \n\nWe thank you for mentioning FisherFlow, which is a concurrent work closely related to ours. We note that the FisherFlow paper was uploaded on May 23 after the NeurIPS submission ddl. Therefore, **SFM and FisherFlow should be considered concurrent work**.\n\nBoth SFM and FisherFlow explored the Riemannian structure of the statistical manifold to establish flows and vector fields by constructing isomorphic mappings between the simplex and the sphere. We both demonstrated the theoretical connection to natural gradient and explored the minibatch-OT formulation during training. **We took a step further to derive the exact likelihood over the probability space and a tight ELBO for the likelihood of discrete data**. We carefully choose the prior, reference measure, and divergence to ensure comparability with other discrete diffusion/flow models.\n\nWhile FisherFlow predominantly focused on bioinformatic tasks of DNA design, **we conducted more comprehensive experiments across various domains with additional baselines** to explore the superior performance and versatility of our method. We believe both works are important explorations on the use of Riemannian geometry structure on discrete generative tasks.\n\n\n## Q4 Conditional on measure\n\nWe apologize for the typo. We meaned \"condition on $x_1$\" where the conditional probability path has $p_1(x|x_1)\\approx \\delta_{x_1}$ at $t=1$, as described in the flow matching paper. In our case, each target point $x_1$ is a one-hot distribution $\\mu_1$. \n\n\n## Q5 data dimension\n\nThe data dimension refers to the length of data. As an example, a DNA sequence with $D$ bases has data dimension $D$. We can represent this data as a matrix $X \\in [0,1]^{D \\times n}$, where $n=4$ reflects the four different categories (A, T, C, G). For simplicity, we assumed $D=1$ in the derivations, which is a standard practice. It is straightforward to extend this to multi-dimension by jointly modeling $D$ categorical probability measures.\n\n## Q6 Connection to natural gradient\n\nPlease see the common rebuttal.\n\n\n## Q7 Optimal transport\n\nWe consider minibatch-OT similar to [20,43,62] but on the statistical manifold. We match between a minibatch of the noises from the prior distribution and the samples from the target distribution with the smallest transport cost (defined line 189) based on the statistical distance defined in Eq.3. A thorough investigation of the theoretical benefit of minibatch OT can be found in the original paper. For Markov-chain based models like D3PM and MultiFlow, it is not possible to derive such a distance measure due to the discrete jump between Markov states. \n\n## Q8 Significance of NLL calculation\n\nGenerative models are trained to capture data distribution $p_\\theta\\approx p_\\text{data}$. Therefore, it is natural and crucial to calculate the likelihood for a given data sample $p_\\theta(x)$ as both **the evaluation metric for generative models and confidence quantification for data**. When evaluated on ground truth data, the NLL serves as a natural evaluation metric of how closely the generative model fits the data. Many common evaluation metrics including cross-entropy, perplexity, and bits-per-dimension are derived from likelihood estimation. With the ability to calculate likelihood, we can provide an intrinsic evaluation of flow models. \n\nAdditionally, NLL can be used to measure the confidence of given inputs and facilitate RLHF of flow models. Many policy-based RL requires explicit log-likelihood, and ELBOs that are loose (e.g. with impromptu noise schedule) could impact the effectiveness of RLHF. Our exact NLL formulation is more accurate and our ELBO definition over discrete data is arguably much tighter, potentially benefiting applications that explicitly rely on likelihood.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear all Reviewers,


We sincerely appreciate your reviews that help make our work more concrete and comprehensive. We thank you for recognizing our novel geometric perspective of discrete generation and our paper's clear presentation. Here we address some of the common questions.


## More details regarding connection to natural gradient

As per request by Reviewer xYE5, WPWv, and hmQs, we elaborate more on the connection to natural gradient in addition to Sec.3.3. We want to demonstrate **the direction of the Riemannian vector field induced by the Fisher information metric coincides with the steepest direction of minimizing local KL divergence** (the natural gradient). 

**From the optimization view point**, one objective for generative modeling of categorical distributions would be the KL divergence $D_\text{KL}(\tilde{\mu}\\|\mu)$ between the generated distribution $\tilde{\mu}$ and the target distribution $\mu$. The fact that the Fisher information metric is the Hessian of KL divergence allows us to locally expand the change of KL divergence in Eq.10 in terms of Fisher information metric as $D_\text{KL}(\mu(\theta)\\|\mu(\theta_1))\approx \frac{1}{2}\sum_{jk}\Delta\theta_j\Delta\theta_j g_{jk}=\frac{1}{2}\\|\Delta\theta\\|^2_g$ where $\Delta\theta=\theta-\theta_1$. The steepest direction $\Delta\theta$ decreasing the KL divergence is known as the *natural gradient* in the existing literature.

**From the geometric view point**, 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\theta\\|^2_g$, we are indeed following the steepest direction of natural gradient that minimizes the local KL divergence. This theoretical connection may contribute to the better performance of SFM.

## Additional results for Text8

We first noted that, though we cited MultiFlow, their Text8 experiment was not added until Jun 5 (>2 weeks after NeurIPS ddl). Nevertheless, as requested by Rewiewer SAto, we now add results with the GPT-J-6B model in addition to BPC, following the setting in MultiFlow. We emphasize that **pseudo-LL is only used in Text8 BPC** and all discrete NLL comparisons follow the similar ELBO derivation in DDSM. We only compared NLLs to diffusion/flow models with comparable ELBO formulations (see the response to SAto) and used consistent divergence (see the response to xYE5). We also provided **non-NLL metrics** (FID, Sei MSE) for MNIST/DNA data which are **non-toy high-dimensional real-world datasets**, where SFM outperformed all baselines with a margin. We sincerely hope that these contributions can still be valued.

As GPT-J-6B was not trained on Text8, its NLL does not necessarily reflect Text8 distribution fitness. From the MultiFlow results, such an NLL can be **made artificially low by duplicating high-frequency words e.g. numbers**, thus joint comparison of GPT NLL and token entropy was used, where an entropy closer to data is preferred. We show in the table below that SFM tends to produce samples with data entropy closer to data, and the GPT NLLs are still comparable to (slightly worse) MultiFlow/D3PM when entropy is close. As SFM directly samples categorical probabilities instead of making multiple discrete jumps with logits, it is incompatible with temperature-tuning. Thus in the table below, we compare with $T=1$ which matches the entropy of SFM. We also noted that GPT NLL can be **easily fooled by randomly generated strings** based on letter frequency. Additionally, low-T MultiFlow variants achieved lower NLLs than the ground truth data, making this metric less credible. We also noticed that the **low NLL samples for MultiFlow often consist of repeated numbers with little semantic meaning**. To enable fair comparison, we provide multiple generated samples (above/below/just-on-average NLL for each model) and their GPT NLLs in the PDF. Different temperature settings for MultiFlow are also included using checkpoint provided on Github. We noted the big impact of temperature on MultiFlow as it generated more numbers repeatedly with lower temperature. In contrast, SFM achieved perceptually similar or even better results with more diverse vocabulary.

We thank the authors of MultiFlow for generously providing their raw data in Fig.2. The complete figure is provided in the pdf.

| Model  | GPT-J-6B score ↓ | Token entropy ↑ |
| - | - | - |
| Data | 4.099 | 7.479 |
| Random | **5.827** | 5.519 |
| SFM w/ OT | 7.071 | 7.347 |
| SFM w/o OT | 7.154 | **7.388** |
| LinearFM | 7.490 | 7.118 |
| SEDD mask | 6.49 | 7.17 |
| SEDD absorb | 6.28 | 6.97 |
| MultiFlow | 6.728 | 7.387 |
| MultiFlow η = 0 | 6.729 | 7.325 |
| D3PM | 6.929 | 7.379 |
| Autoregressive | 6.296 | 7.165 |


We thank Reviewer SAto for proposing an alternative evaluation. We will include these results on Text8 for more comprehensive evaluation, and tone down on our original claims. Nevertheless, we'd like to note that the **general formulation of SFM over continuous statistical manifold has potential beyond NLP discrete sampling**, and could be of interest to Bayesian inference, distributional RL, and other tasks that involve sampling over distributions. 

NLL used to be one of the common metrics for generative models as it reflects closeness to the data distribution. It is until recently that this metric became underlooked due to arguably loose and diverse ELBO derivation in diffusion, and researchers resort to third-party models (InceptionV3, GPT, AF2, etc.) fitted on different data that may be subject to artifact and bias. We'd like to call for better community efforts on this issue. We provided exact NLL and a tight ELBO formulation with a goodwill that this could be a step towards better evaluation of flow models, and we have tried our best to make it as comparable as possible.

Meta-review: 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.
Decision: Accept (poster)