One-Step Generation via Riemannian Wasserstein Gradient Flows

David Li; Chanhyuk Lee; Jaehoon Yoo; Nikita Gushchin; Floor Eijkelboom; Eric Moulines; Maxim Panov; Alexander Korotin

cs.LGOct 3, 2026

One-Step Generation via Riemannian Wasserstein Gradient Flows

Authors: David Li, Chanhyuk Lee, Jaehoon Yoo, Nikita Gushchin, Floor Eijkelboom, Eric Moulines, Maxim Panov, Alexander Korotin,

Organizations: Mohamed bin Zayed University of Artificial Intelligence · Korea Advanced Institute of Science and Technology (KAIST) · Applied AI Institute · AXXX ·

Abstract

Recently, Drifting Models and Wasserstein Gradient Flows have attracted substantial attention because they move iterative distributional refinement to training and amortize it into a generator, enabling fast inference. However, existing formulations have been developed largely for continuous Euclidean domains, such as image spaces, where particles admit unconstrained additive updates. On constrained spaces, these updates can leave the valid domain or ignore its geometry, making them unsuitable targets for training. Recent work has adapted updates to these spaces, but has focused on particular fields or offered limited empirical comparison. We derive and compare several geometry-aware fields within a common training framework for one-step generators. We test the method on data with different structures and obtain competitive one-step results in each setting. The best-performing field varies by task, showing why the choice of objective matters in practice.

Figures & tables

Figure 1: Fixed-

z

generation on a product of simplices. Outputs lie in

\mathcal{M}=(\Delta^{2})^{3}

, with one simplex per word position. Each color tracks

f_{\theta_{k}}(z_{i})

for a fixed

z_{i}

\theta_{k}

is updated, ending at Cats are cute , Penguins wear pajamas , and I am alive .

Manifold $\mathcal{M}$	$\mathcal{S}_{\varepsilon}(q,p)$	$\tfrac{1}{2}\operatorname{MMD}_{\kappa}^{2}(q,p)$	$\operatorname{KL}(q\,\\|\,p)$
Euclidean: $\mathbb{R}^{D}$	$\nabla_{x}\!\left(f_{q,q}^{\varepsilon}-f_{q,p}^{\varepsilon}\right)$	$\nabla_{x}\!\left(m_{p}^{\kappa}-m_{q}^{\kappa}\right)$	$\nabla_{x}\log\!\frac{p(x)}{q(x)}$
Round sphere: $\mathbb{S}^{D-1}$	$P_{x}\nabla_{x}\!\left(f_{q,q}^{\varepsilon}-f_{q,p}^{\varepsilon}\right)$	$P_{x}\nabla_{x}\!\left(m_{p}^{\kappa}-m_{q}^{\kappa}\right)$	$P_{x}\nabla_{x}\log\!\frac{p(x)}{q(x)}$
Product simplex: $(\Delta_{\circ}^{d-1})^{m}$	$G_{x}\nabla_{x}\!\left(f_{q,q}^{\varepsilon}-f_{q,p}^{\varepsilon}\right)$	$G_{x}\nabla_{x}\!\left(m_{p}^{\kappa}-m_{q}^{\kappa}\right)$	$G_{x}\nabla_{x}\log\!\frac{p(x)}{q(x)}$

Table 1: Intrinsic drift

V_{p,q}(x)

for each objective–geometry pair considered in this work. Definitions and drift derivations for Sinkhorn, MMD, and KL are given in Appendices D.4 , D.5 , and D.6 .

Figure 2: Left: R-WGF training algorithm. Right: Visualization of training dynamics of $q$ on $\mathcal{M}=\mathbb{S}^{2}$ toward a spherical target $p$ .

Dataset	R-WGF-MMD ( Ours )	R-WGF-KL ( Ours )	R-WGF-Sinkhorn ( Ours )	KGDM
Earthquake	$\mathbf{0.032\pm 0.0068}$	$0.050\pm 0.0015$	$0.041\pm 0.0048$	$0.044\pm 0.0033$
Volcano	$0.13\pm 0.011$	$\mathbf{0.099\pm 0.024}$	$0.12\pm 0.047$	$0.14\pm 0.035$
Fire	$0.036\pm 0.012$	$\mathbf{0.022\pm 0.0067}$	$0.031\pm 0.0034$	$0.034\pm 0.014$
Flood	$0.060\pm 0.014$	$0.074\pm 0.013$	$\mathbf{0.046\pm 0.011}$	$0.054\pm 0.0085$

