Categorical Flow Maps

• Categorical Flow Maps are generative models that extend continuous flows to discrete data by leveraging simplex geometry and ODE-based trajectories.
• They use variational flow matching and self-distillation to achieve rapid few-step sampling while preserving tractable likelihood estimation.
• The framework scales to billions of parameters, enabling state-of-the-art performance in text, image, and molecular graph generation.

Categorical Flow Maps are a class of generative models that extend the theory and practice of continuous normalizing flows and diffusion/flow matching to categorical (discrete) data domains. They facilitate accelerated few-step generation, enable flexible model parameterizations for categorical state spaces, and retain the tractable likelihood estimation and guidance techniques characteristic of continuous-flow generative models. Categorical flow maps have been applied to text, images, molecular graphs, and beyond, scaling to billions of parameters and achieving state-of-the-art performance in few-step and single-step sampling regimes (Roos et al., 12 Feb 2026, Davis et al., 8 May 2026).

1. Mathematical Formulation and State Space

Let $K$ be the cardinality of the categorical space. The fundamental state space of Categorical Flow Maps is the probability simplex $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$, which represents relaxed “one-hot” encodings of categorical variables. Each $x \in \Delta^{K-1}$ can be interpreted either as a categorical probability vector or as a soft relaxation of a discrete class.

The generative process is described by a continuous-time trajectory in $\Delta^{K-1}$, typically governed by an ordinary differential equation:

$$\frac{dx_t}{dt} = f_\theta(x_t, t), \quad x_0 \sim p_0, \quad x_1 \sim p_1$$

where $f_\theta$ is a learnable vector field parameterized so that the trajectories remain in the simplex, $p_0$ is a simple base distribution (often uniform or Gaussian), and $p_1$ is the categorical data distribution supported on simplex vertices (Roos et al., 12 Feb 2026, Davis et al., 8 May 2026, Williams et al., 31 Oct 2025, Dunn et al., 2024).

A key parameterization is in terms of a flow map:

$$X_{s,t}(x_s) = x_s + \gamma_{s,t} \big( \pi_{s,t}(x_s) - x_s \big), \quad \gamma_{s,t} = \frac{t-s}{1-s}$$

with $\pi_{s,t}(x)$ acting as a partial denoiser or endpoint predictor in $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$0. This preserves the geometry of the simplex and provides an efficient way to model endpoint-conditioned transitions for accelerated generation (Roos et al., 12 Feb 2026).

2. Variational Flow Matching and Training Objectives

The dominant training paradigm for modern categorical flow maps is variational flow matching (VFM). In this framework, one introduces a variational posterior $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$1—typically a product of categorical distributions with logits parameterized by the model—so the endpoint mean/minimizer instantly falls on the simplex (Eijkelboom et al., 2024, Roos et al., 12 Feb 2026):

$$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$2

The optimization objective is the time-averaged reverse KL:

$$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$3

Equivalently, this becomes a simple cross-entropy on the logits output by $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$4 (Eijkelboom et al., 2024, Roos et al., 12 Feb 2026, Davis et al., 8 May 2026). The associated learned vector field is

$$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$5

so training the cross-entropy loss drives $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$6 to the true conditional mean and ensures consistency of the learned transport.

In flow-based implementations that begin with a continuous base (e.g., Gaussian), simplex-to-Euclidean bijections such as the Isometric Log-Ratio (ILR) or Stick-Breaking (SB) transforms are used to parameterize flows in $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$7 while respecting the Aitchison geometry. Models may employ a Dirichlet interpolation (“dequantization”) to inject mass into the interior of the simplex for more stable and expressive density estimation (Williams et al., 31 Oct 2025).

3. Self-Distillation and Accelerated Sampling

Standard flow-matching or diffusion-based models typically require hundreds of function evaluations for high-fidelity samples. Categorical flow maps, via self-distillation, can compress the transport into a one-step or few-step flow map $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$8 (Roos et al., 12 Feb 2026, Davis et al., 8 May 2026):

• Progressive Self-Distillation (PSD): Enforces compositionality of short- and long-range flow maps by minimizing the KL between chained and direct transitions.
• Endpoint Consistency Lagrangian Distillation (ECLD): Penalizes inconsistency in endpoint predictions across pairs $$\Delta^{K-1} = \{ p \in \mathbb{R}^K : p_k \geq 0, \sum_k p_k = 1 \}$$9 using cross-entropy losses and temporal drift regularization (Roos et al., 12 Feb 2026).

At inference, flow-map samplers iterate a small number of times $x \in \Delta^{K-1}$0 (often 1–4) using the closed-form flow update:

$x \in \Delta^{K-1}$1

and discretize the output (via argmax) to obtain a categorical sample. This reduces the number of function evaluations by two orders of magnitude at near-equivalent sample quality (Davis et al., 8 May 2026, Roos et al., 12 Feb 2026).

4. Simplex Geometry, Priors, and Discrete-to-Continuous Lifting

Key to the fidelity and theoretical guarantees of categorical flow maps is honoring the geometry of the probability simplex:

