Categorical Flow Matching on Statistical Manifolds

We thank you for your high recognition of our work's novel geometric insight and clear delivery. We will address your questions and concerns as follows.

Q1 Advantage of Riemannian structure

Unlike LinearFM which assumes a flat simplex and straight flows, SFM considers the Riemannian structure induced by the Fisher information metric for categorical distributions. The naive Euclidean assumption may fail to capture the true geometry of the statistical manifold. Furthermore, we demonstrated in Sec.3.3 that following the vector field induced by the Fisher information metric coincides with the steepest direction of the natural gradient that decreases the KL divergence, which may also contribute to a better performance from the optimization point of view. Another possible advantage comes from the curved geodesic in Fig.1. Unlike straight geodesic under the Euclidean setting, the Riemannian vector fields are curved towards the boundary (more parallel to the boundary), making it hard to overshoot across the boundary.

Q2 Comparison to DirichletFM

We have included comparisons to DirichletFM in Appendix A.4 to provide some insights into SFM's better performance. DirichletFM considered the specific probability path in which the target distributions are always assumed to be Dirichlet. Such an assumption is not applicable to more complex target distributions and fails the Swiss roll toy example in Fig.3. Additionally, its probability path does not reflect the shortest geodesic distance on the statistical manifold, which may lead to suboptimal flows and vector fields. In contrast, our SFM framework constructs flows based on the geodesic on the statistical manifold, where vector fields follow the steepest direction of decreasing KL divergence. Our method can applied to any source and target distribution, making it more flexible and efficient.

Q3 FisherFlow

We thank you for mentioning FisherFlow, which is a concurrent work closely related to ours. We note that the FisherFlow paper was uploaded on May 23 after the NeurIPS submission ddl. Therefore, SFM and FisherFlow should be considered concurrent work.

Both SFM and FisherFlow explored the Riemannian structure of the statistical manifold to establish flows and vector fields by constructing isomorphic mappings between the simplex and the sphere. We both demonstrated the theoretical connection to natural gradient and explored the minibatch-OT formulation during training. We took a step further to derive the exact likelihood over the probability space and a tight ELBO for the likelihood of discrete data. We carefully choose the prior, reference measure, and divergence to ensure comparability with other discrete diffusion/flow models.

While FisherFlow predominantly focused on bioinformatic tasks of DNA design, we conducted more comprehensive experiments across various domains with additional baselines to explore the superior performance and versatility of our method. We believe both works are important explorations on the use of Riemannian geometry structure on discrete generative tasks.

Q4 Conditional on measure

We apologize for the typo. We meaned "condition on $x_1$" where the conditional probability path has $p_1(x|x_1)\approx \delta_{x_1}$ at $t=1$, as described in the flow matching paper. In our case, each target point $x_1$ is a one-hot distribution $\mu_1$.

Q5 data dimension

The data dimension refers to the length of data. As an example, a DNA sequence with $D$ bases has data dimension $D$. We can represent this data as a matrix $X \in [0,1]^{D \times n}$, where $n=4$ reflects the four different categories (A, T, C, G). For simplicity, we assumed $D=1$ in the derivations, which is a standard practice. It is straightforward to extend this to multi-dimension by jointly modeling $D$ categorical probability measures.

Q6 Connection to natural gradient

Please see the common rebuttal.

Q7 Optimal transport

We consider minibatch-OT similar to [20,43,62] but on the statistical manifold. We match between a minibatch of the noises from the prior distribution and the samples from the target distribution with the smallest transport cost (defined line 189) based on the statistical distance defined in Eq.3. A thorough investigation of the theoretical benefit of minibatch OT can be found in the original paper. For Markov-chain based models like D3PM and MultiFlow, it is not possible to derive such a distance measure due to the discrete jump between Markov states.

Q8 Significance of NLL calculation

Generative models are trained to capture data distribution $p_\theta\approx p_\text{data}$. Therefore, it is natural and crucial to calculate the likelihood for a given data sample $p_\theta(x)$ as both the evaluation metric for generative models and confidence quantification for data. When evaluated on ground truth data, the NLL serves as a natural evaluation metric of how closely the generative model fits the data. Many common evaluation metrics including cross-entropy, perplexity, and bits-per-dimension are derived from likelihood estimation. With the ability to calculate likelihood, we can provide an intrinsic evaluation of flow models.

Additionally, NLL can be used to measure the confidence of given inputs and facilitate RLHF of flow models. Many policy-based RL requires explicit log-likelihood, and ELBOs that are loose (e.g. with impromptu noise schedule) could impact the effectiveness of RLHF. Our exact NLL formulation is more accurate and our ELBO definition over discrete data is arguably much tighter, potentially benefiting applications that explicitly rely on likelihood.

I thank the authors for the rebuttal, but I am not understanding one of the answers.

About the sphere metric, I don't see how taking a simplex and curving the interior (mapping from a flat simplex to a curved one) can prevent overshooting over the boundary. Could you please elaborate?

To clarify my doubt, think of the case of 3 categories. We then compare a flat equilateral triangle with straight lines, to a "quadrant" of a sphere, where geodesics on a sphere are so to say, "pieces of equators", or maximal circunferences on the sphere. Then the segment going from a vertex of the straight triangle to the opposite side, makes an angle $\alpha\in[\pi/3, \pi/2]$ with the boundary whereas, in the curved setting, the geodesic from a vertex of the quadrant to the opposite side makes an angle $\alpha=\pi/2$ with the boundary. So why does one trajectory overshoot less than the other?

I can see the performance of your experiments, but to me the geometry of the sphere quadrant and of a simplex still seem quite similar. In fact, adding to what I said in the review: the two metric spaces are not just $C^1$-diffeomorphic, but actually have finite distortion one with respect to the other, and I think the distortion constant is smaller than $2$. So, what would warrant such a big difference of performance then?

Given the above comparison of agles of geodesics, I don't think the answer is related to overshooting, it must be something else.

Emphasizing again, to be clear: this is just an important curiosity for me, because I can see the better performance of the experiments. But still I'm not convinced about the actual principle/reason beyond the better performance at the moment.

Thank you for your high recognition of our work's significance in the realm of discrete generative modeling and for raising interesting questions and suggestions. We will fix the typo and citation in the revision. Our responses to your questions are as follows:

Q1 Motivations

Sampling over categorical distributions instead of from a single categorical distribution

We agree that both formulations can be effectively used on discrete sampling, but their compatibility with different generative models may vary. It is important to note that SFM (and the majority of diffusion/flow models, VAEs, and GANs) are not autoregressive (AR). Unlike LM which predicts logits for the next single token deterministically given context, diffusion and flow models operate directly on the joint distributions of all tokens. Therefore, it is inefficient to learn a single set of deterministic categorical distributions to capture diverse joint distributions. Alternatively, the majority of non-AR methods seek to transform from prior noise distribution (e.g. multivariant Gaussians) to target joint distribution, effectively resulting in a two-step formulation where stochastic samples over joint distribution space are drawn first (i.e., logits) and then decoded with a second sampling step (argmax or softmax). Secondly, diffusion/flow models are naturally defined over continuous space. To leverage such a framework, existing discrete diffusion/flow models utilize relaxation to operate on logits which are then projected to probability simplex, similar to the statistical manifold we defined but without a properly defined geometry. In other words, the two-step sampling in SFM is not a niche design choice but rather universal for discrete diffusion and flow models (e.g., D3PM, SEDD, BFN, LinearFlow).

Our distinction to prior works is the proposal of a better objective for training the flow model over the categorical distributions (first step), where we equip the manifold with a Riemannian metric. For the second step, we always used argmax when decoding discrete samples.

Fisher-Rao metric

We would first like to clarify a potential misunderstanding. During training, we did not use the Fisher-Rao metric for calculating the distance between two sequences (i.e., discrete samples). Instead, a metric is necessary for defining probability paths along the geodesic from the prior distribution (usually uniform over the simplex) towards the target distribution (concentrated around the vertex). The role of the Fisher metric is to equip the manifold of categorical distributions with a Riemannian structure on which vector fields can be calculated as the learning target. The Fisher information metric is the natural canonical Riemannian metric on the statistical manifold. It also enjoys the benefit of following the "steepest" direction of the natural gradient of decreasing KL divergence, as we have mentioned in Sec.3.3.

