Discrete MeanFlow: One-Step Generation via Conditional Transition Kernels

2026-05-12

AI Summary

Existing continuous flow matching methods are ill-suited for single-step generation in discrete state spaces due to the absence of smooth trajectories and spatial derivatives. This work proposes a discrete MeanFlow identity that models probability mass transport via the conditional transition kernel of a continuous-time Markov chain, replacing the chain rule with the Kolmogorov forward equation to define an average discrete velocity. A boundary-aware, self-constructing neural parameterization is further introduced to enforce probability conservation and exact boundary conditions. The method accurately recovers analytical solutions on finite-state Markov chains and demonstrates effective single-step generation across synthetic tasks with varying alphabet sizes and sequence lengths.

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.

Problem

• discrete state spaces
• one-step generation
• continuous-time Markov chain
• transition kernel
• flow matching

Innovation

• Discrete MeanFlow
• conditional transition kernel
• continuous-time Markov chain
• one-step generation
• boundary-by-construction