Discrete MeanFlow: One-Step Generation via Conditional Transition Kernels

TL;DR

Discrete MeanFlow introduces a novel approach for one-step generation in discrete spaces by learning a transition kernel that directly models probability mass transport, eliminating the need for iterative refinement.

Contribution

It develops the Discrete MeanFlow framework, including a theoretical identity and a boundary-structured kernel parameterization for efficient, exact probability-based generation in discrete state spaces.

Findings

Successfully recovers ground truth transition kernels in finite-state Markov chains.

Achieves high-precision results on synthetic sequence generation tasks.

Enables one-step generation without iterative processes or ODE integration.

Abstract

MeanFlow enables one-step generation in continuous spaces by learning an average velocity over a time interval rather than the instantaneous velocity field of flow matching. However, discrete state spaces do not have smooth trajectories or spatial derivatives, so the continuous formulation does not directly apply. We introduce Discrete MeanFlow, which replaces the motion of a point with the transport of probability mass over finite states. Our key object is the conditional transition kernel of a continuous-time Markov chain (CTMC), from which we define a mean discrete rate that measures the average change in transition probability over a time interval. We prove a Discrete MeanFlow identity that relates this finite-interval rate to the instantaneous CTMC generator at the endpoint, with the Kolmogorov forward equation replacing the spatial chain rule of continuous MeanFlow. Based on this identity, we parameterize the transition kernel directly using a boundary-by-construction design that guarantees valid probability outputs and exact boundary conditions without auxiliary losses. Since the learned kernel is itself a probability distribution, generation reduces to a single forward pass followed by one categorical draw meaning no iterative denoising, ODE integration, or multi-step refinement is required. We validate the framework on exact finite-state Markov chains, where the learned kernel recovers the analytical ground truth to high precision, and on factorized synthetic sequence generation tasks with varying alphabet sizes and sequence lengths.