PaperlayerView exact revision on arXiv

Source-linked AI summary

Categorical Flow Maps

Daan Roos, Oscar Davis, Floor Eijkelboom, Michael Bronstein, Max Welling, İsmail İlkan Ceylan, Luca Ambrogioni, Jan-Willem van de Meent

arXiv:2602.12233v1cs.LG

TL;DR

Discrete generative models need accelerated inference, but discrete interpolations are costly and difficult to combine with flow-map distillation. Categorical Flow Maps use continuous simplex-constrained endpoint trajectories and endpoint-consistency self-distillation for categorical generation. The method delivers strong few-step results across molecular graphs, images, and text, while also supporting test-time steering.

  • Problem

    Discrete-interpolation generative models require many steps and have stochastic trajectories, making direct flow-map self-distillation difficult for discrete data.

  • Method

    Categorical Flow Maps transport a continuous prior to a discrete target through simplex-constrained endpoint predictions and use an endpoint-consistency distillation objective.

  • Results

    CFMs achieve strong few-step generation across molecular graphs, images, and text, including 93.5% ZINC validity at 1 NFE and up to 3.5× lower PPL than UDLM in low-NFE LM1B evaluation.

  • Takeaways & Limitations

    The continuous flow-map formulation provides a practical route to self-distillable categorical generation and enables test-time steering with adapted guidance and reweighting.

  • Takeaways & Limitations

    For ECLD, Euler integration can match or outperform flow-map sampling at higher step counts, indicating a gap between the learned map and velocity field outside the low-NFE regime.

Abstract

from arXiv · show

We introduce Categorical Flow Maps, a flow-matching method for accelerated few-step generation of categorical data via self-distillation. Building on recent variational formulations of flow matching and the broader trend towards accelerated inference in diffusion and flow-based models, we define a flow map towards the simplex that transports probability mass toward a predicted endpoint, yielding a parametrisation that naturally constrains model predictions. Since our trajectories are continuous rather than discrete, Categorical Flow Maps can be trained with existing distillation techniques, as well as a new objective based on endpoint consistency. This continuous formulation also automatically unlocks test-time inference: we can directly reuse existing guidance and reweighting techniques in the categorical setting to steer sampling toward downstream objectives. Empirically, we achieve state-of-the-art few-step results on images, molecular graphs, and text, with strong performance even in single-step generation.

1 Introduction

Accelerated inference has advanced for continuous generative models, but discrete interpolations remain costly and difficult to distill. Categorical Flow Maps address this by using continuous simplex-constrained trajectories, enabling self-distillation and strong few-step generation across discrete modalities.

  • Motivation: Continuous-domain methods have achieved strong few-step generation, motivating whether accelerated inference can extend to inherently discrete text, sequences, and graphs.The section frames discrete data as an important setting beyond images and other continuous domains.
  • Challenges: Discrete-interpolation models require many generation steps, while their stochastic discrete trajectories prevent direct application of flow-map objectives.These costs can make test-time generation resemble permuted autoregressive generation.
  • Method: Categorical Flow Maps use continuous trajectories and simplex-constrained endpoint predictions to make self-distillation applicable to categorical data.The endpoint parametrisation also supports a cross-entropy loss that provably bounds the Lagrangian residual.
  • Results: Single-step CFM generation achieves over 95.8% valid molecular graphs on QM9, 93.5% on ZINC, 10.1 FID on Binary MNIST, 5.33 NLL on Text8, and 274.87 Gen-PPL on LM1B.These results cover molecular graphs, images, and text, demonstrating strong performance even at one step.

2 Background

