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

Abstract

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({\epsilon }^2)$ when using $N\lesssim \log (1/\epsilon )$ many JKO steps ($N$ Residual Blocks in the flow) where $\epsilon$ 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-$W_2$ mixed error guarantees. The non-asymptotic convergence rate of the JKO-type $W_2$-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.

References in corpus (15)

• An Introduction to Variational Autoencoders
• NICE: Non-linear Independent Components Estimation
• Normalizing Flows: An Introduction and Review of Current Methods
• A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems
• Stochastic Interpolants: A Unifying Framework for Flows and Diffusions
• Convergence of denoising diffusion models under the manifold hypothesis
• Bridging Mean-Field Games and Normalizing Flows with Trajectory Regularization
• An error analysis of generative adversarial networks for learning distributions
• Rectified Flow: A Marginal Preserving Approach to Optimal Transport
• Classification Logit Two-sample Testing by Neural Networks
• Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks
• Improved Convergence of Score-Based Diffusion Models via Prediction-Correction
• Error Bounds for Flow Matching Methods
• Distribution learning via neural differential equations: a nonparametric statistical perspective
• Computing high-dimensional optimal transport by flow neural networks

Cited by in corpus (1)

• Integrating Dynamical Systems Modeling with Spatiotemporal scRNA-seq Data Analysis