# Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport

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

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 76 references that are in the map.
- First posted on arXiv 8 Feb 2024.

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Training Generative Adversarial Networks with Limited Data](https://mlsanity.com/paper/2006.06676/) (2020)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Pixel Recurrent Neural Networks](https://mlsanity.com/paper/1601.06759/) (2016)
- [Which Training Methods for GANs do actually Converge?](https://mlsanity.com/paper/1801.04406/) (2018)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Score-based Generative Modeling in Latent Space](https://mlsanity.com/paper/2106.05931/) (2021)
- [Tackling the Generative Learning Trilemma with Denoising Diffusion GANs](https://mlsanity.com/paper/2112.07804/) (2021)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Score-Based Generative Modeling with Critically-Damped Langevin Diffusion](https://mlsanity.com/paper/2112.07068/) (2021)
- [Adversarial Latent Autoencoders](https://mlsanity.com/paper/2004.04467/) (2020)
- [StyleSwin: Transformer-based GAN for High-resolution Image Generation](https://mlsanity.com/paper/2112.10762/) (2021)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Learning Energy-Based Models by Diffusion Recovery Likelihood](https://mlsanity.com/paper/2012.08125/) (2020)
- [Stabilizing Training of Generative Adversarial Networks through Regularization](https://mlsanity.com/paper/1705.09367/) (2017)
- [VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Models](https://mlsanity.com/paper/2010.00654/) (2020)
- [AutoGAN: Neural Architecture Search for Generative Adversarial Networks](https://mlsanity.com/paper/1908.03835/) (2019)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)

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

- [Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling](https://mlsanity.com/paper/2504.10612/) (2025)
- [One-Step Generative Modeling via Wasserstein Gradient Flows](https://mlsanity.com/paper/2605.11755/) (2026)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [Gradient Flow Drifting: Generative Modeling via Wasserstein Gradient Flows of KDE-Approximated Divergences](https://mlsanity.com/paper/2603.10592/) (2026)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Improving Neural Optimal Transport via Displacement Interpolation](https://mlsanity.com/paper/2410.03783/) (2024)
- [Overcoming Spurious Solutions in Semi-Dual Neural Optimal Transport: A Smoothing Approach for Learning the Optimal Transport Plan](https://mlsanity.com/paper/2502.04583/) (2025)
- [Are We Really Learning the Score Function? Reinterpreting Diffusion Models Through Wasserstein Gradient Flow Matching](https://mlsanity.com/paper/2509.00336/) (2025)
- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)
- [Weighted Conditional Flow Matching](https://mlsanity.com/paper/2507.22270/) (2025)
- [Towards Understanding Gradient Dynamics of the Sliced-Wasserstein Distance via Critical Point Analysis](https://mlsanity.com/paper/2502.06525/) (2025)
- [QUATRO: Query-Adaptive Trust Region Policy Optimization for LLM Fine-tuning](https://mlsanity.com/paper/2602.04620/) (2026)
- [Learning Discrete Diffusion of Graphs via Free-Energy Gradient Flows](https://mlsanity.com/paper/2604.11311/) (2026)
- [Rate-Optimal Noise Annealing in Semi-Dual Neural Optimal Transport: Tangential Identifiability, Off-Manifold Ambiguity, and Guaranteed Recovery](https://mlsanity.com/paper/2602.04110/) (2026)
- [Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow](https://mlsanity.com/paper/2608.29647/) (2026)
- [A Differential Equation Approach for Wasserstein GANs and Beyond](https://mlsanity.com/paper/2405.16351/) (2024)
- [Guiding Time-Varying Generative Models with Natural Gradients on Exponential Family Manifold](https://mlsanity.com/paper/2502.07650/) (2025)
- [Helping CLIP See Both the Forest and the Trees: A Decomposition and Description Approach](https://mlsanity.com/paper/2507.03458/) (2025)
- [Rethinking the Flow-Based Gradual Domain Adaptation: A Semi-Dual Optimal Transport Perspective](https://mlsanity.com/paper/2602.01179/) (2026)
- [Analyzing and Improving Diffusion Models for Time-Series Data Imputation: A Proximal Recursion Perspective](https://mlsanity.com/paper/2602.01182/) (2026)
- [UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems](https://mlsanity.com/paper/2605.21094/) (2026)
- [One-Step Generative Modeling via Unbalanced Optimal Transport](https://mlsanity.com/paper/2609.32708/) (2026)
- [Unbalanced W-Flow: Unbalanced Wasserstein Drifting under Finite-Batch Marginal Fluctuations](https://mlsanity.com/paper/2610.18552/) (2026)

## Abstract (arXiv)

Wasserstein Gradient Flow (WGF) describes the gradient dynamics of probability density within the Wasserstein space. WGF provides a promising approach for conducting optimization over the probability distributions. Numerically approximating the continuous WGF requires the time discretization method. The most well-known method for this is the JKO scheme. In this regard, previous WGF models employ the JKO scheme and parametrize transport map for each JKO step. However, this approach results in quadratic training complexity O(K2) with the number of JKO step K. This severely limits the scalability of WGF models. In this paper, we introduce a scalable WGF-based generative model, called Semi-dual JKO (S-JKO). Our model is based on the semi-dual form of the JKO step, derived from the equivalence between the JKO step and the Unbalanced Optimal Transport. Our approach reduces the training complexity to O(K). We demonstrate that our model significantly outperforms existing WGF-based generative models, achieving FID scores of 2.62 on CIFAR-10 and 5.46 on CelebA-HQ-256, which are comparable to state-of-the-art image generative models.