Flow matching transports a simple prior to data through a continuous probability path, while flow maps learn finite-interval transport directly to reduce costly numerical integration. Variational formulations parameterize endpoint predictions, and categorical endpoints naturally use simplex-valued probabilities with cross-entropy training.

  • 2.1 Stochastic Interpolants: Stochastic interpolants continuously connect samples from an easy prior p0 to target data p1 and induce a probability path and flow.The standard setup uses differentiable interpolation functions, with straight-line interpolation as the main case considered.
  • 2.1 Stochastic Interpolants: Velocity-field learning minimizes a mean-squared-error objective, with the drift given by a conditional expectation along the interpolated path.The probability flow is defined by the differentiable path, and the velocity field and conditional mean determine each other.
  • 2.2 Flow Maps: A flow map X_s,t transports a state at time s directly to time t, allowing generation by applying X_0,1 to a sample from p0 instead of numerically integrating the velocity field.This directly targets the tens to hundreds of network evaluations often required for high-quality numerical integration.
  • 2.2 Flow Maps: Flow-map self-distillation trains finite-interval transport maps together with instantaneous velocity predictions, using Lagrangian self-distillation and a stabilizing flow-matching loss.The reviewed Lagrangian objective samples t uniformly and s uniformly between 0 and t.
  • 2.3 Variational Flow Matching: Variational Flow Matching estimates the conditional endpoint distribution rather than directly estimating velocity, and categorical endpoints yield simplex-preserving cross-entropy training.The simplex is the set of nonnegative class probabilities summing to one; matching its conditional mean recovers the velocity field.

3 Categorical Self-Distillation

Categorical Flow Maps use simplex-constrained endpoint predictions to make continuous flow-map self-distillation applicable to categorical data. The resulting endpoint-consistency objective supports stable training and test-time steering through existing stochastic-interpolant techniques.

  • 3.1 Constructing Categorical Flow Maps: Endpoint parametrisation replaces unconstrained velocity prediction with a partial denoiser supported on the data simplex.This preserves flow-map structure while incorporating the known support of categorical distributions.
  • 3.1 Constructing Categorical Flow Maps: The tangent condition ensures that the partial denoiser at s = t recovers the instantaneous velocity field, while the Lagrangian condition enables self-distillation.These two properties allow flow-map objectives for s < t and the VFM endpoint loss for s = t.
  • 3.2 Distillation through Endpoint Consistency: The ECLD objective uses endpoint cross-entropies to align teacher and student predictions while regularising temporal drift.Its combined form is L_ECLD := 4L_CE-EC + 2L_TD, and both terms operate on endpoint predictions rather than velocities.
  • 3.2 Distillation through Endpoint Consistency: When endpoint consistency and temporal drift both vanish, the Lagrangian residual is zero and the learned map is the unique flow map.The consistency term compares teacher and student endpoint predictions, while the drift term captures variation of the endpoint predictor with target time.
  • 3.2 Distillation through Endpoint Consistency: The authors report that the endpoint parametrisation substantially outperforms unconstrained velocity-field flow maps, especially on higher-dimensional ZINC.Improvements are described as modest on QM9 but substantial on ZINC.
  • 3.3 Test-time inference in Categorical Flow Maps: Continuous categorical flow maps permit test-time guidance by using endpoint lookahead to steer sampling toward differentiable rewards, with SMC correcting the dynamics for unbiased tilted sampling.Rewards can be evaluated using hard straight-through discretisation or directly on the soft simplex state.

4 Experiments

The experiments evaluate Categorical Flow Maps across molecular graphs, binarised images, and text, emphasizing few-step quality, endpoint parametrisations, and test-time guidance. Results show strong low-NFE generation, competitive molecular quality, improved image metrics, and favorable low-NFE text performance.

  • Experimental scope: Experiments cover molecular graphs on QM9 and ZINC250k, binarised MNIST images, and Text8 and LM1B text.The molecular experiments compare against multi-step diffusion and flow methods as well as one-shot and few-step baselines.
  • Graph generation: At two steps on QM9, the sampler reaches 91–97% validity and 96–97% uniqueness while achieving competitive FCD.The comparison is made against strong multi-step flow-based baselines.
  • Graph generation: At 1 NFE on ZINC, ECLD achieves 93.5% validity versus 63.1% for MoFlow and 11.2% for PairFlow.At two steps, CSD and ECLD achieve FCD scores of 12.1–12.2, matching or improving multi-step baselines with a 50–500× reduction in function evaluations.
  • Graph generation: CSD and ECLD outperform unconstrained Flow Maps on validity and FCD, while ECLD favors validity and CSD achieves lower FCD at matched step counts.With up to 100 NFEs, CSD improves more steeply than Naive Flow Matching and ECLD, which level off earlier.
  • Image generation: Few-step MNIST sampling improves FID over Park et al. (2025), reaching 10.1 versus 12.9 at one step and 7.8 versus 8.5 at four steps.Flow-map lookahead provides the best class-conditional FID and highest target-class accuracy for reward-guided generation toward digit zeroes.
  • Text generation: On Text8, ECLD outperforms DFM in NLL across all NFEs, while on LM1B CFM outperforms UDLM in the low-NFE regime by up to 3.5× lower PPL.The LM1B comparison uses an identical evaluation pipeline for both models, although CFM has significantly lower entropy.

5 Related Works

Prior work on discrete generative modeling includes discrete and continuously parametrised diffusion, but stochastic discrete trajectories make consistency-style self-distillation and few-step inference difficult. Accelerated methods therefore motivate continuous formulations for categorical data.

  • Discrete diffusion models categorical distributions through discrete-time or continuous-time Markov chains.
  • Continuously parametrised approaches evolve discrete data through a continuous space while retaining discrete targets, including simplex- and metric-based constructions.
  • Accelerated diffusion and flow methods have demonstrated one- or few-step generation in continuous domains, raising whether similar acceleration is possible for discrete data.
  • At inference, independently sampling sequence elements can lead to low-quality generation, while existing approaches only partially remove the underlying stochasticity.

6 Conclusion

Categorical Flow Maps provide a self-distillable flow-based framework for accelerated discrete generation by transporting a continuous prior to a discrete target through simplex-constrained endpoint predictions. The formulation supports endpoint-consistency training and strong few-step generation across categorical tasks.

  • Categorical Flow Maps transport a continuous prior to a discrete target while constraining endpoint predictions to the probability simplex.
  • The simplex-constrained endpoint parametrisation yields a categorical cross-entropy learning signal and makes flow-map self-distillation applicable to discrete domains.
  • Endpoint-consistency self-distillation controls the Lagrangian residual directly through simplex-valued endpoint predictions.
  • Categorical Flow Maps achieve strong few-step generation across multiple categorical tasks and support test-time steering through adapted guidance and reweighting.
  • The paper develops flow-map self-distillation from flow-map jump, tangent, Lagrangian, Eulerian, and semigroup conditions.

B Proofs

The proofs establish how categorical endpoint predictions define constrained flow maps, how endpoint-consistency relates to the Lagrangian residual, and how training and sampling implement the resulting framework. The construction extends to categorical graphs through factorised node and edge predictions.

  • Under the linear endpoint-to-velocity decoder, endpoint-consistency controls the Lagrangian residual through a reverse-KL bound.
  • The endpoint predictor is a probability vector defining a categorical distribution over one-hot endpoints.
  • Endpoint parametrisation is practically necessary because the true vector field becomes degenerate near t = 1 when continuous prior mass must concentrate on discrete simplex vertices.
  • Simplex-constrained endpoint predictions geometrically confine flow-map states to line segments from the starting point to the data support.
  • Training combines variational flow matching with categorical or endpoint-consistency self-distillation, enforcing accurate endpoint predictions and trajectory self-consistency.
  • At inference, the learned flow map transports samples from the base distribution through successive time points, ending with argmax or categorical sampling.
  • For graphs, categorical node and edge variables are factorised, with separate node and edge losses combined using an edge-weight coefficient.

E.2 Additional results

Additional experiments compare CSD, ECLD, and Naive Flow Matching across higher evaluation budgets and Euler versus flow-map sampling. CSD benefits most from increased evaluations, while the preferable sampler depends on the loss.

  • CSD sample quality increases particularly strongly with higher NFE, whereas Naive Flow Matching and ECLD improve more slowly.
  • CSD significantly outperforms ECLD in FCD when samples use Euler integration with the learned instantaneous velocity field.
  • At higher NFEs, flow-map sampling generally outperforms Euler integration for Naive and CSD losses, but Euler integration achieves better FCD for ECLD.

E.3 Datasets

The molecular experiments use QM9 and Zinc250k, evaluate validity, uniqueness, and FCD, and compare CSD and ECLD against one-shot, discrete-flow, accelerated-inference, and diffusion baselines.

  • QM9: QM9 contains molecules with up to 9 heavy atoms, using 100K training, 20K validation, and 13K test molecules.The dataset has 4 atom types and 3 bond types.
  • Zinc250k: Zinc250k contains 249,455 molecules with up to 38 heavy atoms from 9 element types, split into 213,912 training and 23,768 validation molecules.
  • Metrics: Validity measures sanitizable RDKit molecules, uniqueness measures distinct largest-fragment SMILES among valid molecules, and FCD measures ChemNet activation distance from the data.
  • Baselines: CSD and ECLD are compared with non-diffusion one-shot, discrete flow-matching, few-step acceleration, and continuous diffusion baselines.The listed baselines include Set2GraphVAE, MoFlow, DeFoG, PairFlow, GDSS, GruM, and CatFlow.

E.6 Architecture

The molecular models use a graph-transformer architecture with separate node, edge, and global features, magnitude-preserving time conditioning, and learned loss weighting for flow-map training.

  • Backbone: The graph transformer processes noisy node, edge, and global feature tensors, zero-padding graphs below the maximum node count while masking padded nodes and edges.
  • Conditioning: Time conditioning embeds both the interval start s and duration ∆=t−s with magnitude-preserving sinusoidal embeddings and combines them through a magnitude-preserving sum.This replaces the DiGress backbone's MLP applied to global features.
  • Backbone: Each transformer layer uses 8 attention heads, ReLU MLPs, and feature-specific hidden dimensions for node, edge, and global transformations.The corresponding hidden dimensions are 256, 128, and 128.
  • Backbone: QM9 models use 9 transformer layers, while the ZINC model uses 12 layers.
  • Loss weighting: Learned uncertainty-based weighting improves flow-map models, whereas standard flow matching omits it because it learns only the instantaneous velocity field.The weighted losses use γs,t=(t−s)/(1−s), with wt≡1.
  • Training: Training uses AdamW, gradient clipping, exponential moving averages, cosine learning-rate schedules, and dataset-specific batch sizes and epoch counts.Models were trained on RTX A6000 GPUs for approximately 2 days on QM9 and 6 days on ZINC.

F Image datasets

The image experiments replace a detrimental softmax-based continuous-input embedding with an RMS-normalized, noise-level-stable representation conditioned through residual and FiLM components.

  • Embedding layer: Applying softmax to continuous inputs in an embedding layer was found particularly detrimental to training.A simple alternative uses a linear projection without softmax and normalizes its variance.
  • Results: Large improvements were observed, especially on the lower-dimensional Text8 dataset, after approximately normalizing the embedding variance.
  • Embedding layer: The proposed image embedding uses an RMS-normalized projection of noisy continuous inputs, a shallow residual MLP, and FiLM conditioning.The design targets noise-level-stable representations for DiT models.
  • Training: Training uses AdamW, warmup, gradient clipping, bfloat16 mixed precision, dropout, and dataset-specific batch sizes with gradient accumulation for Text8.

G.3 Qualitative samples

The paper reports qualitative single-step samples for molecular graphs and LM1B text, using uncurated or non-cherry-picked examples to illustrate generated outputs.

  • Molecular samples: The molecular experiments include uncurated qualitative samples generated with a single function evaluation using ECLD-trained models.These samples are shown in Figures 9 and 10.
  • Text samples: The displayed LM1B samples contain packed text segments separated by [CLS] markers.
  • Text samples: The LM1B model provides non-cherry-picked qualitative samples generated with a single sampling step.The [CLS] strings delimit packed sequence segments.
Loading 2602.12233v1…