Table 2: Held-out MMD on the Earth benchmark (

\downarrow

), with the three best results highlighted

Binary-MNIST

Method NFE FID

\downarrow

PairFlow ( Park et al., 2026 ) 4 8.50 PairFlow + ReDi ( Park et al., 2026 ) 4 8.54 CFM ( Roos et al., 2026 ) 4 7.8 Euclidean Drifting ( Eijkelboom et al., 2026 ) 1

29.0\pm 1.2

CatDrift ( Eijkelboom et al., 2026 ) 1

8.0\pm 0.1

R-WGF-MMD ( Ours ) 1

11.2\pm 0.4

R-WGF-KL ( Ours ) 1

\mathbf{6.5\pm 0.3}

R-WGF-Sinkhorn ( Ours ) 1

7.0\pm 0.2

Table 3: Discretized image generation FID (

\downarrow

) and IS (

\uparrow

), with the three best results per column highlighted. Full sample grids are in Figs. 3 and 5 .

Panel of Table 4 — Table 4: One-step LM1B results, with the three best results highlighted. Full generation examples are in Appendix B.4.1 .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	LR	$T_{\mathrm{KGDM}}$	$\varepsilon_{\mathrm{OT}}$	$\gamma_{\mathrm{MMD}}/\gamma_{\mathrm{KL}}$
Earthquake	$5.17\!\times\!10^{-4}$	.7178	.0322	2/2
Volcano	$5.24\!\times\!10^{-4}$	.3250	.0264	1/1
Fire	$9.70\!\times\!10^{-4}$	.6332	.2005	1/1
Flood	$5.52\!\times\!10^{-4}$	.1291	.0333	4/4

Appendix

Table 5: Final Earth settings.

Method	$\epsilon_{\mathrm{data}}$	Scale
R-WGF-Sinkhorn	0	$\varepsilon_{\mathrm{OT}}=.005$
R-WGF-KL	$.001$	$\rho=.05$
R-WGF-MMD	$.001$	$\rho=.05$

Appendix

Table 6: Binary-MNIST objective implementations.

Panel of Figure 3 — Figure 3: Full Binary-MNIST sample grids. Fixed, uncurated $8\times 8$ argmax grids from the selected EMA checkpoints.

Panel of Figure 4 — Figure 4: Discretized CIFAR-10 pilot runs. FID-5k (left) and IS (right) of R-WGF-Sinkhorn, R-WGF-KL, and R-WGF-MMD trained from scratch with the same recipe, evaluated with one NFE and argmax decoding.

Figure 5: Full discretized CIFAR-10 sample grid. Fixed, uncurated $8\times 8$ argmax grid from the selected R-WGF-Sinkhorn EMA checkpoint with one NFE.

Figure 6: LM1B objective screening. Labels mark the lowest observed mean GenPPL.

One-Step Generative Modeling

One-step generative modeling maps noise or conditions to data in one network evaluation, replacing iterative denoising or numerical integration.

See all 146 papers

Discrete Wasserstein Flows for One-Step Generative Modeling
Oct 1, 2026Alessandro Micheli, Andrea Zerio, Samir BhattWasserstein Gradient Flows Markov Chains
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
One-Step Generative Modeling via Wasserstein Gradient Flows
May 12, 2026Jiaqi Han, Puheng Li, Qiushan Guo +3Wasserstein Gradient Flows Image Generation
Stanford University · ByteDance
Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow
Aug 30, 2026Hoseong Hwang, Woorim Han, Joungin Chun +2Image Generation One-Step Generative Modeling
Sungkyunkwan University, Seoul, Republic of Korea · Korea Institute for Advanced Study, Seoul, Republic of Korea

Oct 1, 2026J/K move · Enter open · S save

Discrete Wasserstein Flows for One-Step Generative Modeling

Alessandro Micheli, Andrea Zerio, Samir Bhatt

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

Wasserstein Gradient Flows Markov Chains One-Step Generative Modeling Generative Modeling Flow-Based Generative Modeling

Open paper AskPDF arXiv

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.

Abstract

One-Step Generative Modeling

Explore similar work

Discrete Wasserstein Flows for One-Step Generative Modeling

One-Step Generative Modeling via Wasserstein Gradient Flows

Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow