# Categorical Flow Matching on Statistical Manifolds

Chaoran Cheng
University of Illinois Urbana-Champaign
<EMAIL_ADDRESS>
&Jiahan Li
Peking University
<EMAIL_ADDRESS>
Jian Peng
University of Illinois Urbana-Champaign
<EMAIL_ADDRESS>
&Ge Liu
University of Illinois Urbana-Champaign
<EMAIL_ADDRESS>
###### 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.

##  1 Introduction

Recently, conditional flow matching (CFM) models [36] have achieved remarkable
success in various generative domains including image generation [36, 18, 28],
molecule [55, 54, 30] and protein design [62, 11], and sequence generation
[56, 37, 13]. While attempts to generalize CFM and diffusion models to
discrete categorical data have been made, they typically exert ad hoc
assumptions on the structure of the discrete distribution. One group of work
relies on stochastic jumps of Markov chains in either the discrete-time [6,
37, 1] or continuous-time setting [12, 51, 13] that discards the continuous
nature of the underlying categorical distributions. Other work directly works
with the probability simplex [7, 56] or the corresponding logit space [27, 38,
23] with potentially imperfect assumptions that fail to capitulate the
underlying true geometry of the statistical manifold. Furthermore, likelihood
is often approximated by variational bounds in previous discrete generative
models due to the lack of tractable exact likelihood.

We propose to incorporate the intrinsic geometry of the _statistical manifold_
by viewing categorical data as points on the statistical manifold of
categorical distributions. Inspired by the mathematical results from
information theory, we utilize the Fisher information metric [48] to naturally
equip such a manifold with a Riemannian structure and develop an efficient
generative training scheme for learning the vector fields without stability
issues. We summarize our contributions as the following:

(1) We propose _Statistical Flow Matching_ (SFM), a novel and mathematically
rigorous generative framework on the manifold of parameterized probability
measures. SFM does not pose any prior assumptions on the statistical manifold
but instead deduces its intrinsic geometry via mathematical tools. To tackle
the discrete generation problem, we instantiate SFM on the manifold of
categorical distributions. We deduce closed-form exponential and logarithm
maps and develop an efficient flow-matching training algorithm that avoids
numerical issues by leveraging a diffeomorphism between manifolds. SFM
effectively leverages the intrinsic geometric properties by following the
shortest paths of geodesics between the noise and target distributions on the
statistical manifold.

(2) Our distinctive geometric perspective of the statistical manifold allows
us to further apply optimal transport during training and derive tractable
exact likelihood for any given sample of probability measure, both of which
are unachievable for most existing methods. We also introduce new theoretical
insights by establishing connections among Riemannian flow matching,
information geometry, and natural gradient descent, which allows us to
interpret SFM as following the steepest descent of the _natural gradient_ from
the optimization angle.

(3) We demonstrated with a toy example on simplex that SFM can learn more
complex patterns on the statistical manifold where existing models often fail
due to impromptu prior assumptions. We further conducted extensive experiments
on diverse real-world discrete generation tasks involving computer vision,
natural language processing, and bioinformatics. SFM consistently outperformed
existing diffusion or flow-based models and also achieved comparable results
with autoregressive models on character-level generation.

##  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 structure known as _statistical
manifold_. Mathematically, consider probability densities
$p=\frac{d\mu}{d\nu}:\mathcal{X}\to\mathbb{R}$ defined by the Radon-Nikodym
derivative where $\mu$ is a probability measure on the sample space
$\mathcal{X}$ and $\nu$ is the reference measure on $\mathcal{X}$. Suppose the
statistical manifold
$\mathcal{P}=\mathcal{P}(\mathcal{X})=\\{p:\int_{\mathcal{X}}d\mu=\int_{\mathcal{X}}p(x;\theta)d\nu=1\\}$
is parameterized by
$\theta=(\theta_{1},\theta_{2},\dots,\theta_{n})\in\Theta$, this
parameterization naturally provides a coordinate system for $\mathcal{P}$ on
which each point is a probability measure $\mu$ with the corresponding
probability density function $p(x;\theta)$. The Fisher information metric is
defined as:

$g_{jk}(\theta)=\mathbb{E}_{X}\left[\frac{\partial\log
p(X;\theta)}{\partial\theta_{j}}\frac{\partial\log
p(X;\theta)}{\partial\theta_{k}}\right]=\int_{\mathcal{X}}\frac{\partial\log
p(x;\theta)}{\partial\theta_{j}}\frac{\partial\log
p(x;\theta)}{\partial\theta_{k}}p(x;\theta)d\nu$ (1)

Rao demonstrated in his seminal paper [48] that statistical manifold can be
equipped with the Fisher information metric to obtain a Riemannian structure,
the study of which is known as _information geometry_ [5, 4, 8]. This
geometric view of statistical manifolds allows us to derive key geometric
concepts for our statistical flow matching framework. For example, a
_geodesic_ $\gamma(t):[0,1]\to\mathcal{P}$ defines a “shortest” path (under
the Riemannian metric) connecting two probability measures on the statistical
manifold. The _geodesic distance_ between two probability measures, also known
as the Fisher-Rao distance [48], measures the similarity between them. The
_tangent space_ $T_{\mu}(\mathcal{P})$ at a point $\mu\in\mathcal{P}$ can be
naturally identified with the affine subspace
$T_{\mu}(\mathcal{P})=\\{\nu|\int_{\mathcal{X}}d\nu=0\\}$ where each element
$\nu$ is a signed measure over $\mathcal{X}$. The _exponential map_
$\exp_{\mu}:T_{\mu}(\mathcal{P})\to\mathcal{P}$ and _logarithm map_
$\log_{\mu}:\mathcal{P}\to T_{\mu}(\mathcal{P})$ can also be defined on the
statistical manifold. While the geodesic for a parameterized family of
probability measures can be obtained numerically by solving the geodesic
equation when closed-form expression is unknown (see Appendix A.1), it usually
requires expensive simulations. Fortunately, closed-form geodesic distances
are available for many common distributions including categorical,
multinomial, and normal distributions [42], which motivate our method.

###  2.2 Conditional Flow Matching on Riemannian Manifold

The conditional flow matching (CFM) framework [36] provides a simple yet
powerful approach to generative modeling by learning a time-dependent vector
field that pushes the prior noise distribution to any target data
distribution. Such a flow-based model can be viewed as the continuous
generalization of the score matching (diffusion) model [52, 53, 25] while
allowing for a more flexible design of the denoising process. The Riemannian
flow matching framework [15] further extends CFM to general manifolds on which
a well-defined distance metric can be efficiently computed.

Consider a smooth Riemannian manifold $\mathcal{M}$ with the Riemannian metric
$g$, a _probability path_ $p_{t}:[0,1]\to\mathcal{P(M)}$ is a curve of
probability densities over $\mathcal{M}$. A _flow_
$\psi_{t}:[0,1]\times\mathcal{M}\to\mathcal{M}$ is a time-dependent
diffeomorphism defined by a time-dependent vector field
$u_{t}:[0,1]\times\mathcal{M}\to T\mathcal{M}$ via the ordinary differential
equation (ODE): $\frac{d}{dt}\psi_{t}(x)=u_{t}(\psi_{t}(x))$. The flow
matching objective directly regresses the vector field $u_{t}$ with a time-
dependent neural net $v(x_{t},t)$ where $x_{t}:=\psi_{t}(x)$. However, this
objective is generally intractable. Both [36, 15] demonstrated that a
tractable objective can be derived by conditioning on the delta measure
$\delta_{x_{1}}$ at $t=1$ of the probability path. The Riemannian flow
matching objective can be formulated as [15]

$\mathcal{L}=\mathbb{E}_{t\sim U[0,1],x_{0}\sim p_{0}(x),x_{1}\sim
q(x)}[\|v(x_{t},t)-u_{t}(x_{t}|x_{0},x_{1})\|_{g}^{2}]$ (2)

where $q$ is the data distribution, $p_{0}$ is the prior distribution, and
$x_{t}:=\psi_{t}(x|x_{0},x_{1})$ denotes the conditional flow. [15] further
demonstrated that if the exponential map and logarithm map can be evaluated in
closed-form, the condition flow can be defined as
$x_{t}=\exp_{x_{0}}(t\log_{x_{0}}x_{1})$, and the corresponding vector field
can be calculated as
$u_{t}(x_{t}|x_{0},x_{1})=\frac{d}{dt}x_{t}=\log_{x_{t}}(x_{1})/(1-t)$. We
also adapt this formulation for our statistical flow, which we will elaborate
on in the next section. We note that, since we are working with manifolds of
probability measures $\mathcal{M}=\mathcal{P(X)}$, we will use $p,q$ for
probability densities _over_ the manifold $\mathcal{P(X)}$ and $\mu,\nu$ for
probability measures (or probability masses for discrete distributions) _on_
the manifold $\mathcal{P(X)}$ to avoid confusion.

##  3 Method

Different from previous work that treated each distribution separately, we
adopt an integral viewpoint toward the manifold of probability distributions.
In this section, we focus on the statistical manifold of categorical
distributions to demonstrate the application of our method on discrete
generation tasks. However, we emphasize that our proposed SFM is applicable to
any statistical manifold with a closed-form geodesic distance and can be
broadly used in generative tasks targeting probability measures on the
statistical manifold.

###  3.1 Statistical Manifold of Categorical Distributions

Consider the discrete sample space $\mathcal{X}=\\{1,2,\dots,n\\}$, an
$n$-class categorical distribution over $\mathcal{X}$ can be parameterized by
$n$ parameters $\mu_{1},\mu_{2},\dots,\mu_{n}$ such that
$\sum_{i=1}^{n}\mu_{i}=1,\mu_{i}\geq 0$. In this way, the reference measure
$\nu$ is the counting measure and the probability measure $\mu$ can be written
as the convex combination of the canonical basis of delta measures
$\\{\delta^{i}\\}_{i=1}^{n}$ over $\mathcal{X}$:
$\mu=\sum_{i=1}^{n}\mu_{i}\delta^{i}$. Geometrically, this manifold can be
visualized as the $(n-1)$-dimensional simplex $\Delta^{n-1}$. The geodesic
distance between two categorical distributions is given in [48, 42] as

$d_{\text{cat}}(\mu,\nu)=2\arccos\left(\sum_{i=1}^{n}\sqrt{\mu_{i}\nu_{i}}\right)$
(3)

The target space at a point $\mu$ can be identified with
$T_{\mu}(\mathcal{P})=\\{u|\sum_{i=1}^{n}u_{i}=0\\}$ and the corresponding
inner product at point $\mu$ is defined as

$\langle
u,v\rangle_{\mu}=\sum_{i=1}^{n}\frac{u_{i}v_{i}}{\mu_{i}},\quad\mu\in\mathcal{P}^{\circ},u,v\in
T_{\mu}(\mathcal{P})$ (4)

where $\mathcal{P}^{\circ}$ denotes the interior of the statistical manifold
$\mathcal{P}$. Note that the inner product is ill-defined on the boundary,
causing numerical issues near the boundary. To circumvent this issue, we
introduce the following diffeomorphism:

$\pi:\mathcal{P}\to S^{n-1}_{+},\quad\mu_{i}\mapsto x_{i}=\sqrt{\mu_{i}}$ (5)

which maps the original statistical manifold to the positive orthant of a unit
$(n-1)$-sphere $S^{n-1}_{+}=\\{x|\sum_{i=1}^{n}x_{i}^{2}=1,x_{i}\geq 0\\}$
(see Fig.2). Note that we have
$\|\pi(\mu)\|_{2}^{2}=\sum_{i=1}^{n}\sqrt{\mu_{i}}^{2}=\sum_{i=1}^{n}\mu_{i}=1$.
The geodesic on $S^{n-1}_{+}$ follows the great circle, and the following
proposition holds between the geodesic distance on $S^{n-1}_{+}$ and
$\mathcal{P}$:

######  Proposition 1.

$d_{S}(\pi(\mu),\pi(\nu))=\frac{1}{2}d_{\mathrm{cat}}(\mu,\nu),\quad\mu,\nu\in\mathcal{P}$
(6)

A proof is provided in Appendix A.3). This indicates that we can work with the
geodesic distance on the unit sphere instead with up to a constant factor:

$d_{S}(x,y)=\arccos(\langle x,y\rangle),\quad x,y\in S^{n-1}_{+}$ (7)

The geodesic distance $d_{S}$ and the inner product
$\langle\cdot,\cdot\rangle$ are well-defined for the boundary, and we found
this transform led to the practical stabilized training of the flow model.
Visualizations of the Riemannian geometry on the statistical manifold of
3-class categorical distributions and the corresponding Euclidean geometry on
the simplex are provided in Fig.1 for comparison. The straight lines under the
Euclidean assumptions fail to capture the true curved geometry of the
statistical manifold.

Figure 1: 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 $\mu_{0}=(1/3,1/3,1/3)$. Middle:
Exponential maps (geodesics) from $\mu_{0}$ to different points near the
boundary. Right: Logarithm maps (vector fields) to $\mu_{0}$.

###  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 it
is possible to directly learn the vector field following the loss in Eq.2,
such a direct parameterization has numerical issues near the boundary. As
described in Sec.3.1, we apply the diffeomorphism $\pi$ in Eq.5 to derive a
more stable training objective on the spherical manifold:

$\mathcal{L}_{\text{SFM}}=\mathbb{E}_{t\sim
U[0,1],x_{0}\sim\pi(p_{0}(\mu)),x_{1}\sim\pi(q(\mu))}\left[\|v(x_{t},t)-u^{S}_{t}(x_{t}|x_{0},x_{1})\|_{2}^{2}\right]$
(8)

where $p_{0}$ is the prior noise distribution over $\mathcal{P}$ and $q$ is
the data distribution. $v:S^{n-1}_{+}\times[0,1]\to TS^{n-1}_{+}$ is a
learnable time-dependent vector field network that maps the interpolant
$x_{t}$ on the unit sphere to a tangent vector in the tangent bundle
$TS^{n-1}_{+}$. The ground truth conditional vector field $u^{S}$ is
calculated using the exponential and logarithms maps on the sphere (Appendix
A.2). The overall training and inference scheme is visualized in Fig.2 and
described in Appendix Alg. 2 \- 3.

We further implement a Euclidean flow matching model on the probability
simplex with linear interpolation between the noises and the target
distributions. Though linear flow matching offers “straight” lines under the
Euclidean assumption, it is unaware of the intrinsic geometric properties of
the statistical manifold and turns out to trace longer paths in terms of the
Riemannian metric. The objective for linear flow matching can be described as

$\mathcal{L}_{\text{LinearFM}}=\mathbb{E}_{t\sim U[0,1],\mu_{0}\sim
p_{0}(\mu),\mu_{1}\sim
q(\mu)}\left[\|v(\mu_{t},t)-(\mu_{1}-\mu_{0})\|_{2}^{2}\right]$ (9)

In the discussion above, we assume a single data dimension on the statistical
manifold. This can be extended to any data dimension by treating them as
independent channels of the input. In practice, each probability measure on
the simplex is represented by a matrix $X\in[0,1]^{D\times n}$ where $D$ is
the data dimension and each row sums to 1. The priors are independently
sampled from the uniform distribution over the $(n-1)$-simplex and each data
dimension is independently interpolated with the same timestep $t\sim U[0,1]$.
The flow model, on the other hand, takes all the data dimensions as input to
capture the dependence between different dimensions. During sampling, existing
ODE solvers and the simple Euler method are used to integrate the vector field
through time to obtain the final categorical distributions. Discrete samples
are then drawn from the generated categorical distributions for evaluation.
Details regarding model parameterization and sampling are further described in
Appendix C.1 and C.2.

Figure 2: Statistical flow matching (SFM) framework. (a) During training
(Sec.3.2), probability measures on $\mathcal{P}$ are mapped to
$\smash{S^{n-1}_{+}}$ via diffeomorphism $\pi$ to compute the time-dependent
vector field (in red). During inference, the learned vector field generates
the trajectory on $\smash{S^{n-1}_{+}}$ and we map the outcome of ODE back to
$\mathcal{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 delta measure to avoid undefined behaviors at 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 an optimization problem. For a generative
model on a statistical manifold, we want to minimize the discrepancy between
the target distributions and the generated distributions. Naturally, the
Kullback-Leibler divergence $D_{\text{KL}}(\mu(\theta)\|\mu(\theta_{1}))$ can
be used as a measure of statistical distance. We note that the Fisher
information metric can be obtained as the Hessian of the KL divergence with
respect to the parameter $\theta$. This can be demonstrated by the fact that
the KL divergence reaches the global minimum of 0 at $\theta=\theta_{1}$, so
all first-order partial derivatives are zero and the Hessian is positive semi-
definite. Therefore, Taylor expansion of KL divergence at $\theta_{1}$ gives

$D_{\text{KL}}(\mu(\theta)\|\mu(\theta_{1}))\approx\frac{1}{2}\sum_{jk}\Delta\theta_{j}\Delta\theta_{k}\left.\frac{\partial^{2}}{\partial\theta_{j}\partial\theta_{k}}D_{\text{KL}}(\mu(\theta)\|\mu(\theta_{1}))\right|_{\theta=\theta_{1}}=\frac{1}{2}\sum_{jk}\Delta\theta_{j}\Delta\theta_{k}g_{jk}(\theta_{1})$
(10)

Therefore, the Fisher information metric defines a second-order optimization
scheme for the KL divergence by following the “steepest” direction of the
Hessian. In fact, this idea has already been explored by previous work on
optimization known as _natural gradient descent_ [2, 3] and has been applied
to training machine learning models [46, 40, 44]. Instead of optimizing along
the normal gradient, the natural gradient is defined as
$\tilde{\nabla}\mathcal{L}=F^{-1}\nabla\mathcal{L}$ where $F$ is the Fisher
information matrix estimated from a batch of sampled data. Results established
in [2, 40] demonstrated that stochastic natural gradient descent is
asymptotically “Fisher efficient” and is indeed the steepest descent direction
in the manifold of distributions where distance is measured in small local
neighborhoods by the KL divergence. Following the geodesic defined by the
Fisher information metric, our SFM framework also shares these benefits with
an additional advantage of analytical expressions for geodesics, as we focus
on the family of categorical distributions instead of general statistical
models.

###  3.4 Optimal Transport on Statistical Manifold

The geometric viewpoint of the statistical manifold offers a continuous and
differentiable formulation of generative modeling and also provides a robust
way to measure the distance between two categorical distributions via the
well-defined geodesic distance in Eq.(3). In contrast, the Markov chain-based
methods cannot establish a robust distance measure due to the stochastic jumps
between discrete states. Inspired by the optimal transport formulation in
previous work [43, 20, 62], we naturally extend it to our statistical setting
in which the cost is defined by averaging the statistical distances over data
dimensions as
$d_{\text{cat}}(X,Y)=\frac{1}{D}\sum_{k=1}^{D}d_{\text{cat}}(\mu^{(k)},\nu^{(k)})$
for $X=\\{\mu^{(k)}\\}_{k=1}^{D},Y=\\{\nu^{(k)}\\}_{k=1}^{D}$. 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.

###  3.5 Exact Likelihood Calculation

Unlike diffusion-based models which rely on variational lower bounds for
likelihood estimation, our proposed method shares the continuous normalizing
flow’s capability of exact likelihood calculation. For an arbitrary test
sample $x\in\mathcal{M}$, using the change of measure formula, the likelihood
can be modeled by the continuity equation [15, 41]:

$\frac{\partial}{\partial t}\log
p_{t}(x_{t})+\mathrm{div}_{g}(v_{t})(x_{t})=0$ (11)

where $\mathrm{div}_{g}$ is the Riemannian divergence and
$v_{t}(x_{t}):=v(x_{t},t)$ is the time-dependent vector field. In this way,
the push-forward probability measures $p_{t}(x_{t})$ defined via the learned
flow can be obtained as the integral of the Riemannian divergence back through
time on the simulated trajectory $x_{t}$. Following [10, 7], we define the ODE
log-likelihood as:

$\log p^{\text{ODE}}(x_{1})=\int_{1}^{0}\mathrm{div}_{g}(v_{t})(x_{t})dt$ (12)

where the trajectory $x_{t}$ is obtained by solving the differential equation
$\frac{\partial}{\partial t}x_{t}=v_{t}(x_{t})$ reverse through time with the
initial condition $x_{1}$ at $t=1$. In this way, we have $\log p(x_{1})=\log
p^{\text{ODE}}+\log p_{0}(x_{0})$ where $\log p_{0}(x_{0})$ is the log-
likehood of the prior distribution at data point $x_{0}$. In practice, we
follow previous work [7] to use Hutchinson’s trace estimator [29] to
efficiently obtain an unbiased estimation of the divergence using standard
normal random variables. To further account for the transform $\pi$ from
$\mathcal{P}$ to $S^{n-1}_{+}$ and the reverse transform $\pi^{-1}$,
additional change of measure formulae need to be applied. Consider the push-
forward of the probability measure $p$ over $\mathcal{P}$ under the
diffeomorphism $\pi$ defined by $\pi_{*}p(V):=p(\pi^{-1}V)$, we have the
change of measure identity $\pi_{*}p(\pi(\mu))|\det d\pi(\mu)|=p(x)$ for
$x=\pi(\mu)$. Therefore, adding the two log-determinant of the push-forward
measures, the log-likelihood can be formulated as

$\log p_{1}(\mu_{1})=\log|\det d\pi^{-1}(x_{1})|+\log
p^{\text{ODE}}(x_{1})+\log|\det d\pi(\mu_{0})|+\log p_{0}(\mu_{0})$ (13)

The above formula is well-defined for all interior points of $\mathcal{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 [58].
Following [7], we derive a variational lower bound for the likelihood as the
marginalized probability over the small neighborhood of the delta measure
$\mu_{1}=\delta$ at the boundary:

$\log p(\delta)\geq\mathbb{E}_{q(\tilde{\mu}_{1}|\delta)}[-\log
q(\tilde{\mu}_{1}|\delta)+\log p(\delta|\tilde{\mu}_{1})+\log
p_{1}(\tilde{\mu}_{1})]$ (14)

where $q(\tilde{\mu}_{1}|\delta)$ can be viewed as the forward noising
probability at $\tilde{\mu}_{1}$ which is close to $\delta$ with a closed-form
likelihood. $p(\delta|\tilde{\mu}_{1})$ is the categorical likelihood (cross-
entropy) and $p_{1}(\tilde{\mu}_{1})$ is the model likelihood in Eq.(13). The
overall workflow of calculating NLL is demonstrated in Fig.2, and more details
regarding likelihood computation can be found in Appendix B.1. It is worth
mentioning that the continuous likelihood calculated here is defined with
respect to the probability distribution over the statistical manifold
$\mathcal{P}$. In contrast, the bits-per-dimension score for autoregressive
models is usually defined with respect to a specific categorical distribution
on $\mathcal{P}$ and therefore not comparable, as was also noted in [7].

##  4 Experiments

Our proposed SFM framework can be leveraged to approximate arbitrary
distribution over $\mathcal{P}$, i.e., any distribution over the parameterized
family of probability measures. We first demonstrate this with the toy example
of a Swiss roll distribution on the probability simplex. We then conduct
extensive experiments on real-world datasets for discrete generation across
various domains including computer vision, natural language processing, and
bioinformatics to demonstrate the effectiveness of our proposed model. We
train and evaluate two different versions of the SFM framework: SFM w/ OT for
our proposed model with optimal transport during training and SFM w/o OT for
the model without optimal transport. We also implement a linear flow matching
model on the probability simplex using the loss in Eq.(9) as an additional
baseline, for which we dub LinearFM. For discrete data, we always use Eq.(14)
to obtain finite negative log-likelihood (NLL). More details regarding model
architectures and the experimental setup can be found in Appendix C.

###  4.1 Toy Example: Swiss Roll on Simplex

We project the Swiss roll dataset onto the 2-simplex with some additional
margins to make sure no point resides on the boundary. We used a time-
dependent MLP to model the vector field for all models. The generative points
on the simplex and NLL are calculated using the Dopri5 ODE solver [9] and are
shown in Fig.3.

Figure 3: Generated samples of the Swill roll on simplex dataset and NLL
(lower is better).

We note that most existing models assume the target data to be delta measures
(one-hot distributions) and cannot be applied to this task. Both DDSM [7] and
DirichletFM [56] imposed strong prior assumptions on the data distribution as
Dirichlet distributions, which failed to capture the complex geometric shape
of the Swill roll data points as demonstrated in the generated samples. On the
contrary, directly built upon the geometry of the statistical manifold, our
SFM model successfully captured the detailed data geometry. As no data resides
on the boundary, exact NLL can be obtained using Eq.13 and was averaged over
all points in the dataset. Though linear flow matching could also capture the
geometry of the data, our SFM outperformed it in terms of NLL.

###  4.2 Binarized MNIST

We extend our SFM to model discrete generation tasks in computer vision. The
binarized MNIST dataset [49] is the binarized version of the original MNIST
dataset [33] by thresholding the original continuous value to be either 0 or
1, thus can be viewed as a 2-class generation task with a data dimension of
$28^{2}=784$. We also compared D3PM [6] with an additional absorbing state. We
used a convolutional neural net adopted from [53] with additional additive
Fourier-based time embeddings. The quantitative evaluation of NLL and Fréchet
inception distance (FID) are shown in Tab.1.

Table 1: NLL and FID of different discrete models on binarized MNIST. * is
from [7].

| SFM w/ OT | SFM w/o OT | LinearFM | D3PM | DDSM | DirichletFM
---|---|---|---|---|---|---
NLL$\downarrow$ | -1.687 $\pm$ 0.020 | -1.631 $\pm$ 0.027 | 0.622 $\pm$ 0.022 | 0.141 $\pm$ 0.021 | 0.100 $\pm$ 0.001* | NA
FID$\downarrow$ | 4.62 | 5.15 | 5.91 | 10.35 | 7.79 | 40.60

All the generated results and NLL calculations were based on the Dopri5 ODE
solver. Additional ablation results can be found in Appendix D.1 and generated
images can be found in Appendix D.2. The NLL for diffusion-based models (D3PM
and DDSM) was also calculated based on the variational bounds derived in their
papers. Our proposed model consistently outperformed both the linear flow
matching and other diffusion baselines in terms of FID and NLL, achieving
higher sample quality and likelihood.

###  4.3 Text8

The Text8 dataset [39] is a media-size character-level corpus consisting of a
small vocabulary of 27, which includes 26 lowercase letters and the whitespace
token. We followed the convention in previous work [6, 23] to use random
chunks of length 256 in both training and evaluation without any
preprocessing. We used a 12-layer transformer-based predictor modified from
[47] with additional time embeddings. As our exact likelihood is not directly
comparable to bits-per-character (BPC) reported in previous work, we
additionally calculated the pseudo-BPC inspired by [50]. See Appendix B.2 for
additional details. All generated samples and NLL were estimated using 256
Euler steps. Quantitative results of BPC and NLL are provided in Tab.3 and
generated texts are provided in Appendix D.2. The results for the transformer-
based autoregressive language models are also provided as a reference. Our
proposed SFM achieved the second-best performance among all diffusion or flow
baselines. We also noted a significant performance gap between SFM and the
naive linear flow matching on simplex, which highlighted the importance of
capturing the intrinsic geometric properties of the underlying statistical
manifold. Optimal transport does not significantly affect the final
performance here, which we presume might be due to the long training stage in
which vector fields were well-explored.

Table 2: BPC and NLL of different models on Text8. Results marked * are taken from corresponding papers. Model | BPC$\downarrow$ / (NLL$\downarrow$)
---|---
SFM w/ OT | 1.291 (-0.776 $\pm$ 0.114)
SFM w/o OT | 1.298 (-0.856 $\pm$ 0.132)
LinearFM | 1.935 (25.834 $\pm$ 0.186)
D3PM-absorb | 1.47*
BFN[23] | 1.41*
SEDD-absorb[37] | 1.32*
Discrete Flow[59] | 1.23*
Argmax Flow[27] | 1.80*
Transformer[6] | 1.23*
Transformer XL[17] | 1.08*

Table 3: SP-MSE on the generated promoter DNA sequences. * are from [7] and † are from [56]. Model | SP-MSE$\downarrow$
---|---
SFM w/ OT | 0.0279
SFM w/o OT | 0.0258
LinearFM | 0.0282
DDSM | 0.0334*
D3PM-uniform | 0.0375*
Bit-Diffusion (one-hot) [16] | 0.0395*
Bit-Diffusion (bit) [16] | 0.0414*
Language Model | 0.0333†
DirichletFM | 0.0269†

###  4.4 Promoter Design

We further apply SFM to the practical task of promoter DNA sequence design in
the bioinformatics realm. The promoter is a key element in gene transcription
and regulation, and the generation of desired promoters can better help us
understand the intricate interactions between human genes. [7] proposed a
human promoter sequence dataset containing 100k promoter sequences with the
corresponding transcription initiation signal profiles. Each promoter sequence
is 1024 base pairs long and is centered at the annotated transcription start
site position [26]. Therefore, we can model promoter design as a generation
task with 4 categories (the base pair ATGC) conditioned on the given
transcription signals. We follow [7] to concatenate the signal as additional
input to the vector field predictor built upon 20 stacks of 1-dimensional
convolutional layers on the input sequence. Optimal transport can also be
applied for conditional generation, as we can pair the conditions with the
input to make sure that the conditional arguments align with the target one-
hot distributions.

To provide a quantitative evaluation of the generated promoter sequences, we
follow [7] to apply the pre-trained deep learning sequence model Sei [14] to
predict active promoters based on predictions of the chromatin mark H3K4me3.
As the dataset spans the whole range of human promoter activity levels from
highly expressed promoters to those with very low expression, the evaluation
is based on comparing the scores between the generated sequences and the
ground truth human genome promoter sequences on the test chromosomes. The
metric SP-MSE is the mean squared error between the predicted promoter
activity of the generated sequences and human genome sequences (lower is
better) and is demonstrated in Tab.3. Our SFM was able to achieve the lowest
SP-MSE score compared with other baselines. It is worth noting, though, that
optimal transport produced slightly sub-optimal results. We hypothesize that
this is because, in conditional generative tasks, the final generation should
rely heavily on the conditions. Simply matching inputs with the targets using
distance measures may discard important information in the conditional
arguments and may not be the best practice.

##  5 Related Work

As we put a special interest in discrete generation, we start with the
existing work on discrete generation. We first list a few traditional
autoregressive models [17, 47, 59] and masked language models [19, 50] but
will skip them as we mainly focus on the diffusion and flow matching models.
Existing diffusion and flow-based discrete generative models can be
categorized into two main groups. The first group relies on stochastic jumps
of Markov chains by treating discrete classes as Markov states. By choosing a
proper transition matrix, the target one-hot distribution can be gradually
noised into the stationary distribution. Noticeably, this approach can also
adopt an additional absorbing state to mimic the mask token in masked language
modeling [6]. D3PM [6] adapted the discrete-time Markov chain with the
diffusion framework and SEDD [37] proposed a more efficient training scheme by
learning the score entropy between different states. [12, 51] extended it to
the continuous-time Markov chain by using the infinitesimal generator (rate
matrix) instead, and [13] further applied the flow matching framework.
Although the transition matrix or the rate matrix determines the entire
diffusion trajectory, there is no canonical way of choosing it for different
generation tasks to control the mixing rate. Also, due to the presence of
discrete jumps of the Markov chain, exact likelihood calculation is no longer
feasible. Only variational bounds were derived in [6, 13]. The other group
directly works with the probability simplex or the logit space with certain
assumptions of the underlying geometric structure. As an example, our linear
flow matching assumes a Euclidean geometry on the probability simplex and is
often used as a baseline in previous work [56]. The multinomial diffusion [27]
assumed a Euclidean geometry on the logits space so interpolation became
multiplication back on the probability simplex. Dirichlet diffusion (DDSM) [7]
and Dirichlet flow matching (DirichletFM) [56] exerted a strong prior
assumption on the probability path as the Jacobi diffusion process. We provide
a more detailed analysis of these models in Appendix A.4. For models that
directly operate on the logit space, the Bayesian flow network (BFN) assumed
Gaussian distributions on the logit space with Bayesian update rules. [38]
also assumed a Euclidean geometry in the logit space with targets as soft one-
hot logits. The assumptions made in these models may not always hold, e.g.,
for the Swill 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 [48] discussing the Fisher information metric. Most
previous work utilizing information geometry focuses on optimization [2, 3]
with a few exceptions on representation learning. [60] utilized the geodesic
distance between two univariate Gaussians for function shape matching for
retrosynthesis of gold nanoparticles. [34] treated point cloud data as
probability measures over the 3D Euclidean space and considered the pullback
metric on the latent space to obtain optimal latent coordinates for the
autoencoder. [45] further applied such a method on single-cell RNA sequence
trajectories. These models demonstrated the advantage of following the proper
geometry compared with the vanilla Euclidean setting, but the pullback metric
needed to be evaluated with expensive Jacobian-vector products during
training. In contrast, our proposed SFM, for the first time, leverages
mathematical tools from information geometry to generative modeling. SFM
directly operates on the statistical manifold with closed-form geodesics,
providing a 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 on the statistical manifold of
categorical distributions in this work and derived optimal transport and exact
likelihood formulae. We applied SFM to diverse downstream discrete generation
tasks across different domains to demonstrate our framework’s effectiveness
over the baselines. We noted that SFM can be further extended to non-discrete
generative tasks whose targets are probability distributions, which we will
leave as future work.

We are also aware of the limitations of our SFM framework. As a special case
of the flow matching model, the generation is an iterative process of
refinement that cannot modify the size of the initial input. This may pose
limitations to generation compared with autoregressive models. Furthermore, we
adopted the assumption of independence between classes so that the canonical
Riemannian structure can be induced by the Fisher metric. However, discretized
data like CIFAR-10 [31] (256 ordinal pixel values) do not follow this
assumption, and results on such data might be suboptimal.

## 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–2048, 1998.
  * [4] S.-i. Amari and H. Nagaoka.  Methods of information geometry, volume 191.  American Mathematical Soc., 2000.
  * [5] C. Atkinson and A. F. Mitchell.  Rao’s distance measure.  Sankhyā: The Indian Journal of Statistics, Series A, pages 345–365, 1981.
  * [6] J. Austin, D. D. Johnson, J. Ho, D. Tarlow, and R. Van Den Berg.  Structured denoising diffusion models in discrete state-spaces.  Advances in Neural Information Processing Systems, 34:17981–17993, 2021.
  * [7] P. Avdeyev, C. Shi, Y. Tan, K. Dudnyk, and J. Zhou.  Dirichlet diffusion score model for biological sequence generation.  In International Conference on Machine Learning, pages 1276–1301. PMLR, 2023.
  * [8] N. Ay, J. Jost, H. Vân Lê, and L. Schwachhöfer.  Information geometry, volume 64.  Springer, 2017.
  * [9] R. Bank, R. Graham, J. Stoer, and R. Varga.  Springer series in 8.  1993\.
  * [10] H. Ben-Hamu, S. Cohen, J. Bose, B. Amos, A. Grover, M. Nickel, R. T. Chen, and Y. Lipman.  Matching normalizing flows and probability paths on manifolds.  arXiv preprint arXiv:2207.04711, 2022.
  * [11] 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.02391, 2023.
  * [12] A. Campbell, J. 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–28279, 2022.
  * [13] A. Campbell, J. 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 arXiv:2402.04997, 2024.
  * [14] K. M. Chen, A. K. Wong, O. G. Troyanskaya, and J. Zhou.  A sequence-based global map of regulatory activity for deciphering human genetics.  Nature genetics, 54(7):940–949, 2022.
  * [15] R. T. Chen and Y. Lipman.  Riemannian flow matching on general geometries.  arXiv preprint arXiv:2302.03660, 2023.
  * [16] T. Chen, R. Zhang, and G. Hinton.  Analog bits: Generating discrete data using diffusion models with self-conditioning.  arXiv preprint arXiv:2208.04202, 2022.
  * [17] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov.  Transformer-xl: Attentive language models beyond a fixed-length context.  arXiv preprint arXiv:1901.02860, 2019.
  * [18] Q. Dao, H. Phung, B. Nguyen, and A. Tran.  Flow matching in latent space.  arXiv preprint arXiv:2307.08698, 2023.
  * [19] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [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] A. Graves.  Generating sequences with recurrent neural networks.  arXiv preprint arXiv:1308.0850, 2013.
  * [23] A. Graves, R. K. Srivastava, T. Atkinson, and F. Gomez.  Bayesian flow networks.  arXiv preprint arXiv:2308.07037, 2023.
  * [24] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [25] J. Ho, A. Jain, and P. Abbeel.  Denoising diffusion probabilistic models.  Advances in neural information processing systems, 33:6840–6851, 2020.
  * [26] 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 rnas with accurate 5’ ends.  Nature, 543(7644):199–204, 2017.
  * [27] E. Hoogeboom, D. Nielsen, P. Jaini, P. Forré, and M. Welling.  Argmax flows and multinomial diffusion: Learning categorical distributions.  Advances in Neural Information Processing Systems, 34:12454–12465, 2021.
  * [28] 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 Intelligence, volume 38, pages 2247–2255, 2024.
  * [29] M. F. Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  Communications in Statistics-Simulation and Computation, 18(3):1059–1076, 1989.
  * [30] L. Klein, A. Krämer, and F. Noé.  Equivariant flow matching.  Advances in Neural Information Processing Systems, 36, 2024.
  * [31] A. Krizhevsky, G. Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [32] A. Le Brigant, S. C. Preston, and S. Puechmorel.  Fisher-rao geometry of dirichlet distributions.  Differential Geometry and its Applications, 74:101702, 2021.
  * [33] Y. LeCun, C. Cortes, and C. Burges.  Mnist handwritten digit database.  ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2, 2010.
  * [34] 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.
  * [35] 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.
  * [36] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le.  Flow matching for generative modeling.  arXiv preprint arXiv:2210.02747, 2022.
  * [37] 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.
  * [38] R. K. Mahabadi, J. Tae, H. Ivison, J. Henderson, I. Beltagy, M. E. Peters, and A. Cohan.  Tess: Text-to-text self-conditioned simplex diffusion.  arXiv preprint arXiv:2305.08379, 2023.
  * [39] M. Mahoney.  Large text compression benchmark, 2011.
  * [40] J. Martens.  New insights and perspectives on the natural gradient method.  Journal of Machine Learning Research, 21(146):1–76, 2020.
  * [41] E. Mathieu and M. Nickel.  Riemannian continuous normalizing flows.  Advances in Neural Information Processing Systems, 33:2503–2515, 2020.
  * [42] H. K. Miyamoto, F. C. Meneghetti, and S. I. Costa.  On closed-form expressions for the fisher-rao distance.  arXiv preprint arXiv:2304.14885, 2023.
  * [43] K. Nguyen, D. Nguyen, T. Pham, N. Ho, et al.  Improving mini-batch optimal transport via partial transportation.  In International Conference on Machine Learning, pages 16656–16690. PMLR, 2022.
  * [44] 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–A1655, 2023.
  * [45] A. Palma, S. Rybakov, L. Hetzel, and F. J. Theis.  Modelling single-cell rna-seq trajectories on a flat statistical manifold.  In NeurIPS 2023 AI for Science Workshop, 2023.
  * [46] H. Park, S.-I. Amari, and K. Fukumizu.  Adaptive natural gradient learning algorithms for various stochastic models.  Neural Networks, 13(7):755–764, 2000.
  * [47] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al.  Language models are unsupervised multitask learners.  OpenAI blog, 1(8):9, 2019.
  * [48] 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.
  * [49] 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.
  * [50] J. Salazar, D. Liang, T. Q. Nguyen, and K. Kirchhoff.  Masked language model scoring.  arXiv preprint arXiv:1910.14659, 2019.
  * [51] J. E. Santos, Z. R. Fox, N. Lubbers, and Y. T. Lin.  Blackout diffusion: generative diffusion models in discrete-state spaces.  In International Conference on Machine Learning, pages 9034–9059. PMLR, 2023.
  * [52] Y. Song and S. Ermon.  Generative modeling by estimating gradients of the data distribution.  Advances in neural information processing systems, 32, 2019.
  * [53] Y. Song and S. Ermon.  Improved techniques for training score-based generative models.  Advances in neural information processing systems, 33:12438–12448, 2020.
  * [54] 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.
  * [55] H. Stark, B. Jing, R. Barzilay, and T. Jaakkola.  Harmonic prior self-conditioned flow matching for multi-ligand docking and binding site design.  In NeurIPS 2023 AI for Science Workshop, 2023.
  * [56] 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.
  * [57] C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna.  Rethinking the inception architecture for computer vision.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818–2826, 2016.
  * [58] L. Theis, A. v. d. Oord, and M. Bethge.  A note on the evaluation of generative models.  arXiv preprint arXiv:1511.01844, 2015.
  * [59] D. Tran, K. Vafa, K. Agrawal, L. Dinh, and B. Poole.  Discrete flows: Invertible generative models of discrete data.  Advances in Neural Information Processing Systems, 32, 2019.
  * [60] K. Vaddi, H. T. Chiang, and L. D. Pozzo.  Autonomous retrosynthesis of gold nanoparticles via spectral shape matching.  Digital Discovery, 1(4):502–510, 2022.
  * [61] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [62] 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 arXiv:2310.05297, 2023.

Supplementary Material

##  Appendix A Riemannian Manifold and Information Geometry

In this section, we give a more detailed mathematical background on Riemannian
manifold and information geometry related to this work. A more comprehensive
background on the Riemannian manifold can be found in [21]. For information
geometry, [8, 4] provide more comprehensive analyses and rigorous
formulations.

###  A.1 Riemannian Manifold

A Riemannian manifold $\mathcal{M}$ is a real, smooth manifold equipped with a
positive definite inner product $g$ on the tangent space $T_{x}(\mathcal{M})$
at each point $x\in\mathcal{M}$. Let
$T\mathcal{M}=\bigcup_{x\in\mathcal{M}}T_{x}(\mathcal{M})$ be the _tangent
bundle_ of the manifold $\mathcal{M}$, a time-dependent _vector field_ on
$\mathcal{M}$ is a mapping $u_{t}:[0,1]\times\mathcal{M}\to T\mathcal{M}$
where $u_{t}(x)\in T_{x}(\mathcal{M})$. A _geodesic_ is a locally distance-
minimizing curve on the manifold. The existence and the uniqueness of the
geodesic state that for any point $x\in\mathcal{M}$ and for any tangent vector
$u\in T_{x}(\mathcal{M})$, there exists a unique geodesic
$\gamma:[0,1]\to\mathcal{M}$ such that $\gamma(0)=x$ and
$\gamma^{\prime}(0)=u$. The _exponential map_ $\exp:\mathcal{M}\times
T\mathcal{M}\to\mathcal{M}$ is uniquely defined to be
$\exp_{x}(u):=\gamma(1)$. The _logarithm map_
$\log:\mathcal{M}\times\mathcal{M}\to T\mathcal{M}$ is defined as the inverse
mapping of the exponential map such that $\exp_{x}(\log_{x}(y))\equiv
y,\forall x,y\in\mathcal{M}$111Note that the formulation here is slightly
different from Sec.2.1 where we fix a point $\mu$ on $\mathcal{P}$.. With the
exponential map and logarithm map, the time-dependent flow can be compactly
written as time interpolation along the geodesic:

$x_{t}:=\psi_{t}(x_{t}|x_{0},x_{1})=\exp_{x_{0}}(t\log_{x_{0}}x_{1}),\quad
t\in[0,1]$ (15)

It can be demonstrated that the above flow indeed traces the geodesic between
$x_{0},x_{1}$ with linearly decreasing geodesic distance
$d_{g}(x_{t},x_{1})=(1-t)d_{g}(x_{0},x_{1})$ [15]. For an $n$-dimensional
Riemannian manifold, the geodesic can be obtained by solving the geodesic
equation

$\sum_{j=1}^{n}g_{ji}\frac{d^{2}x_{j}}{dt^{2}}+\sum_{j=1}^{n}\sum_{k=1}^{n}\Gamma_{ijk}\frac{dx_{j}}{dt}\frac{dx_{k}}{dt}=0,\quad
i=1,2,\dots,n$ (16)

where $x_{i}$ are local coordinates of the geodesic and $\Gamma_{ijk}$ are the
Christoffel symbols of the first kind defined by

$\Gamma_{ijk}=\frac{1}{2}\left(\frac{\partial g_{ji}}{\partial
x_{k}}+\frac{\partial g_{ji}}{\partial x_{j}}-\frac{\partial g_{jk}}{\partial
x_{i}}\right),\quad i,j,k=1,2,\dots,n$ (17)

The closed-form expressions for the geodesic are generally not available. [42]
curated a wide range of common statistical manifolds for parameterized
families of both the discrete and the continuous distributions. We will focus
on the statistical manifold of categorical distributions and also the
spherical manifold, as the latter is closely related to the former via the
diffeomorphism $\pi$ in Eq.(5).

###  A.2 Spherical Manifold

The positive orthant of the unit $(n-1)$-sphere $S^{n-1}_{+}$ is a
$(n-1)$-dimensional manifold. The sphere can be embedded into the ambient
Euclidean manifold $\mathbb{R}^{n}$ so it inherits the canonical inner product
as

$\langle u,v\rangle_{S}=\langle u,v\rangle=\sum_{i=1}^{n}u_{i}v_{i},\quad
u,v\in T_{x}(S^{n-1}_{+})$ (18)

The tangent space $T_{x}(S^{n-1}_{+})=\\{u|\langle u,x\rangle=0\\}$ is a
$(n-1)$-dimensional hyperplane perpendicular to the vector $x$. The geodesic
on the sphere follows the great circle between two points, and the geodesic
distance can be calculated in Eq.7. The corresponding exponential map can be
calculated as:

$\exp_{x}(u)=x\cos\|u\|_{2}+u\operatorname{sinc}\|u\|_{2}$ (19)

where $\operatorname{sinc}(\theta)=\sin\theta/\theta$ is the unnormalized sinc
function. The logarithm map can be calculated as:

$\log_{x}(y)=\arccos(\langle x,y\rangle)\frac{P_{x}(y-x)}{\|P_{x}(y-x)\|_{2}}$
(20)

where $P_{x}(w)=w-\langle x,w\rangle x$ is the projection of vector $w$ onto
the tangent space $T_{x}(S^{n-1}_{+})$.

###  A.3 Statistical Manifold of Categorical Distributions

The statistical manifold of categorical distributions
$\mathcal{P}(\mathcal{X})$ is closely related to $S^{n-1}_{+}$ by the
diffeomorphism $\pi$ in Eq.(5). In fact, $\pi$ is an isometry between the
Fisher information metric on $\mathcal{P}$ and the standard induced metric on
$S^{n-1}_{+}$ up to a constant scaling factor of $1/4$. To see this, we have
the following proposition:

######  Proposition 2.

$\left\langle\frac{\partial\pi(\mu)}{\partial
u},\frac{\partial\pi(\mu)}{\partial v}\right\rangle=\frac{1}{4}\langle
u,v\rangle_{\mu},\quad u,v\in T_{\mu}(\mathcal{P})$ (21)

###### Proof.

$\left\langle\frac{\partial\pi(\mu)}{\partial
u},\frac{\partial\pi(\mu)}{\partial
v}\right\rangle=\left\langle\left.\frac{d}{dt}\pi(\mu+tu)\right|_{t=0},\left.\frac{d}{dt}\pi(\mu+tv)\right|_{t=0}\right\rangle=\sum_{i=1}^{n}\frac{u_{i}}{2\sqrt{\mu_{i}}}\frac{v_{i}}{2\sqrt{\mu_{i}}}=\frac{1}{4}\langle
u,v\rangle_{\mu}$

∎

Note that the infinitesimal arc length on the geodesic can be expressed as
$ds^{2}=\|dx\|_{g}^{2}=\sum_{jk}dx_{j}g_{jk}(x)dx_{k}$ where $\|\cdot\|_{g}$
is the canonical Riemannian norm induced by the inner product. Integrating
over the geodesic, we can easily arrive at the result in Proposition 1. Based
on this, the exponential map and logarithm map on this statistical manifold
can be written as

$\displaystyle\exp_{\mu}(u)$
$\displaystyle=\left(\sqrt{\mu}\cos\left\|\frac{u}{2\sqrt{\mu}}\right\|_{2}+\frac{u}{2\sqrt{\mu}}\operatorname{sinc}\left\|\frac{u}{2\sqrt{\mu}}\right\|_{2}\right)^{2}$
(22) $\displaystyle\log_{\mu}(\nu)$
$\displaystyle=\frac{2\arccos(\langle\sqrt{\mu},\sqrt{\nu}\rangle)}{\sqrt{1-\langle\sqrt{\mu},\sqrt{\nu}\rangle}}(\sqrt{\mu\odot\nu}-\langle\sqrt{\mu},\sqrt{\nu}\rangle\mu)$
(23)

where $\odot$ denotes the pairwise multiplication. We mentioned in the main
text that directly using the exponential and logarithm maps on the statistical
manifold has numerical issues near the boundary. Instead, we used the mapping
$\pi$ to work with the spherical manifold. Nonetheless, Eq.(22) and (23)
provide useful tools for visualization of the statistical manifold, as
demonstrated in Fig.1.

We also noted the relation between the Fisher information metric defined in
Eq.1 and the canonical inner product defined in Eq.4 for general statistical
manifolds. Let $\mathcal{M}$ be an $n$-dimensional differentiable manifold and
consider an embedding defined by
$\mathcal{M}\hookrightarrow\mathcal{P(M)},\theta\mapsto
p(\theta)=\sum_{i\in\mathcal{X}}p_{i}(\theta)\delta^{i}$. The pullback of the
Fisher metric $g$ defines a metric on $\mathcal{M}$. For $u,v\in
T_{\theta}(\mathcal{M})$, we have

$\displaystyle g_{\theta}(u,v)$ $\displaystyle:=p^{*}(g)_{\theta}(u,v)$ (24)
$\displaystyle=g_{p(\theta)}\left(\frac{\partial p}{\partial u},\frac{\partial
p}{\partial v}\right)$
$\displaystyle=\sum_{i}\frac{1}{p_{i}(\theta)}\frac{\partial p_{i}}{\partial
u}(\theta)\frac{\partial p_{i}}{\partial v}(\theta)$
$\displaystyle=\sum_{i}p_{i}(\theta)\frac{\partial\log p_{i}}{\partial
u}(\theta)\frac{\partial\log p_{i}}{\partial v}(\theta)$

which gives the more common definition of the Fisher information matrix in
Eq.1. This relation also generally holds for a continuous sample space
$\mathcal{X}$ but requires more careful handling with measure theory tools. We
refer interested readers to mathematical references, e.g., Chapter 3.1 in [8].

###  A.4 Other Statistical Manifolds

The statistical manifold of categorical distributions provides an intuitive
way of modeling discrete targets. However, we noted that other parameterized
distributions can also be used to derive alternative statistical manifolds
with different geometries. We will briefly discuss some manifolds that have
been used in previous generative models (though not from the information
geometry perspective as we do).

#### Dirichlet Distribution

The Dirichlet distribution $\mathrm{Dir}(\alpha)$ is a multivariate
generalization of the beta distribution. It is defined on the
$(n-1)$-dimensional simplex $\Delta^{n-1}$ with $n$ concentration parameters
$\alpha=\\{\alpha_{i}\\}_{i=1}^{n}$. The probability density function is
defined as

$p(x;\alpha)=\frac{1}{\mathrm{B}(\alpha)}\prod_{i=1}^{n}x_{i}^{\alpha_{i}-1},\quad
x\in\Delta^{n-1}$ (25)

where
$\mathrm{B}(\alpha)=\prod_{i=1}^{n}\Gamma(\alpha_{i})/\Gamma(\sum_{i=1}^{n}\alpha_{i})$
is the multivariate beta function. The closed-form expressions for the
geodesics between Dirichlet distributions (and for the marginal beta
distributions, as were used in [7]) are unknown. It is known that, however,
this statistical manifold has non-constant negative sectional curvatures
everywhere [32]. [7] proposed to follow the specific path of Jacobi diffusion
processes on the marginal beta distributions, for which expensive modified
Jacobi polynomials needed to be precomputed. [56] further applied the flow
matching framework by considering the probability path of
$\mathrm{Dir}(1+te_{k})$ for some fixed maximum timestep. Still, expensive
calculations of the regularized incomplete beta functions were required.
Compared with these two generative models, our proposed SFM on categorical
distributions has the following advantages:

  * •

DDSM and DirichletFM have strong prior assumptions that the categorical
distributions follow some specific forward noising process which always
terminates at (or near) one-hot distributions. This assumption generally holds
for discrete generation, but cannot be generalized to more complex geometry on
the probability simplex, as we have demonstrated in our toy example of the
Swill roll on simplex dataset in Fig.3.

  * •

The forward noising paths are not the shortest geodesics in the sense of the
Riemannian metric. In our SFM, the vector fields always follow the geodesic,
thus following the sharpest direction of decreasing KL divergence (see
Sec.3.3).

  * •

Our SFM framework does not require the computation of expensive polynomials
and is more mathematically concise.

  * •

The final one-hot distributions are unreachable in [56], as it would require
an infinite timestep. Therefore, a variational bound needs to be derived for
likelihood estimation, which was not done in the original paper.

#### Multinomial Distribution

Consider $m$ i.i.d. experiments that follow a categorical distribution with
$n$ possible outcomes and probability $\mu=\\{\mu_{i}\\}_{i=1}^{n}$, a
multinomial distribution gives the probability of getting $x_{i}$ times the
$i$-th outcome with $\sum_{i=1}^{n}x_{i}=m$. The geodesic distance for
multinomial distributions is identical to that of the categorical
distributions up to a scaling constant [42]:

$d_{\text{mul}}(\mu,\nu)=\sqrt{n}d_{\text{cat}}(\mu,\nu)=2\sqrt{n}\arccos\left(\sum_{i=1}^{n}\sqrt{\mu_{i}\nu_{i}}\right)$
(26)

Therefore, our model can also be interpreted as multinomial flow matching on
the corresponding statistical manifold. The previous work of multinomial
diffusion (argmax flow) [27] transformed the one-hot distribution to
$ne_{k}-1$ in the logit space (soft one-hot logits) and assumed the common
Euclidean structure of the logit space for constructing the diffusion process.
This assumption did not consider the intrinsic geometry of the underlying
statistical manifold and also had the inaccessibility issue of the delta
measure as DirichletFM. We have demonstrated that SFM consistently
outperformed the naive multinomial diffusion on the Text8 datasets, indicating
the effectiveness of tracing the true geodesics.

##  Appendix B Likelihood Calculation

In this section, we provide more concrete mathematical details regarding the
exact likelihood calculation in the main text. We first make a distinction
between negative log-likelihood (NLL) and bits-per-dimension (BPD, including
bits-per-character, BPC) in this work. As we are dealing with probability
distributions over the statistical manifold of probability measures, it is
important for us to keep the difference between the probability spaces
$\mathcal{P}(\mathcal{P}(\mathcal{X}))$ and $\mathcal{P}(\mathcal{X})$. In our
notation, we use $\mu,\nu$ to denote elements in $\mathcal{P}(\mathcal{X})$,
or categorical distributions, and use $p,q$ to denote elements in
$\mathcal{P}(\mathcal{P}(\mathcal{X}))$, or distributions over categorical
distributions. We further use NLL to refer to $-\log p(\mu)$ for a specific
$\mu$ and BPD for categorical likelihood $-\log_{2}p(\delta|\mu)$ where
$\delta$ is the ground truth delta measure (one-hot distribution). Note this
definition of BPD (BPC) follows the convention in autoregressive language
models where the conditional probabilities for the next token are calculated
and averaged [17, 22]. In contrast, most previous work drew equivalence
between these two concepts.

###  B.1 Negative Log-Likelihood (NLL) Calculation

Our SFM framework enjoys the exact negative log-likelihood calculation as a
continuous normalizing flow model [36, 15]. It is worth noting that negative
NLLs are not uncommon [15]. As a probability density, the NLL can indeed go to
negative infinity at the boundary. Consider the arcsine distribution defined
on the finite support $[0,1]$ with the probability density function

$p(\theta)=\frac{1}{\pi\sqrt{\theta(1-\theta)}},\quad\theta\in(0,1)$ (27)

It is easy to see the density approaches infinity as $\theta$ approaches 0 or
1. Note that the arcsine distribution is a special case of beta distribution
$\mathrm{Beta}(\frac{1}{2},\frac{1}{2})$, which is in turn a special case of
Dirichlet distribution $\mathrm{Dir}(\frac{1}{2},\frac{1}{2})$. Each $\theta$
can be naturally viewed as the parameter for the Bernoulli distribution
(2-class categorical distribution), so the arcsine distribution indeed defines
a distribution over Bernoulli distributions with negative NLLs near the
boundary. In fact, for discrete data, the target data always come in as delta
measures, which indicates that the target distribution must be a convex
combination of delta measures with infinite likelihood:
$q=\sum_{i=1}^{n}\theta_{i}\delta_{i}$ where $\delta_{i}:=\delta_{\delta^{i}}$
is the delta measure over the underlying delta measure $\delta^{i}$ on the
discrete class $i$.

#### Variational Bound

We followed [7] to derive a variational bound as marginalized over a small
neighborhood of the delta measure in Eq.14 to obtain finite NLLs. The
variational bound can be derived via the standard variational approach as

$\displaystyle\log p(\delta)$
$\displaystyle=\mathbb{E}_{q(\mu|\delta)}\left[\log p(\delta)\right]$ (28)
$\displaystyle=\mathbb{E}_{q(\mu|\delta)}\left[\log\left(\frac{p(\delta,\mu)}{p(\mu|\delta)}\right)\right]$
$\displaystyle=\mathbb{E}_{q(\mu|\delta)}\left[\log\left(\frac{p(\delta,\mu)}{q(\mu|\delta)}\frac{q(\delta,\mu)}{p(\mu|\delta)}\right)\right]$
$\displaystyle=\mathbb{E}_{q(\mu|\delta)}\left[\log\left(\frac{p(\delta,\mu)}{q(\mu|\delta)}\right)\right]+D_{\text{KL}}(q(\mu|\delta)\|p(\mu|\delta))$
$\displaystyle\geq\mathbb{E}_{q(\mu|\delta)}\left[\log\left(\frac{p(\delta,\mu)}{q(\mu|\delta)}\right)\right]$
$\displaystyle=\mathbb{E}_{q(\mu|\delta)}\left[\log p(\delta|\mu)+\log
p(\mu)-\log q(\mu|\delta)\right]$

The posterior probability $q(\mu|\delta)$ serves as the forward diffusion
process in [6, 56] with a closed-form expression that depends on the timestep.
In our Riemannian flow matching formulation, however, such a forward process
is implicitly defined via the time-interpolation along the geodesics.
Nonetheless, we do know that the conditional distribution should approach the
delta measure when $t\to 1$ as the geodesic distance approaches 0:
$p_{1}(\mu|\delta)\approx\delta$. Therefore, as we are marginalized over a
small neighborhood, we can define a time-dependent posterior distribution
$q_{t}(\mu|\delta)$ that approaches $\delta$ as $t\to 1$. In principle, any
forward probability that approaches the delta measure at $t=1$ with an
analytical likelihood can be used to derive the lower bound. For example, we
can follow [7] to use the Jacobi diffusion process. To avoid the need for
expensive Jacobi polynomial calculation, we instead use simple indicator
measures over the simplex as

$q_{t}(\mu|\delta^{k})=\begin{cases}\frac{p_{0}}{(1-t)^{n-1}},&\mu_{i}\leq
t,1\leq i\leq n,i\neq k\\\ 0,&\text{otherwise}\end{cases}$ (29)

where $p_{0}$ is the base probability of the uniform distribution over the
simplex, $\mu_{i}=\theta_{i}$ is the $i$-th component of the categorical
distribution of $\mu$. In other words, $q_{t}(\mu|\delta^{k})$ is the time-
interpolated probability between the delta measure and the uniform
distribution: $q_{t}(\mu|\delta^{k})=t\delta_{k}+(1-t)U\Delta^{n-1}$. We have

$\log q_{t}(\mu|\delta^{k})=\log p_{0}-(n-1)\log(1-t)$ (30)

In practice, we use $t=0.995$ and provide more results in the ablation studies
in Appendix D.1. The categorical likelihood $p(\delta^{k}|\mu)$ is the
standard cross-entropy loss, also referred to as the _reconstruction loss_ in
some previous work [23]. It can be calculated as

$\log p(\delta^{k}|\mu)=\log\mu_{k}$ (31)

#### 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 $\Delta^{n-1}$. This can be efficiently
done by normalizing $n$ i.i.d random variables from the exponential
distribution:

$\tilde{\theta}_{i}\sim\mathrm{Exp}(1),\theta_{i}=\tilde{\theta}_{i}/\sum_{k=1}^{n}\tilde{\theta}_{k},\quad
i=1,2,\dots,n$ (32)

The uniform distribution over simplex is exactly the Dirichlet distribution
$\mathrm{Dir}(1,\dots,1)$. Therefore, the log-likelihood can be calculated as

$\log p_{0}=-\log\Gamma(n)$ (33)

where $\Gamma$ is the Gamma function. Note that the prior likelihood is
independent of the samples.

#### Change of Measure

The diffeomorphism $\pi$ in Eq.5 induces a push-forward measure on
$S^{n-1}_{+}$ which can be described by the change of measure identity

$\pi_{*}p(\pi(\mu))|\det d\pi(\mu)|=p(x),x=\pi(\mu)$ (34)

The change of measure term can be calculated using the change of volume
formula with the canonical coordinates $\\{\tilde{x}_{i}\\}_{i=1}^{n-1}$ and
$\\{\varphi_{i}\\}_{i=1}^{n-1}$ as

$\displaystyle\mu_{1}$ $\displaystyle=\tilde{x}_{1},\quad$ $\displaystyle
x_{1}=$ $\displaystyle\cos\varphi_{1}$ (35) $\displaystyle\mu_{2}$
$\displaystyle=\tilde{x}_{2},\quad$ $\displaystyle x_{2}=$
$\displaystyle\sin\varphi_{1}\cos\varphi_{2}$ $\displaystyle\vdots$
$\displaystyle\vdots$ $\displaystyle\mu_{n-1}$
$\displaystyle=\tilde{x}_{n-1},\quad$ $\displaystyle x_{n-1}=$
$\displaystyle\sin\varphi_{1}\dots\sin\varphi_{n-2}\cos\varphi_{n-1}$

where $0\leq\varphi_{i}\leq\pi/2$. The diffeomorphism $x=\pi(\mu)$ is given by

$\displaystyle\tilde{x}_{1}$ $\displaystyle=\cos^{2}\varphi_{1}$ (36)
$\displaystyle\tilde{x}_{2}$
$\displaystyle=\sin^{2}\varphi_{1}\cos^{2}\varphi_{2}$ $\displaystyle\vdots$
$\displaystyle\tilde{x}_{n-1}$
$\displaystyle=\sin^{2}\varphi_{1}\dots\sin^{2}\varphi_{n-2}\cos^{2}\varphi_{n-1}$

The change of measure term is

$\log|\det d\pi^{-1}(x)|=-(n-1)\log 2-\sum_{i=1}^{n}\log x_{i}$ (37)

Similarly, for the inverse mapping, the term is

$\log|\det d\pi(\mu)|=(n-1)\log 2+\frac{1}{2}\sum_{i=1}^{n}\log\mu_{i}$ (38)

#### Model Likelihood

The model likelihood can be viewed as the push-forward measure of the learned
flow:

$p_{t}=(\psi_{t})_{*}p_{0}$ (39)

We demonstrated in Sec.3.5 that this term can be calculated as the integral of
divergence in Eq.(12) plus the base probability in Eq.(33). More specifically,
the NLL can be obtained as the solution to the following ODE system back
through time with the initial condition of $x_{1}$ and $\log
p^{\text{ODE}}_{1}=0$:

$\frac{d}{dt}\left[\begin{array}[]{c}x_{t}\\\ \log
p^{\text{ODE}}_{t}\end{array}\right]=\left[\begin{array}[]{c}v_{t}(x_{t})\\\
-\operatorname{div}_{g}(v_{t})(x_{t})\end{array}\right]$ (40)

For manifolds that can be embedded in the ambient Euclidean space (e.g.,
simplex and sphere), the Riemannian divergence can be calculated as the normal
Euclidean divergence $\mathrm{div}_{g}=\mathrm{div}$ [15]. We further use
Hutchinson’s trace estimator [29] to efficiently estimate the divergence as

$\mathrm{div}\>v_{t}(x)=\mathbb{E}_{\varepsilon\sim\mathcal{N}(0,I)}[\varepsilon^{\top}\nabla
v_{t}(x)\varepsilon]$ (41)

The expectation can be efficiently computed using the vector-Jacobian product.

Combining Eq.(30), (31), (33), (37), (38) with Eq.(14) and (13), we can
efficiently calculate the NLL given arbitrary delta measures as input. The NLL
calculation scheme is visualized in Fig.2 and described in Alg.1.

Algorithm 1 NLL Calculation for Discrete Data

1:Sample $\tilde{\mu}_{1}\sim q_{t}(\mu|\delta)$ in Eq.(29) and calculate
$-\log q_{t}(\tilde{\mu}_{1}|\delta)$ and $\log p(\delta|\tilde{\mu}_{1})$.

2:Apply the diffeomorphism in Eq.5 to obtain
$\tilde{x}_{1}=\pi(\tilde{\mu}_{1})$ and calculate $\log|\det
d\pi^{-1}(\tilde{x}_{1})|$.

3:Solve the ODE system in Eq.(40) to obtain $x_{0}$ and $\log p^{\text{ODE}}$.

4:Apply $\pi^{-1}$ to obtain $\mu_{0}=\pi^{-1}(x_{0})$ and calculate
$\log|\det d\pi(\mu_{0})|$.

5:Calculate the base log probability $\log p_{0}(\mu_{0})$.

6:return NLL as in Eq.(14).

###  B.2 Bits-Per-Character (BPC) Calculation

We have demonstrated that the NLL can reach negative infinity. In contrast, in
our definition of BPC with the cross-entropy loss $-\log_{2}p(\delta|\mu)$, it
is always non-negative and consistent with previous autoregressive language
models. However, as our flow-based model is not autoregressive, we instead
follow [50] to calculate the pseudo-NLL/BPC for comparison. The pseudo-log-
likelihood of a $D$-dimensional data is defined as

$\frac{1}{D}\sum_{d=1}^{D}\log p(x^{d}|x^{-d})$ (42)

where $x^{-d}=\\{x^{k},1\leq k\leq D,k\neq d\\}$ represents the data without
the specific dimension index $d$. BPC can be similarly defined as
$-\frac{1}{d}\sum_{d=1}^{D}\log_{2}p(x^{d}|x^{-d})$. In our evaluation with
Text8, we always pick the middle token with a window of 128 known tokens
before and 127 known tokens after. The model tries to recover the missing
token given the conditional context. As our model only accepts one timestep
for all data dimensions, instead of fixing the context, we still uniformly
randomly sample all the initial noises from the simplex but substitute the
corresponding vector fields with the ground truth ones to make sure the
context tokens converge to the ground truth one-hot distributions.

##  Appendix C Experimental Setup

In this section, we further describe the experimental setup, model
architecture, and datasets.

###  C.1 Model Parameterization

Our SFM architecture is encoder-agnostic and can be applied to arbitrary
discrete generative tasks. Here, we describe the common setting of the flow
model across different datasets. We use the sinusoidal timestep embedding
described in [61] as $\mathrm{Emb}:[0,1]\to\mathbb{R}^{H}$. We follow [15] to
manually project the predicted vector field onto the corresponding tangent
space. For the spherical manifold, the projection can be described as

$v_{t}(x_{t})=\tilde{v}_{t}(x_{t})-\langle x_{t},\tilde{v}_{t}(x_{t})\rangle
x_{t}$ (43)

For linear flow matching on the simplex, the projection can be described as

$v_{t}(x_{t})=\tilde{v}_{t}(x_{t})-\frac{1}{n}\sum_{i=1}^{n}\tilde{v}_{t}(x_{t})_{i}$
(44)

The projection guarantees that the final prediction lies on the corresponding
tangent space, and the data points will stay on the manifold during Euler
sampling. We will always assume the projection is performed in the following
context of using $v_{t}$. The training stage of SFM is described in Alg.2. The
time complexity of our SFM framework is dominated by the underlying vector
field predictor with little overhead. Each model for binarized MNIST and
promoter design was trained on a single NVIDIA A100 GPU for 6-10 hours. Each
model for Text8 was trained on 8 NVIDIA A100 GPUs for about 7 days. We will
further describe the encoders for each generation task in the following
subsections.

Algorithm 2 Training SFM

1:while not converged do

2: Sample noise distribution $\mu_{0}\sim p_{0}(\mu)$ and target distribution
$\mu_{1}\sim q(\mu)$.

3: if optimal transport then

4: Do batch OT assignments of $\mu_{0}$ and $\mu_{1}$ according to the average
statistical distances.

5: end if

6: Apply the diffeomorphism in Eq.5 to obtain
$x_{0}=\pi(\mu_{0}),x_{1}=\pi(\mu_{1})$.

7: Sample $t\sim U[0,1]$ and interpolate
$x_{t}=\exp_{x_{0}}(t\log_{x_{0}}x_{1})$ using Eq.(19) and (20).

8: Calculate the conditional vector field
$u^{S}_{t}(x_{t}|x_{0},x_{1})=\frac{d}{dt}x_{t}=\log_{x_{t}}(x_{1})/(1-t)$.

9: Predict the vector field using $v(x_{t},t)$ and optimize the SFM loss in
Eq.(8).

10:end while

###  C.2 Model Sampling

The sampling process from the trained model can be described as solving the
differential equation $\frac{\partial}{\partial t}x_{t}=v_{t}(x_{t})$ from
$t=0$ to 1 with the initial conditional $x_{0}$ sampled from the prior noise
distribution. Alternatively, we can write the solution as the integral of the
learned time-dependent vector through time as

$x_{1}=x_{0}+\int_{0}^{1}v_{t}(x_{t})dt$ (45)

We always use the Dopri5 ODE solver [9] for our ODE sampling and NLL
calculation. For Euler method with $N$ discrete steps, the timesteps
$0,1/N,2/N,\dots,(N-1)/N$ are used with a step size of $1/N$. The sampling
stage is described in Alg.3.

Algorithm 3 Sampling from SFM

1:Sample noise distribution $\mu_{0}\sim p_{0}(\mu)$.

2:Apply the diffeomorphism in Eq.5 to obtain $x_{0}=\pi(\mu_{0})$.

3:if ODE sampling then

4: Solve $\frac{\partial}{\partial t}x_{t}=v_{t}(x_{t})$ using Dopri5 ODE
solver with initial condition $x_{0}$.

5:else$\triangleright$ Euler method

6: for $t\leftarrow 0,1/N,2/N,\dots,(N-1)/N$ do

7: $x_{t+1/N}=\exp_{x_{t}}(v(x_{t},t)/N)$

8: end for

9:end if

10:return $\mu_{1}=\pi^{-1}(x_{1})$

###  C.3 Swill Roll

The Swill roll on the 2-simplex is generated by normalizing the span of the
2-dimensional Swill roll to $[0+\varepsilon,1-\varepsilon]$ where
$\varepsilon$ is a small margin to make sure no point lies on the boundary.
The third dimensional is automatically obtained as
$\mu_{3}=1-\mu_{1}-\mu_{2}$. We fixed a random seed and generated 1000 samples
for a full batch training. The vector field predictor is based on simple
multi-layer perceptions (MLPs) and can be described as

$v(x_{t},t)=\mathrm{MLP}(\mathrm{MLP}(x_{t})\|\mathrm{MLP}(\mathrm{Emb}(t)))$
(46)

where $\|$ denotes concatenation and we used a hidden dimension of $H=128$.
Each model was trained for 2000 epochs using full batch gradient descent with
an initial learning rate of $10^{-3}$. 1000 samples were sampled using the
Dopri5 ODE solver, and NLL was calculated on the whole data samples 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 [56].

###  C.4 Binarized MNIST

We used the preprocessed binarized MNIST dataset from [49] which has a split
of 50k/10k/10k. We adopted the CNN-based vector field predictor from [53] with
additional additive time embeddings at each convolution layer’s output. The
model has 4 residual layers [24] and 4 refinement layers [35] with a total
number of 29.8M trainable parameters. All models were trained for 100k
iterations with a batch size of 256 (approximately 510 epochs) with an initial
learning rate of $3\times 10^{-4}$.

For the quantitative evaluation of the generated samples, we calculated the
FID score for 1000 generated samples for each model. The statistics for the
binarized MNIST dataset were calculated on the whole training set via the pre-
trained InceptionV3 [57] 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.

###  C.5 Text8

We followed previous work [23, 6] to use a fixed split of 90M/5M/5M with a
fixed sequence length of 256. We used a 12-layer transformer-based model
modified from [47] with an additional additive time embedding. The resultant
model has 85.1M trainable parameters. Noticeably, our model size is smaller
than [23] which used a full 24-layer model, but is similar to [37]. The models
were trained for 1M iterations with a batch size of 512 (approximately 23
epochs), an initial learning rate of $10^{-4}$, and an exponential moving
average (EMA) decay rate of 0.9999. We further used 1000 iterations as the
linear warmup of the learning rate and used a decay rate of 0.8 with a
patience of 3 and a validation interval of 2000 training iterations. The model
snapshot with the lowest validation loss was saved for evaluation. BPC and NLL
were calculated on the first 10k characters of the test set with 256 Euler
steps, and generated texts were obtained using either 256 Euler steps or the
Dopri5 ODE solver (see Appendix D.2).

###  C.6 Promoter Design

We used the splits from the dataset paper [7] that assign Chromosome 10 to the
validation set, Chromosomes 8 and 9 to the test set, and all the other 21
human chromosomes to the training set. This assignment can avoid potential
data leakage on the same chromosome. The transcription signals are provided as
a float number for each base pair position. We also followed [7] to use a
total number of 100k sequences with a context window of 1024 base pairs and
added a random offset of 10 to the sequence position during training. The
vector field predictor we used was identical to that in [7], with 20 stacks of
1D convolutional layers and a total number of 13.3M trainable parameters. Our
models were trained for 200k iterations with a batch size of 256 and an
initial learning rate of $5\times 10^{-4}$. The model snapshot with the lowest
validation SP-MSE on the validation set was saved for evaluation. The test SP-
MSE was evaluated on the generated samples on the full test set’s
transcription signals with 300 Euler steps of generation.

##  Appendix D Additional Result

In this section, we provide additional results of ablation studies as well as
generated samples for each task.

###  D.1 Ablation Study

We provide ablation studies of the effect of different sampling methods,
different sampling steps for the Euler methods, and different
$t_{\text{max}}\approx 1$ for NLL calculation. The results are provided in
Tab.5. The effect of the variational timestep $t_{\text{max}}$ matched the
results from [7], which stated that a closer timestep to the target data would
decrease the NLL. The reasons we choose $t=0.995$ instead of higher values
are: 1) we noticed numerical stability issues at a timestep very close to the
target, 2) we want a fair comparison with [7] which used an equivalent
timestep of 0.9975, and 3) as long as we use a same timestep, the results are
comparable within. We also found that NLL increased with the number of Euler
sampling steps. We hypothesize that it is due to the large divergence change
at $t\approx 1$ and naive Euler steps tend to overestimate its contribution.
Nonetheless, the Euler method with more than 300 steps gave a similar NLL as
the ODE solver. We further provide additional results of the FID scores on
different sample methods in Tab.5 in which we used $t=0.995$ and 300 Euler
steps. The final performance was quite similar between these two sampling
methods.