Nevertheless, we respectfully disagree that treating discrete data as onehot distributions and using Fisher-Rao distance over them is unnatural. In fact, discrete data have been commonly treated as distributions in classic discrete generative model training, when using cross-entropy loss. Fisher-Rao distance is yet another metric between distributions, which indicates the steepest direction for KL divergence minimization. We will make this clear in the revision to better motivate this setup as suggested by the reviewer.

Q2 Name of SFM & title

As mentioned in the previous text and Sec.3, our proposed SFM is a general generative framework for modeling measures (distributions) over the statistical manifold of a family of parameterized distributions. In this work, we presented the realization of our framework on categorical distributions which have wide applications in various discrete generation domains. Therefore, we wish to keep the general naming of statistical FM and the categorical in the title to inform the audience about the application. We will think of a better title to avoid confusion.

Q3 Impromptu assumptions in prior work

For Line 29, we want to convey that the naive Euclidean assumption of the probability simplex or the logit space does not capture the true geometry of the statistical manifold. Similarly, in Line 52, we want to emphasize that the Euclidean assumption in previous existing models does not have solid mathematical grounds, which may lead to worse performance on statistical manifolds.

Q4 Fisher-Rao metric

The Fisher information matrix for categorical distributions is $g_{jk}(\mu)=\frac{\delta_{jk}}{\mu_j}+\frac{1}{\mu_n}$ for $1\le i,j,\le n-1$, where $\delta_{jk}$ is the Kronecker delta. Substituting this into $\langle u,v\rangle_\mu$ leads to the Riemannian inner product in Eq.4.

Q5 Sampling in Eq.8

We will rewrite the sampling as $x_0\sim\pi_*(p_0(\mu))$ where $\pi_*(p_0)$ denotes the standard pushforward measure induced by the diffeomorphism $\pi$. This is equivalent to first sampling $\mu_0\sim p_0(\mu)$ and then taking $x_0=\pi(\mu_0)$.

Q6 Connection to natural gradient

Please see the common rebuttal for details.

Q7 Non one-hot data

We really appreciate your recognition of the potential usage of SFM on general categorical distributions besides one-hot. In the paper, we have provided the Swiss Roll example as a demonstration, where other baselines with strong priors may fail. We believe that SFM could be of interest to Bayesian inference, distributional RL, and other tasks that involve sampling over distributions. We defer these use cases to our future work to extend SFM due to limited time and resources.

We appreciate your recognition of our work's clear presentation and novel geometric perspective for discrete generation. We'd like to clarify that Text8 results of MultiFlow were added on Jun 5 after the NeurIPS ddl, thus we were unable to compare them although it was cited.

We wish to emphasize that our major contributions are beyond the Text8 experiment. SFM is a general generative framework for measures of categorical distributions, and we have demonstrated the effectiveness of SFM across diverse domains including computer vision (binarized MNIST), bioinformatics (DNA design), and NLP(Text8). We also demonstrated performance improvements using evaluation metrics independent of likelihood calculation (FID & Sei MSE). We further established the underlying connection of SFM to natural gradient and proposed likelihood formulation with comparability in mind.

We respectfully disagree that Text8 is the only non-toy dataset. DNA design has a tangible societal impact, and DDSM & DirichletFM all predominantly focused on this task. The data dimension of MNIST (784) and DNA (1024) is also larger than Text8 (256). Nonetheless, we understand the importance of text8 evaluation and now include additional results in the common rebuttal following your suggestions. Below we provide more details and justifications in response to your comments.

Q1 Choice of evaluation metrics

We're afraid that you might have misinterpreted our evaluation setup. For Swiss roll and MNIST, all baselines are diffusion or flow models with continuous parameterization. We never compared NLLs with discrete-space likelihood from discrete models (e.g. autoregressive) in these tasks. Secondly, all NLLs used in discrete tasks (Tab.1,2,4,5) are NOT pseudo-likelihoods. They are based on the variational bond described in Eq.14 and are comparable to ELBOs from the diffusion models we compared. Only BPC for Text8 (Tab.2) was calculated based on pseudo-likelihoods as described in Appendix B.2. According to the original PLL paper, it provides a reasonable evaluation metric between non-autoregressive models.