• Aitchison Geometry: Used for defining distances and vector fields in $x \in \Delta^{K-1}$2 (Williams et al., 31 Oct 2025).
• Simplex-to-Euclidean Transforms: Orders and bijections (ILR, SB) ensure invertibility and tractable Jacobians, supporting exact density estimation and likelihood evaluation.
• Dirichlet Interpolations: Discrete categorical data (one-hot) is “dequantized” via Dirichlet mixtures so training flows can be performed in the relative interior of the simplex (Williams et al., 31 Oct 2025).
• Choice of Priors: The prior $x \in \Delta^{K-1}$3 can be uniform on the simplex, empirical marginals, barycentric, or Gaussian (latent flows). Choice of prior affects volume coverage, mode mixing, and ultimate sample validity (Dunn et al., 2024).

5. Theoretical Guarantees and Recovery

The categorical flow map framework yields several key theoretical properties:

• Perfect Discrete Recovery: Under mild conditions (e.g., Dirichlet interpolation with $x \in \Delta^{K-1}$4), the limiting induced discrete distribution from the model matches the true categorical distribution when discretized via argmax (Williams et al., 31 Oct 2025).
• Consistency and Likelihood: The reweighted time-integral of the VFM loss upper-bounds the negative log-likelihood, providing likelihood estimates even in the semi-discrete setting (Davis et al., 8 May 2026, Eijkelboom et al., 2024).
• Isometry and Scaling: ILR is guaranteed to be an exact isometry between the simplex with Aitchison geometry and Euclidean space, ensuring the flow dynamics preserve intrinsic distances (Williams et al., 31 Oct 2025).
• Scalability: Recent models scale to 1.7B parameters and vocabulary sizes $x \in \Delta^{K-1}$5, with practical adaptations for memory, time schedules, and adaptive weighting yielding robust training at scale (Davis et al., 8 May 2026).

6. Empirical Performance, Applications, and Impact

Categorical flow maps have achieved strong empirical results across multiple domains:

• Language Modeling: Achieve generative perplexity close to autoregressive baselines with as few as 4 steps; tractable semi-discrete ELBO enables likelihood-based evaluation and scoring for multiple-choice QA (Davis et al., 8 May 2026).
• Molecular Graph Generation: Single-step models rival multi-hundred-step baselines in FCD, uniqueness, and validity metrics on QM9 and ZINC datasets; flow map guidance and endpoint parameterization crucially improve few-step performance (Roos et al., 12 Feb 2026, Eijkelboom et al., 2024).
• Binarized Images: State-of-the-art few-step FID scores on MNIST; guidance mechanisms enable highly accurate conditional sampling (Roos et al., 12 Feb 2026).
• Text: Lower NLLs in few-step regimes on Text8 and LM1B compared to previous discrete diffusion baselines (Roos et al., 12 Feb 2026).

The capacity for few-step, non-autoregressive, parallelizable sampling makes categorical flow maps especially suitable for high-throughput and low-latency generation tasks.

7. Comparison to Alternative Approaches

Categorical flow maps unify and generalize several classes of generative models:

• Categorical Diffusion Models: These operate directly in discrete space via Markov chains; categorical flow maps retain the flexibility of continuous flows and can integrate denoising/diffusion mechanics as a limit (Hoogeboom et al., 2021, Eijkelboom et al., 2024).
• Categorical Normalizing Flows: Earlier approaches used continuous encoders with factorized decoders, relying on invertible flows and efficient latent partitioning but without geometric constraints—recent CFM approaches provide sharper geometric alignment and improved sample quality (Lippe et al., 2020).
• SimplexFlow/Dirichlet Flow: Explicit simplex-constrained flows can sometimes underperform Gaussian-latent flows in practice, raising ongoing questions about geometry vs. volume and the optimality of latent prior choice (Dunn et al., 2024).

A comparative summary is given below.

Approach	Geometry	Sampling	Likelihood	Empirical Regime
Categorical Flow Maps (CFM)	Simplex	1–4 steps (flow)	Yes (ELBO)	Text, graphs, images
Categorical Diffusion	Simplex	100–200 steps	Yes	Images, graphs
Categorical Normalizing Flows	Euclidean	Parallel, invert	Yes	Graphs, text, images
Gaussian-Latent Flows	Euclidean	1 step	Yes	Molecules, text

CFMs provide the tightest bridge between continuous and discrete modeling, balancing geometric faithfulness and sampling efficiency, supporting downstream inference tasks, and scaling effectively to large data and model sizes (Davis et al., 8 May 2026, Roos et al., 12 Feb 2026, Williams et al., 31 Oct 2025).

---

References:

• "Categorical Flow Maps" (Roos et al., 12 Feb 2026)
• "Scaling Categorical Flow Maps" (Davis et al., 8 May 2026)
• "Simplex-to-Euclidean Bijections for Categorical Flow Matching" (Williams et al., 31 Oct 2025)
• "Variational Flow Matching for Graph Generation" (Eijkelboom et al., 2024)
• "Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions" (Hoogeboom et al., 2021)
• "Categorical Normalizing Flows via Continuous Transformations" (Lippe et al., 2020)
• "Mixed Continuous and Categorical Flow Matching for 3D De Novo Molecule Generation" (Dunn et al., 2024)