Categorical Flow Maps

2026-02-12Machine Learning

Machine Learning
AI summary

The authors present Categorical Flow Maps, a new method to quickly generate categorical data, like text or images, in just a few steps. Their approach creates smooth paths that move probabilities toward a target, which helps models make better predictions. This method also works well with existing training tricks and lets users guide the output during generation. Experiments show it performs very well on several types of data, including images, molecular graphs, and text, even when generating everything in one step.

Categorical dataFlow matchingSelf-distillationProbability simplexContinuous trajectoriesDiffusion modelsAccelerated inferenceEndpoint consistencyGuidance reweightingMolecular graphs
Authors
Daan Roos, Oscar Davis, Floor Eijkelboom, Michael Bronstein, Max Welling, İsmail İlkan Ceylan, Luca Ambrogioni, Jan-Willem van de Meent
Abstract
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.