We used the identically variational formulation in the DDSM paper to be comparable. We also provided FID scores that do not rely on NLL calculation, for which our models still outperformed other baselines by a margin. More discussion on the validity of such a variational likelihood can be found in Sec.3.5 of DDSM.

We justified our choice of $t_\text{max}=0.995$ in the calculation of the variational bound for NLL in Appendix D.1. We carefully chose this value to make sure it was comparable to DDSM settings. Our ablation study on the choice of $t_\text{max}$ agreed with that in DDSM, where $t$ closer to data gives a tighter estimated bound (it can't get arbitrarily low). We believe it is reasonable to compare DDSM and LinearFM using similar $t_\text{max}$.

Q2 Additional results on Text8

We provide GPT-NLL vs token entropy, and representative samples for SFM and MultiFlow with $T=0.8,0.9,1$ for visual quality check. SFM has closer to data token entropy and comparable GPT-NLL with other methods in MultiFlow's result, although slightly worse on the latter. We respectfully disagree that our samples "look extremely poor", especially when compared with MultiFlow results with matching entropy ($T=1$). The perceptual quality of SFM and MultiFlow are very close, with noticeable misspellings in both. We also found GPT-NLL very unreliable and can be easily fooled by random strings. As an example, the following sample from MultiFlow at $T=1$ looks even worse despite a GPT NLL of 6.648.

she hill obhalarnnach eochrans eileann munthar cearlomha mhaonna  tardare mho mord tore lian linn mu phaile gael cangallauig laothuis guilleith leois glion guildh lara gall innonte tilbonne guilecht shuachtansert guillaste guatnaoic asthache cuichant conai

We will include these results on Text8 for a more comprehensive evaluation as you suggested, and tone down on our original claims. We sincerely hope that our contributions besides text8 can be valued properly.

Q3 Change of measure term

The change of measure term describes the change of the log-likelihood between different manifolds as $\log p_S(x)=\log p_\Delta(\mu)+\log |\det d\pi(\mu)|$ where $x=\sqrt{\mu}$. Although the density is always well-defined, the log-likelihood is undefined at the boundary, as it involves taking the logarithm of a zero probability density (see Eq.37 & 38).

Q4 Forward probability $q(\mu|\delta)$

The detailed definition of the forward probability $q(\mu|\delta)$ was provided in Eq.29 & 30 in Appendix B.1. The forward diffusion process $q_t$ defines a small neighborhood for variational estimation. In DDSM, the authors followed the fixed probability path of the Jacobi diffusion process with known forward probability. In our flow-based setting, we can use simpler indicator measures as the linear interpolation between the delta measure and the uniform distribution on the simplex.

Q5 ODE solver

As described in Sec.4.2, all generated results and NLL calculations are based on the Dopri5 solver unless otherwise specified. The Dopri5 solver is an adaptive solver with a good balance between accuracy and efficiency.

Q6 LinearFM Sensitivity to solvers

We noted that LinearFM tends to have very negative divergence around the one-hot distribution. Therefore, Euler step may overestimate the contribution of divergence in this case while adaptive solvers can provide a more accurate result. In contrast, as the SFM vector field is defined on the sphere, it is less sensible to solvers and results in a more stable divergence.

Q7 Typos & minor issues

The results in Tab.5 are NLLs. The ODE likelihood defined in Eq.12 is indeed the change of log-likelihood: $\log p^\text{ODE}:=\log p(\mu_1)-\log p(\mu_0)$. We thank you for pointing out these issues, and we will fix them in the revised manuscript.

Thank you for your appreciation of our work's clear presentation as a natural and timely extension of Riemannian FM on discrete data. We thank your thoughtful suggestions and will fix typos and clarify minibatch-OT in the revision. We'd like to address your comments as follows.

Q1 Regarding likelihood definition

We mentioned the Radon-Nikodym derivative in 2.1 as the definition of probability density functions on general manifolds. In our likelihood derivation (Sec.3.5), the density $\frac{dP}{dV}$ can be well-defined over the statistical manifold which is finite-dimensional ($n-1$ dimensional for $n$ categories) using the Riemannian volume as the reference measure. This is in contrast to works like functional FM where data points are continuous functions lying on the infinite-dimensional Hilbert space. We will clarify this as you suggested.

In terms of the choice of measure, we make sure to use a consistent reference measure for all change-of-measures (COM) in Sec 3.5. Note that the Riemannian volume on the sphere coincides with the Euclidean one (see response to Q2), and only the COM for the two transformations $\pi$,$\pi^{-1}$ will be impacted by the choice of measure. To make the NLL comparable to Euclidean flows, we have kept the reference measure consistent as the Euclidean volume form in the equations in Appendix B.1. We will make these mathematical details clearer and explicitly state the choice of volume in the revised manuscript.

For categorical distributions $\mu$, we used the canonical counting measure over the discrete sample space $\mathcal{X}=\{0,1,\dots,n\}$.

Q2 Regarding likelihood computation

We thank your insightful question regarding the divergence in NLL calculation. You are completely right that even though vector field is learned with Riemannian geometry, we can still compute likelihood with Euclidean divergence. As a matter of fact, we have taken comparability into consideration in our derivation of likelihood. Note that for manifolds that can be embedded in the ambient Euclidean space (e.g., simplex and sphere), the Riemannian divergence is the same as Euclidean divergence (for embeddable manifolds, we have $\log|\det g|\equiv0$). By using the Euclidean volume form as the reference measure in all COMs and computing prior likelihood on Euclidean simplex, we effectively used Euclidean divergence on the simplex for SFM and LinearFM to guarantee that the results are comparable.

It is possible, though, to calculate the Riemannian likelihood for SFM. To do this, we need to make several adaptations. The prior likelihood should be adjusted to $p_0=\Gamma(n)/\sqrt{|\det g|}$ such that the integral $\int_\Delta p_0dV=\int_\Omega p_0\sqrt{|\det g|}d\mu_1\cdots d\mu_{n-1}=1$. The change of measure term in Eq.37 & 38 also needs to be adjusted as $\log|\det d\pi|=(n-1)\log 2$ and $\log|\det d\pi^{-1}|=-(n-1)\log 2$.

Q3 Regarding choice of metric in designing conditional $u_t$

Thank you for mentioning the nice property of Riemannian $u_t$ being parallel to the boundary. We'd like to mention that our visualization of the logarithm map can be viewed as a demonstration of $u_t$'s direction (up to a constant scaling factor that depends on $t$), as the vector field can be calculated in terms of the logarithm map as $u_t(\mu_t|\mu_0,\mu_1)=\log_{\mu_t}(\mu_1)/(1-t)$ (assuming a linearly decreasing geodesic distance). We indeed noticed that the Riemannian structure induced by the Fisher information metric leads to vector fields more parallel to the boundaries. This can also be demonstrated mathematically by looking at the logarithm map in Eq.23. Consider the direction term $\sqrt{\mu\odot\nu}-\langle\sqrt{\mu},\sqrt{\nu}\rangle \mu$. For $\mu$ close to the boundary with $\mu_k\approx 0$, its corresponding vector field will also have a close to $u_k\approx 0$ component, which is different from linear flow matching's fixed $\nu-\mu$.

We hypothesize that one potential benefit of such a curved geometry over the naive Euclidean geometry is that the former helps prevent overshooting across the boundaries. Specifically, consider a target point on the boundary. The Euclidean vector field will continue to push the points outside the manifold, whereas the Riemannian vector field tends to travel parallel to that boundary to prevent going across the manifold. The sphere exponential map also naturally ensures the point (after transformation back) lies on the simplex. Once overshooting happens during sampling, the model may exhibit undefined behavior as it was never trained on points outside the manifold. We will emphasize this benefit more clearly in our future revision.

Q4 Regarding the connection to natural gradient

Please see the common rebuttal for details.