# One-Step Generative Modeling via Wasserstein Gradient Flows

arXiv: https://arxiv.org/abs/2605.11755 · HTML: https://mlsanity.com/paper/2605.11755/

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 56 of its 70 references that are in the map.
- First posted on arXiv 12 May 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 56)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://mlsanity.com/paper/1809.11096/) (2018)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://mlsanity.com/paper/1511.06434/) (2015)
- [ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation](https://mlsanity.com/paper/2305.16213/) (2023)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction](https://mlsanity.com/paper/2404.02905/) (2024)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [Improved Distribution Matching Distillation for Fast Image Synthesis](https://mlsanity.com/paper/2405.14867/) (2024)
- [Autoregressive Image Generation without Vector Quantization](https://mlsanity.com/paper/2406.11838/) (2024)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets](https://mlsanity.com/paper/2202.00273/) (2022)
- [Scaling up GANs for Text-to-Image Synthesis](https://mlsanity.com/paper/2303.05511/) (2023)
- [Improved Techniques for Training Consistency Models](https://mlsanity.com/paper/2310.14189/) (2023)
- [Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models](https://mlsanity.com/paper/2501.01423/) (2025)

## Built on by (most cited first, 24 of 24)

- [Representation Distribution Matching for One-Step Visual Generation](https://mlsanity.com/paper/2607.02375/) (2026)
- [Three-Body Scattering for Generative Modeling](https://mlsanity.com/paper/2607.18198/) (2026)
- [Difficulty-Calibrated Interpolation Paths for Conditional Flow Matching](https://mlsanity.com/paper/2608.21286/) (2026)
- [Unifying Distributional Training for One-Step Visual Generation](https://mlsanity.com/paper/2609.35763/) (2026)
- [Improved Distributional Diffusion Models](https://mlsanity.com/paper/2609.37147/) (2026)
- [Condition-Wise Sinkhorn Drifting for One-Shot Learned Channel Simulation](https://mlsanity.com/paper/2606.17893/) (2026)
- [DriftSE: Speech Enhancement with Generative Drifting](https://mlsanity.com/paper/2609.12252/) (2026)
- [Stochastic Gradient Descent over P2](https://mlsanity.com/paper/2609.13343/) (2026)
- [FlowTSFM: Turning Encoder Depth into Quantile Transport](https://mlsanity.com/paper/2609.13640/) (2026)
- [One-Step Generative Surrogate Models via Block-Triangular Joint Drifting](https://mlsanity.com/paper/2609.26435/) (2026)
- [ViRDM: Taming Representation Distribution Matching for Few-Step Causal Video Generation](https://mlsanity.com/paper/2609.28923/) (2026)
- [Aligning One-Step Generative Models with Reward-Weighted Transport Distillation](https://mlsanity.com/paper/2609.30840/) (2026)
- [An End-to-End Latent-Rollout Approach for Pushing Few-Step ImageNet-256 Generation to FID 1.11 without Fr'echet Losses](https://mlsanity.com/paper/2609.32376/) (2026)
- [One-Step Generative Modeling via Unbalanced Optimal Transport](https://mlsanity.com/paper/2609.32708/) (2026)
- [RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows](https://mlsanity.com/paper/2609.39271/) (2026)
- [PMosFM: Preconditioned Manifold Matching for One-Step Physics-Constrained Generation](https://mlsanity.com/paper/2609.40287/) (2026)
- [One-Step Generative Modeling via Training Dynamics Action](https://mlsanity.com/paper/2610.00518/) (2026)
- [One-Step Generation via Riemannian Wasserstein Gradient Flows](https://mlsanity.com/paper/2610.04454/) (2026)
- [Learning Discriminative Geometry for Drifting Models](https://mlsanity.com/paper/2610.04703/) (2026)
- [Multi-Bandwidth Distribution Matching Distillation: On the Equivalence of Distribution Matching Distillation and Drifting Models](https://mlsanity.com/paper/2610.10989/) (2026)
- [Amortised Inference through One-Step Implicit Sampling](https://mlsanity.com/paper/2610.12679/) (2026)
- [Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning](https://mlsanity.com/paper/2610.15662/) (2026)
- [Unbalanced W-Flow: Unbalanced Wasserstein Drifting under Finite-Batch Marginal Fluctuations](https://mlsanity.com/paper/2610.18552/) (2026)
- [Are We Overthinking One-Step Generation?](https://mlsanity.com/paper/2610.20331/) (2026)

## Abstract (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×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.