# Convergence of flow-based generative models via proximal gradient descent in Wasserstein space

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 79 references that are in the map.
- First posted on arXiv 26 Oct 2023.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (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)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [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)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Maximum Likelihood Training of Score-Based Diffusion Models](https://mlsanity.com/paper/2101.09258/) (2021)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization](https://mlsanity.com/paper/1606.00709/) (2016)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [Rectified Flow: A Marginal Preserving Approach to Optimal Transport](https://mlsanity.com/paper/2209.14577/) (2022)
- [Convergence of denoising diffusion models under the manifold hypothesis](https://mlsanity.com/paper/2208.05314/) (2022)
- [Nearly d-Linear Convergence Bounds for Diffusion Models via Stochastic Localization](https://mlsanity.com/paper/2308.03686/) (2023)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions](https://mlsanity.com/paper/2211.01916/) (2022)
- [Convergence of score-based generative modeling for general data distributions](https://mlsanity.com/paper/2209.12381/) (2022)
- [Convergence for score-based generative modeling with polynomial complexity](https://mlsanity.com/paper/2206.06227/) (2022)
- [The probability flow ODE is provably fast](https://mlsanity.com/paper/2305.11798/) (2023)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Theoretical guarantees for sampling and inference in generative models with latent diffusions](https://mlsanity.com/paper/1903.01608/) (2019)

## Built on by (most cited first, 30 of 34)

- [A Sharp Convergence Theory for The Probability Flow ODEs of Diffusion Models](https://mlsanity.com/paper/2408.02320/) (2024)
- [Convergence Analysis for General Probability Flow ODEs of Diffusion Models in Wasserstein Distances](https://mlsanity.com/paper/2401.17958/) (2024)
- [Low-dimensional adaptation of diffusion models: Convergence in total variation](https://mlsanity.com/paper/2501.12982/) (2025)
- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)
- [Provable Acceleration for Diffusion Models under Minimal Assumptions](https://mlsanity.com/paper/2410.23285/) (2024)
- [Convergence of Continuous Normalizing Flows for Learning Probability Distributions](https://mlsanity.com/paper/2404.00551/) (2024)
- [Convergence of Score-Based Discrete Diffusion Models: A Discrete-Time Analysis](https://mlsanity.com/paper/2410.02321/) (2024)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Broadening Target Distributions for Accelerated Diffusion Models via a Novel Analysis Approach](https://mlsanity.com/paper/2402.13901/) (2024)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [Faster Diffusion Models via Higher-Order Approximation](https://mlsanity.com/paper/2506.24042/) (2025)
- [Flow-based Distributionally Robust Optimization](https://mlsanity.com/paper/2310.19253/) (2023)
- [Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions](https://mlsanity.com/paper/2409.20547/) (2024)
- [Integrating Dynamical Systems Modeling with Spatiotemporal scRNA-seq Data Analysis](https://mlsanity.com/paper/2503.11347/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Theory on Score-Mismatched Diffusion Models and Zero-Shot Conditional Samplers](https://mlsanity.com/paper/2410.13746/) (2024)
- [Worst-case generation via minimax optimization in Wasserstein space](https://mlsanity.com/paper/2512.08176/) (2025)
- [Posterior sampling via Langevin dynamics based on generative priors](https://mlsanity.com/paper/2410.02078/) (2024)
- [Local Flow Matching Generative Models](https://mlsanity.com/paper/2410.02548/) (2024)
- [Assured autonomy: How operations research powers and orchestrates generative AI systems](https://mlsanity.com/paper/2512.23978/) (2025)
- [Linear convergence of proximal descent schemes on the Wasserstein space](https://mlsanity.com/paper/2411.15067/) (2024)
- [Pathway to O(√d) Complexity bound under Wasserstein metric of flow-based models](https://mlsanity.com/paper/2512.06702/) (2025)
- [Generative models for decision-making under distributional shift](https://mlsanity.com/paper/2604.04342/) (2026)
- [Dimension-free error estimate for diffusion model and optimal scheduling](https://mlsanity.com/paper/2512.01820/) (2025)
- [Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity](https://mlsanity.com/paper/2501.14993/) (2025)
- [Flow-based generative models as iterative algorithms in probability space](https://mlsanity.com/paper/2502.13394/) (2025)
- [Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow](https://mlsanity.com/paper/2601.05583/) (2026)
- [New algorithms for sampling and diffusion models](https://mlsanity.com/paper/2406.09665/) (2024)
- [Diffusion Models for Adaptive Sequential Data Generation](https://mlsanity.com/paper/2606.06007/) (2026)
- [Embedding Empirical Distributions for Computing Optimal Transport Maps](https://mlsanity.com/paper/2504.17740/) (2025)

## Abstract (arXiv)

Flow-based generative models enjoy certain advantages in computing the data generation and the likelihood, and have recently shown competitive empirical performance. Compared to the accumulating theoretical studies on related score-based diffusion models, analysis of flow-based models, which are deterministic in both forward (data-to-noise) and reverse (noise-to-data) directions, remain sparse. In this paper, we provide a theoretical guarantee of generating data distribution by a progressive flow model, the so-called JKO flow model, which implements the Jordan-Kinderleherer-Otto (JKO) scheme in a normalizing flow network. Leveraging the exponential convergence of the proximal gradient descent (GD) in Wasserstein space, we prove the Kullback-Leibler (KL) guarantee of data generation by a JKO flow model to be O(varepsilon2) when using N lesssim log (1/varepsilon) many JKO steps (N Residual Blocks in the flow) where varepsilon  is the error in the per-step first-order condition. The assumption on data density is merely a finite second moment, and the theory extends to data distributions without density and when there are inversion errors in the reverse process where we obtain KL-W2 mixed error guarantees. The non-asymptotic convergence rate of the JKO-type W2-proximal GD is proved for a general class of convex objective functionals that includes the KL divergence as a special case, which can be of independent interest. The analysis framework can extend to other first-order Wasserstein optimization schemes applied to flow-based generative models.