mlsanityData to 7 Oct 2026

Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport

arXiv 2402.05443 ↗First posted 8 Feb 2024cs.LG cs.CVcited by 22

Not on the map. Why?

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.

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

Builds on

56 of 76 references.

References · 56

  1. 2020Denoising Diffusion Probabilistic Modelscited by 20,851
  2. 2020Denoising Diffusion Implicit Modelscited by 9,655
  3. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 8,964
  4. 2022Generative Adversarial Networkscited by 5,603
  5. 2022Flow Matching for Generative Modelingcited by 5,426
  6. 2019Generative Modeling by Estimating Gradients of the Data Distributioncited by 4,067
  7. 2014Deep Learning Face Attributes in the Wildcited by 3,679
  8. 2020Taming Transformers for High-Resolution Image Synthesiscited by 3,229
  9. 2017Progressive Growing of GANs for Improved Quality, Stability, and Variationcited by 2,967
  10. 2017Improved Training of Wasserstein GANscited by 2,223
  11. 2016Density estimation using Real NVPcited by 2,029
  12. 2018Glow: Generative Flow with Invertible 1x1 Convolutionscited by 1,539
38 more
  1. 2020Training Generative Adversarial Networks with Limited Datacited by 1,184
  2. 2023Improving and generalizing flow-based generative models with minibatch optimal transportcited by 835
  3. 2016Pixel Recurrent Neural Networkscited by 814
  4. 2018Which Training Methods for GANs do actually Converge?cited by 672
  5. 2020NVAE: A Deep Hierarchical Variational Autoencodercited by 635
  6. 2021Score-based Generative Modeling in Latent Spacecited by 634
  7. 2021Tackling the Generative Learning Trilemma with Denoising Diffusion GANscited by 515
  8. 2016Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithmcited by 498
  9. 2016Input Convex Neural Networkscited by 407
  10. 2021Score-Based Generative Modeling with Critically-Damped Langevin Diffusioncited by 230
  11. 2020Adversarial Latent Autoencoderscited by 155
  12. 2021StyleSwin: Transformer-based GAN for High-resolution Image Generationcited by 152
  13. 2015Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measurescited by 138
  14. 2020Learning Energy-Based Models by Diffusion Recovery Likelihoodcited by 122
  15. 2017Stabilizing Training of Generative Adversarial Networks through Regularizationcited by 118
  16. 2020VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Modelscited by 103
  17. 2019AutoGAN: Neural Architecture Search for Generative Adversarial Networkscited by 94
  18. 2021Proximal Optimal Transport Modeling of Population Dynamicscited by 91
  19. 2020Robust Optimal Transport with Applications in Generative Modeling and Domain Adaptationcited by 86
  20. 2021Generative Modeling with Optimal Transport Mapscited by 80
  21. 2021Large-Scale Wasserstein Gradient Flowscited by 76
  22. 2022Subspace Diffusion Generative Modelscited by 70
  23. 2020A Contrastive Learning Approach for Training Variational Autoencoder Priorscited by 65
  24. 2020SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergencecited by 64
  25. 2020Dual Contradistinctive Generative Autoencodercited by 58
  26. 2021Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networkscited by 53
  27. 2021Variational Wasserstein gradient flowcited by 53
  28. 2019Sinkhorn Divergences for Unbalanced Optimal Transportcited by 52
  29. 2020Refining Deep Generative Models via Discriminator Gradient Flowcited by 45
  30. 2022Normalizing flow neural networks by JKO schemecited by 39
  31. 2023Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transportcited by 39
  32. 2021KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Supportcited by 37
  33. 2023Convergence of flow-based generative models via proximal gradient descent in Wasserstein spacecited by 34
  34. 2019Deep Generative Learning via Variational Gradient Flowcited by 32
  35. 2021Efficient Gradient Flows in Sliced-Wasserstein Spacecited by 21
  36. 2023Learning Energy-Based Models by Cooperative Diffusion Recovery Likelihoodcited by 19
  37. 2022Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Schemecited by 17
  38. 2021Wasserstein Proximal of GANscited by 15

50 most cited of 56 All 56 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 2; 2025: 10; 2026: 10.

22 papers.

Sort:

By area:

  1. 2025Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modelingcited by 25
  2. 2025Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approachcited by 18
  3. 2026One-Step Generative Modeling via Wasserstein Gradient Flowscited by 18
  4. 2026Gradient Flow Drifting: Generative Modeling via Wasserstein Gradient Flows of KDE-Approximated Divergencescited by 15
  5. 2024Improving Neural Optimal Transport via Displacement Interpolationcited by 7
  6. 2025Overcoming Spurious Solutions in Semi-Dual Neural Optimal Transport: A Smoothing Approach for Learning the Optimal Transport Plancited by 6
  7. 2025Learning of Population Dynamics: Inverse Optimization Meets JKO Schemecited by 6
  8. 2025Are We Really Learning the Score Function? Reinterpreting Diffusion Models Through Wasserstein Gradient Flow Matchingcited by 5
  9. 2025DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flowscited by 4
  10. 2025Towards Understanding Gradient Dynamics of the Sliced-Wasserstein Distance via Critical Point Analysiscited by 3
  11. 2026QUATRO: Query-Adaptive Trust Region Policy Optimization for LLM Fine-tuningcited by 2
  12. 2025Weighted Conditional Flow Matchingcited by 1
10 more
  1. 2026Rate-Optimal Noise Annealing in Semi-Dual Neural Optimal Transport: Tangential Identifiability, Off-Manifold Ambiguity, and Guaranteed Recoverycited by 1
  2. 2026Learning Discrete Diffusion of Graphs via Free-Energy Gradient Flowscited by 1
  3. 2024A Differential Equation Approach for Wasserstein GANs and Beyondcited by 0
  4. 2025Guiding Time-Varying Generative Models with Natural Gradients on Exponential Family Manifoldcited by 0
  5. 2025Helping CLIP See Both the Forest and the Trees: A Decomposition and Description Approachcited by 0
  6. 2026Rethinking the Flow-Based Gradual Domain Adaptation: A Semi-Dual Optimal Transport Perspectivecited by 0
  7. 2026Analyzing and Improving Diffusion Models for Time-Series Data Imputation: A Proximal Recursion Perspectivecited by 0
  8. 2026UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problemscited by 0
  9. 2026Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flowcited by 0
  10. 2026One-Step Generative Modeling via Unbalanced Optimal Transportcited by 0