Hyemin Gu: Well-posed generative flows via combined Wasserstein-1 and Wasserstein-2 proximals of $f$-divergences
Abstract
This talk focuses on well-posed continuous-time generative flows for learning distributions that are supported on low-dimensional manifolds, using Wasserstein proximal regularizations of $f$-divergences to formulate these flows. Wasserstein-1 proximal operators regularize $f$-divergences so that singular distributions can be compared. Meanwhile, Wasserstein-2 proximal operators regularize the paths of the generative flows by adding an optimal transport cost, i.e., a kinetic energy penalization. Via mean-field game theory, it is shown that the \emph{combination} of the two proximals is critical for formulating well-posed generative flows. Generative flows can be analyzed through optimality conditions of a mean-field game (MFG), a system of a backward Hamilton-Jacobi (HJ) and a forward continuity partial differential equations (PDEs) whose solution characterizes the optimal generative flow. For learning distributions that are supported on low-dimensional manifolds, the MFG theory shows that the Wasserstein-1 proximal, which addresses the HJ terminal condition, and the Wasserstein-2 proximal, which addresses the HJ dynamics, are both necessary for the corresponding backward-forward PDE system to be well-defined and have a unique solution with provably \emph{linear} flow trajectories. This implies that the corresponding generative flow is also unique and can therefore be learned in a robust manner even for learning high-dimensional distributions supported on low-dimensional manifolds. The generative flows are learned through \emph{adversarial training} of continuous-time flows, which bypasses the need for reverse simulation. Numerical experiments demonstrate the efficacy of our approach for generating high-dimensional images without the need of autoencoders or specialized architectures for the normalizing flow. Our work demonstrates how ideas from the analysis of multi-agent systems contribute to the design and understanding of generative flows.