Categorical Flow Matching on Statistical Manifolds

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{\mathop{}\!\mathrm{d}\mu}{\mathop{}\!\mathrm{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}}\mathop{}\!\mathrm{d}\mu=\int_{\mathcal{X}}p(x;\theta)\mathop{}\!\mathrm{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

Rao demonstrated in his seminal paper [51] 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 [51], 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}}\mathop{}\!\mathrm{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 [44], which motivates our method.

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 a flow-based model can be viewed as the continuous generalization of the score matching (diffusion) model [56, 57, 26] while allowing for a more flexible design of the denoising process. The Riemannian flow matching framework [14] further extends CFM to general manifolds on which a well-defined distance metric can be efficiently computed.

Consider a smooth Riemannian manifold $\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{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t}\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 [38, 14] demonstrated that a tractable objective can be derived by conditioning on the target data $x_{1}$ at $t=1$ of the probability path. The Riemannian flow matching objective can be formulated as [14]

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. [14] further demonstrated that if the exponential map and logarithm map can be evaluated in closed-form, the condition flow can be defined as $x_{t}=\exp_{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{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t}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.

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 Dirac 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 [51, 44] as

The tangent 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

where $\mathcal{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

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}$:

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:

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.

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:

where $p_{0}$ is the prior noise distribution over $\mathcal{P}$ and $q$ is the data distribution; $\pi_{*}$ denotes the standard pushforward measure. $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 Alg.2 and 3 in Appendix C.

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

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.

We further provide an interpretation of our proposed statistical flow matching framework from the perspective of an optimization problem. From the optimization viewpoint, for a generative model on a statistical manifold, we want to minimize the discrepancy between the target distributions and the generated distributions. Naturally, the Kullback-Leibler divergence $D_{\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}$ with $\Delta\theta=\theta-\theta_{1}$ gives

From the geometric viewpoint, the geodesic, by definition, is a (locally) length-minimizing curve with respect to the corresponding Riemannian metric. Therefore, by following the direction of the vector field that decreases the geodesic element $\mathop{}\!\mathrm{d}s^{2}=\|\mathop{}\!\mathrm{d}\theta\|^{2}_{g}\approx\|\Delta\theta\|^{2}_{g}$, we are indeed following the steepest direction that minimizes the local KL divergence. In this sense, the Fisher information metric defines a second-order optimization scheme for the KL divergence by following the “steepest” direction of the Hessian. Indeed, The steepest direction $\Delta\theta$ that decreases the KL divergence is known as the natural gradient in the existing literature [2, 3] and has been explored in optimization known as natural gradient descent[48, 42, 46]. 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, 42] demonstrated that stochastic natural gradient descent is asymptotically “Fisher efficient” and is indeed the steepest descent direction in the manifold of distributions where distance is measured in small local neighborhoods by the KL divergence. Following the geodesic defined by the Fisher information metric, our SFM framework also shares these benefits with an additional advantage of analytical expressions for geodesics, as we focus on the family of categorical distributions instead of general statistical models. Such a theoretical connection may contribute to the better performance of SFM.

The geometric viewpoint of the statistical manifold offers a continuous and differentiable formulation of generative modeling and also provides a robust way to measure the distance between two categorical distributions via the well-defined geodesic distance in Eq.(3). In contrast, the Markov chain-based methods cannot establish a robust distance measure due to the stochastic jumps between discrete states. Inspired by the optimal transport formulation in previous work [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 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.

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 [14, 43]:

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 pushforward probability measures $p_{t}(x_{t})$ defined via the learned flow can be obtained as the integral of the Riemannian divergence back through time on the simulated trajectory $x_{t}$. Following [9, 7], we define the ODE log-likelihood as the change of the log-likelihood as:

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 [30] 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 pushforward 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\mathop{}\!\mathrm{d}\pi(\mu)|=P(x)$ for $x=\pi(\mu)$. Therefore, by adding the two log-determinant of the pushforward measures, the log-likelihood can be formulated as

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 [62]. Following [7], we derive a variational lower bound for the likelihood as the marginalized probability over the small neighborhood of a Dirac measure $\mu_{1}=\delta$ at the boundary:

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].

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

We note that most existing models assume the target data to be Dirac measures (one-hot distributions) and cannot be applied to this task. Both DDSM [7] and DirichletFM [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 Eq.(13) and was averaged over all points in the dataset. Though linear flow matching could also capture the geometry of the data, our SFM outperformed it in terms of NLL.

We extend our SFM to model discrete generation tasks in computer vision. The binarized MNIST dataset [53] is the binarized version of the original MNIST dataset [34] 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] and DFM [22] as examples of masked diffusion models for discrete generation. We used a 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 Tab.1.

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

The Text8 dataset [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-per-character (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.3 and generated texts are provided in Appendix D.3. The results for the autoregressive language models are also provided as a reference (Discrete Flow [64] is based on autoregressive normalizing flow).

Note that, unlike all of the other baselines, SFM and LinearFM do not directly optimize such a BPC as a training objective, so such an evaluation metric is unfavorable for our model (see Appendix B.2). Despite such an unfavorable evaluation, our proposed SFM still achieved the second-best performance among other diffusion and flow baselines that were directly trained with cross-entropy losses. We also noted a significant performance gap between SFM and the naive linear flow matching on simplex, highlighting the importance of capturing the intrinsic geometric properties of the underlying statistical manifold. Optimal transport does not significantly affect the final performance here, which we presume might be due to the long training stage in which vector fields were well-explored. Additional evaluation following [12] is provided in Appendix D.1, in which SFM achieved the best generation entropy as evaluated by the pre-trained GPT-J-6B model [67].

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