Table 4: NLL for different sampling methods, sampling steps, and
$t_{\text{max}}$.

.  Method #step $t_{\text{max}}$ NLL$\downarrow$ Euler 300 0.9995 -2.545 $\pm$
0.029 0.999 -2.442 $\pm$ 0.018 0.995 -1.689 $\pm$ 0.031 0.99 -1.019 $\pm$
0.032 100 0.995 -1.871 $\pm$ 0.020 500 -1.677 $\pm$ 0.022 1000 -1.647 $\pm$
0.024 ODE - 0.995 -1.687 $\pm$ 0.020

Table 5: FID for different sampling methods with $t_{\text{max}}=0.995$. For the Euler method, we used 300 steps.  Method | Model | FID$\downarrow$
---|---|---
ODE | SFM w/ OT | -1.687 $\pm$ 0.020
SFM w/o OT | -1.631 $\pm$ 0.027
LinearFM | 0.622 $\pm$ 0.022
Euler | SFM w/ OT | -1.689 $\pm$ 0.031
SFM w/o OT | -1.653 $\pm$ 0.028
LinearFM | 0.499 $\pm$ 0.022

###  D.2 Generated Samples

Figure 4: 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 models and different sampling settings are demonstrated in
Fig.4. The lower row demonstrates generated samples with different numbers of
Euler steps. Generated texts on the Text8 dataset using different models and
sampling methods are demonstrated below.

