Discrete Wasserstein Flows for One-Step Generative Modeling
Authors: Alessandro Micheli, Andrea Zerio, Samir Bhatt
Organizations: Imperial College London London, United Kingdom · Department of Computer Science, Aalborg University, Copenhagen, Denmark; Centre for Frontier AI Research (CFAR), Institute of Advanced Intelligence and Computing (IAIC), A*STAR, Singapore · University of Copenhagen Copenhagen, Denmark
We introduce a new framework for one-step generative modelling on finite state spaces. To extend drifting beyond continuous domains, we use discrete Wasserstein geometry to define a target-relative KL gradient flow over the transitions of a reversible Markov kernel. We realize this probability flow at the particle level through Markov jumps and amortize the resulting transport updates into a latent-conditioned generator, so that the iterative dynamics are required only during training while inference remains one-step. In a controlled setting where the underlying distributions and transport dynamics can be computed exactly, we verify KL dissipation, consistency between the particle dynamics and the probability flow, and the predicted numerical scaling. We further show that a finite-capacity neural generator can track these exact transport targets while retaining one-step generation. These results validate the basic construction and provide a foundation for scaling Discrete Drifting to structured discrete data.
Figures & tables
Figure 1: Exact tabular flow on the 32×32 torus . First 1600 updates. Top: Relative density ρqt=∣X∣qt (blue) at five stages; orange contours show the fixed target density. The model begins as a single bump separated from the target, transports mass across low-density regions, and converges to the three modes in the correct proportions. Bottom: KL(qt∥p) versus flow time on a logarithmic scale. Numbered markers associate each displayed distribution with its position along the trajectory.
Figure 2: Edgewise decomposition of the relative-entropy current at initialization . Left: Target-attraction current over the target density. Center: Model-repulsion current over the model density; by the logarithmic-mean identity, this component coincides with the linear diffusion current. Right: Net current over log(ρq/ρp) , transporting mass from model surplus toward target deficit.
Figure 3: Left: one frozen transition of equation 28 against the two updates it replaces: fitted slope 2.00 for the residual of equation 31 , and machine precision for the Euler step of the law, which equation 30 makes an identity. The vertical rule marks ηCFL , the largest step size admitted by equation 29 , past which no particle can make the move. Right: final gap to the conditional oracle against batch size b , at four step sizes with ηN fixed; medians over five seeds, band spanning seeds at η=0.8 .
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Component
Setting
Model
One unconstrained logit per state; qg=softmax(g)
Initialization
Exact logits representing q0
Drift evaluation
All 1024 states and directed nearest-neighbour edges
Outer updates
6400
Flow step
η=0.8
Final flow time
5120
Appendix
Table 1: Configuration of the exact tabular validation.
Flow time t
KL(qt∥p)
Step-size bound
Fitted slope
Law-identity residual
0
1.85
3.72
2.0014
3.7×10−16
102.4
3.08×10−1
7.19
2.0001
1.2×10−16
256.0
7.47×10−2
16.06
2.0000
9.3×10−17
576.0
1.20×10−2
38.20
2.0001
6.9×10−17
Appendix
Table 2: Frozen-transition order test at four states of the tabular trajectory. The step-size bound is the largest η admitted by equation 29 , the fitted slope is that of ∥qg+ηΔg−qgPη,g∥1 , and the last column is the largest observed violation of equation 30 over the whole sweep. The bound grows as the flow converges and the rates decay.
Component
Setting
Latent distribution
z∼N(0,I16)
Generator
Three hidden layers of width 128 , with SiLU activations
Output
1024 -way softmax
Training latent bank
B=256 , fixed for the complete run
Held-out latent bank
256 , fixed and never optimized against
Initialization
Statewise output bias calibrated to q0 , then frozen
Appendix
Table 3: Configuration of the finite neural amortization experiment.
Figure 4: End-to-end diagnostics for the exact tabular flow. Top left: relative KL for the stop-gradient chart update, probability-space Euler reference, sampled state–edge implementation, and differentiate-through-target control. Top right: the corresponding total-variation trajectories. Bottom left: predicted instantaneous dissipation and measured finite-difference KL decrease. Bottom right: numerical residuals for four structural identities. The intended chart update, probability-space reference, and sampled implementation agree throughout the resolved descent, whereas the differentiate-through-target control does not.
Figure 5: Final neural marginals across five seeds. The panels show the target, the mean final marginal, and one representative seed. Orange contours mark the fixed target-density levels. Every seed recovers the three modes and their unequal mass allocation, with limited variation across runs (final total variation (5.12±1.09)×10−3 ).
Diffusion models and flow-based methods have shown impressive generative capability, especially for images, but their sampling is expensive because it requires many iterative updates. We introduce W-Flow, a framework for training a generator that transforms samples from a simple reference distribution into samples from a target data distribution in a single step. This is achieved in two steps: we first define an evolution from the reference distribution to the target distribution through a Wasserstein gradient flow that minimizes an energy functional; second, we train a static neural generator to compress this evolution into one-step generation. We instantiate the energy functional with the Sinkhorn divergence, which yields an efficient optimal-transport-based update rule that captures global distributional discrepancy and improves coverage of the target distribution. We further prove that the finite-sample training dynamics converge to the continuous-time distributional dynamics under suitable assumptions. Empirically, W-Flow sets a new state of the art for one-step ImageNet 256×256 generation, achieving 1.29 FID, with improved mode coverage and domain transfer. Compared to multi-step diffusion models with similar FID scores, our method yields approximately 100× faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.
Recently, Deng et al. (2026) proposed Generative Modeling via Drifting (GMD), a novel framework for generative tasks. This note presents an analysis of GMD through the lens of Wasserstein Gradient Flows (WGF), i.e., the path of steepest descent for a functional in the space of probability measures, equipped with the geometry of optimal transport. Unlike previous WGF-based contributions, GMD can be thought of as directly targeting a fixed point of a specific WGF flow. We demonstrate three main results: first, that one algorithm proposed by Deng et al. (2026) corresponds to finding the limiting point of a WGF on the KL divergence, with Parzen smoothing on the densities. Second, that the algorithm actually implemented by Deng et al. (2026) corresponds to a different procedure, which bears some resemblance to the fixed point of a WGF on the Sinkhorn divergence, but lacks certain desirable properties of the latter. Third, the same same idea can be extended to the limiting point of other WGFs, including the Maximum Mean Discrepancy (MMD), the sliced Wasserstein distance, and GAN critic functions.
Arthur Gretton, Li Kevin Wenliang, Alexandre Galashov +3
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.
Fairoz Nower Khan, Nabuat Zaman Nahim, Md Sajid Ahmed +2
Department of Computer Science, University of Kentucky Lexington, Kentucky 40506, USA