Convergence of Flow-Based Generative Models via Proximal Gradient Descent in Wasserstein Space
Year: 2024
Type: article
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(\varepsilon ^{2})$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-$\mathcal {W}_{2}$mixed error guarantees. The non-asymptotic convergence rate of the JKO-type$\mathcal {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.
Authors: Xiuyuan Cheng, Jianfeng Lu, Yixin Tan, Yao Xie
Institutions: Duke University, Georgia Institute of Technology
Language: English
Subfield: Computer Vision and Pattern Recognition
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: ATD: Scanning Dynamic Spatial-Temporal Discrete Events for Threat Detection, Bridging Statistical Hypothesis Tests and Deep Learning for Reliability and Computational Efficiency, CAREER: Quick Detection for Streaming Data Over Dynamic Networks, G3387330441, G4324459455, Data-Driven Optimal Police Patrol Zone Districting and Staffing, Sequential Detection and Prediction for Solar Situation Awareness in Power Networks, AI Institute for Advances in Optimization, G6127310570, G6671297155, Innovative Numerical Methods for High-Dimensional Applications, G7238632903, Innovation of Numerical Methods for High-Dimensional Partial Differential Equations, G7911581280, CAREER: Learning of graph diffusion and transport from high dimensional data with low-dimensional structures, G975068223
OpenAlex ID: W4400276434