Discrete Wasserstein Flows for One-Step Generative Modeling

Alessandro Micheli; Andrea Zerio; Samir Bhatt

cs.LGOct 1, 2026

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

Abstract

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\times 32$ torus . First 1600 updates. Top: Relative density $\rho_{q_{t}}=|\mathcal{X}|q_{t}$ (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: $\mathrm{KL}(q_{t}\,\|\,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(\rho_{q}/\rho_{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

\eta_{\mathrm{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

\eta N

fixed; medians over five seeds, band spanning seeds at

\eta=0.8

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Setting
Model	One unconstrained logit per state; $q_{g}=\operatorname{softmax}(g)$
Initialization	Exact logits representing $q_{0}$
Drift evaluation	All $1024$ states and directed nearest-neighbour edges
Outer updates	$6400$
Flow step	$\eta=0.8$
Final flow time	$5120$

Appendix

Table 1: Configuration of the exact tabular validation.

Flow time $t$	$\mathrm{KL}(q_{t}\,\\|\,p)$	Step-size bound	Fitted slope	Law-identity residual
$0$	$1.85$	$3.72$	$2.0014$	$3.7\times 10^{-16}$
$102.4$	$3.08\times 10^{-1}$	$7.19$	$2.0001$	$1.2\times 10^{-16}$
$256.0$	$7.47\times 10^{-2}$	$16.06$	$2.0000$	$9.3\times 10^{-17}$
$576.0$	$1.20\times 10^{-2}$	$38.20$	$2.0001$	$6.9\times 10^{-17}$

Appendix

Table 2: Frozen-transition order test at four states of the tabular trajectory. The step-size bound is the largest

\eta

admitted by equation 29 , the fitted slope is that of

\|q_{g+\eta\Delta_{g}}-q_{g}P_{\eta,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\sim\mathcal{N}(0,I_{16})$
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 $q_{0}$ , 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\pm 1.09)\times 10^{-3}$ ).

Explore similar work

One-Step Generative Modeling via Wasserstein Gradient Flows
May 12, 2026Jiaqi Han, Puheng Li, Qiushan Guo +3Diffusion Models Generative Models
Stanford University · ByteDance
On the Wasserstein Gradient Flow Interpretation of Drifting Models
May 6, 2026Arthur Gretton, Li Kevin Wenliang, Alexandre Galashov +3Wasserstein Gradient Flows Generative Models
Google DeepMind
Discrete MeanFlow: One-Step Generation via Conditional Transition Kernels
May 12, 2026Fairoz Nower Khan, Nabuat Zaman Nahim, Md Sajid Ahmed +2Meanflow Generative Flow Networks
Department of Computer Science, University of Kentucky Lexington, Kentucky 40506, USA

May 12, 2026, cs.LGJ/K move · Enter open · S save

One-Step Generative Modeling via Wasserstein Gradient Flows

Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu+2

Stanford University · ByteDance

Diffusion Models Generative Models Wasserstein Distance Imagenet Gradient

Open paper AskPDF arXiv

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

\times

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

\times

faster sampling. These results show that Wasserstein gradient flows provide a principled and effective foundation for fast and high-fidelity generative modeling.

Abstract

Explore similar work

One-Step Generative Modeling via Wasserstein Gradient Flows

On the Wasserstein Gradient Flow Interpretation of Drifting Models

Discrete MeanFlow: One-Step Generation via Conditional Transition Kernels