SFM w/ OT, ODE
---
⬇
a_gull_a_clock_as_sing_cowbreaks_diseases_in_forms_the_eilewes_of_clocks_contained_four_two_scarters_to_be_cursed_norse_the_hygg
y_alnanking_wade_of_joul_veine_in_one_ovias_sn_gl_cz_asi_evic_hal_art_diseashe_the_final_cangent_oedesis_spurl_the_last_one_surf
⬇
_cut_much_time_tow_through_the_wine_the_merlury_and_sore_characterography_is_lumaniua_nuzgen_by_khira_gormmaster_johk_herself_ha
s_nuquen_heads_some_conguees_following_row_the_living_core_coy_or_lanon_for_the_processions_and_workspandiios_for_three_e_conque
⬇
_was_lirelying_the_practice_of_finally_assista_he_was_one_of_the_superl_writings_though_using_the_qurdish_sieno_ardio_woo_daagul
itze_joy_to_pann_him_arrived_in_the_muslim_s_liorum_adam_and_were_ardius_who_was_to_have_not_thing_an_author_and_the_descent_hot
SFM w/o OT, ODE
⬇
running_of_experienting_and_preserved_in_to_them_no_longion_to_skin_this_is_not_serving_the_s_ready_pvorq_t_in_pchwadenaiosence_
bastiforces_in_wii_equal_partful_coward_dringetatt_is_often_devotes_to_devine_castles_faitst_gat_i_m_chilishka_fringelfat_to_the
⬇
s_help_to_lead_is_rich_to_his_court_the_acts_of_a_brateful_sharp_act_and_were_neploed_in_guilt_money_is_lost_recents_was_depende
nt_on_sweptumppichard_depeazed_n_using_ly_to_ruinto_war_the_kinbalt_rativation_that_will_snow_beploed_many_scritics_were_not_oft
⬇
_god_eight_one_doe_eight_german_embinor_eranger_and_the_fondam_s_andersen_diemstrop_filof_eight_one_silor_base_neget_gunshting_o
ther_pethors_in_support_the_texas_series_of_bright_pubtle_snap_czech_dogagova_s_thiru_buabing_aviator_space_k_ine_i_line_lines_a
LinearFM, ODE
⬇
_nine_three_zero_directions_requintly_the_post_wio_pranixenel_is_if_josalsted_but_shark_for_the_acchust_zoo_in_featural_acgule_o
f_micvilland_fame_there_two_nixed_by_for_the_aideadeqs_never_re_aine_is_mao_bun_jacubs_miles_for_ceatn_that_lost_not_roduitalist
⬇
form_lateuf_gh_movieg_thingporeefictine_slill_actral_the_dicav_siguraded_or_the_dorlas_oodi_ledic_missimines_a_consumts_tlaw_ref
arndy_by_hross_the_pous_by_darsick_murejfant_folloals_alenic_pnsemide_it_the_gremustate_fictional_and_tricity_to_to_the_obina_te
⬇
n_coan_in_treaty_of_areas_two_info_z_and_a_co_putu_for_the_rumo_stillia_defensive_productia__dit_one_zero_six_two_t_pe_itay_cdon
_ca_seven_as_used_as_mook_rejeired_space_foodrestions_cr_see_that_they_weure_some_pute_a_in_allime_amtc_s_oft_call_authought_pen
SFM w/ OT, Euler 256
⬇
by_keith_pang_ways_go_english_mouyr_hoch_in_one_nine_five_zero_th_ow_people_hugh_era_pizware_constitution_of_its_one_nine_five_o
ne_in_this_midway_co_english_anarchist_formed_clitue_these_america_produced_a_view_of_the_two_baggoots_and_america_mf_ayain_spen
⬇
icies_the_government_it_shot_to_replace_a_work_chagping_is_r_given_only_to_puddue_caparty_adminisment_in_two_zero_zero_three_au_
feez_giorrella_staffie_stewars_alfree_a_bum_often_boalhieved_by_significan_performances_of_axotas_a_in_the_world_wars_a_limit_fo
⬇
s_parented_by_two_ecll_which_repromite_s_tvhq_which_have_been_card_they_are_their_protrant_scorring_us_two_zero_million_raiters_
attempt_to_swing_a_pent_chatizen_of_the_tane_the_west_resattled_this_partement_are_anting_pasent_with_conse_x_dot_n_les_carling_
SFM w/o OT, Euler 256
⬇
art_of_the_dad_celsu_later_history_evea_s_d_gault_nentoci_lan_than_jans_grais_in_tood_s_daily_hobred_panes_than_was_recorded_in_
the_plot_of_dwaneke_s_elee_fate_carter_and_fauda_raylnell_mate_of_eveluaea_it_was_magnitheis_bereal_yearso_s_description_of_his_
⬇
_line_is_barer_edement_breeleys_lords_one_nine_seven_seven_gesadius_one_nine_eight_four_at_mike_cornelly_fonohre_lindon_pantheri
ngs_one_nine_eight_seven_one_line_the_cedenasipcs_the_ground_of_inci_are_gaurie_joining_the_story_and_tried_with_caway_inci_are_
⬇
f_strines_bose_such_as_throoptero_watermade_scuumenis_roche_and_streubs_the_unrbridted_apollophagia_i_two_five_zero_and_attends_
to_one_of_the_samkath_uses_maintaining_the_mosqui_tetonmeters_have_also_published_negau_in_place_ho_coubt_that_the_phoroachorido
LinearFM, Euler 256
⬇
h_jo_leek_the_form_buu_and_a_art_in_the_buel_he_bbu_bluiated_bpc_relations_of_debated_issued_that_age_denintenede_to_advers_adge
n_wmather_alconsensifia_cabip_lod_pyer_vubibax_which_s_vesskn_comesignearclij_ly_was_flee_down_becauid_the_pawer_of_cantix_are_s
⬇
tormut_of_duelibp_cas_unquests_pet_trokety_cal_totat_taiutext_two_logy_their_cylpheighfoding_in_par_has_sor_amale_optets_of_spyw
ork_for_sleampter_re_davi_cadiual_age_gestricational_gaie_the_cases_autharity_ouly_pop_radmitius__mamfare_the_diatuary_s_warking
⬇
w_cerip_mike_the_devisoll_rite_stringri_cometdral_king_hort_eitry_vix_orth_other_centogyika_venin_weads_nine_stan_shiup_extm_ziw
_xtermal_yy_tocharsical_lyvi_in_one_eighteight_peri_e_entoguian_drag_wham_paris_sin_res_nrigidal_ort_centric_gas_live_perby_the_