Research map / Probabilistic and scientific modeling / Variational and generative sampling
Categorical Flow Matching on Statistical Manifolds
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Statistical Flow Matching (SFM), a flow-matching framework on the Fisher-Rao statistical manifold of categorical distributions that uses closed-form geodesics via the sphere diffeomorphism, enabling exact likelihood and discrete generation.
Problem
Discrete generative models either rely on Markov jump processes or assume ad hoc geometry (Euclidean simplex or logit space), lacking the true geometry of categorical distributions and exact likelihood.
Approach
Equip the manifold of categorical distributions with the Fisher information metric and map it diffeomorphically (sqrt map) to the positive orthant of the unit sphere, where closed-form geodesics (great circles) allow simulation-free Riemannian flow matching with stable training; exact likelihood via change-of-measure terms and a variational bound for Dirac data.
Key results
- On binarized MNIST, SFM w/ OT achieves FID 4.62 versus 15.34 for DFM and 7.79 for DDSM.
- On Text8, SFM w/o OT reaches 1.386 BPC, better than D3PM-absorb (1.47) and LinearFM (1.651), but worse than SEDD-absorb (1.32).
- On promoter design, SFM w/o OT achieves the lowest SP-MSE of 0.0258 among listed methods.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use SFM with optimal transport instead of LinearFM | Binarized MNIST generation · Binarized MNIST | FID: 5.91 -> 4.62 | Helped |
| use SFM with optimal transport instead of LinearFM | Binarized MNIST generation · Binarized MNIST | NLL: 0.622 ± 0.022 -> -1.687 ± 0.020 | Helped |
| add optimal transport (ablation) | SFM · Binarized MNIST | FID: 5.15 -> 4.62 | Helped |
| add optimal transport (ablation) | SFM · Binarized MNIST | NLL: -1.631 ± 0.027 -> -1.687 ± 0.020 | Helped |
| add the diffeomorphism π (ablation) | SFM · Binarized MNIST | FID: 8.2073 -> 4.62 | Helped |
| add the diffeomorphism π (ablation) | SFM · Binarized MNIST | NLL: 0.928 ± 0.035 -> -1.687 ± 0.020 | Helped |
| use SFM instead of LinearFM | Text8 character generation · Text8 | BPC: 1.651 ± 0.027 -> 1.399 ± 0.020 | Helped |
| add optimal transport (ablation) | SFM · Text8 | BPC: 1.386 ± 0.033 -> 1.399 ± 0.020 | No clear effect |
| use SFM without optimal transport instead of SFM with optimal transport (ablation) | Promoter design · Promoter Design | SP-MSE: 0.0279 -> 0.0258 | Helped |
| use SFM instead of LinearFM | Promoter design · Promoter Design | SP-MSE: 0.0282 -> 0.0258 | Helped |
Surprising findings
- Optimal transport did not significantly improve Text8 performance and slightly hurt promoter design SP-MSE.
- Exact NLL can be negative and unbounded near the simplex boundary, so NLL is not directly comparable to discrete BPC.
Limitations
- Generation is iterative refinement and cannot change the size of the input, a limitation relative to autoregressive models.
- Assumes class independence so the Fisher metric induces the canonical Riemannian structure; ordinal data such as CIFAR-10 pixel values may be suboptimal.
- Likelihood on Dirac data requires a variational bound, and exact divergence is only computed for D=1.
Experimental setup
Datasets: Swiss roll binarized MNIST Text8 human promoter sequence dataset
Benchmarks: binarized MNIST Text8 promoter DNA design (SP-MSE) Swiss roll on simplex
Metrics: NLL FID BPC SP-MSE
Baselines: LinearFM DDSM DirichletFM D3PM DFM SEDD-absorb BFN MultiFlow Argmax Flow Bit-Diffusion
Builds on: Riemannian flow matching conditional flow matching Dirichlet flow matching D3PM SEDD
Closest papers by meaning
- Rectifying Categorical Flows on Statistical Manifolds for One-Step Generation (NeurIPS 2026 · similarity 0.93)
- Metric Flow Matching for Smooth Interpolations on the Data Manifold (NeurIPS 2024 · similarity 0.90)
- Beckmann Transport Models: From Autonomous Flows to One-Step Maps (NeurIPS 2026 · similarity 0.90)
- Fisher Flow Matching for Generative Modeling over Discrete Data (NeurIPS 2024 · similarity 0.89)
- Discrete Flow Matching (NeurIPS 2024 · similarity 0.89)
- Variational Flow Matching for Graph Generation (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.