A Unifying View of Variational Generative Wasserstein Flows
Abstract
Many modern generative models can be viewed as minimizing divergences between probability distributions, yet they rely on different algorithmic and geometric principles. Wasserstein gradient flows provide a continuous-time formulation for optimizing over distributions, and can be approximated through their implicit discretization via the Jordan–Kinderlehrer–Otto (JKO) scheme. In this work, we present a unified theoretical framework for generative modeling based on Wasserstein gradient flows, which we refer to as Generative Wasserstein Flows. We show that a broad class of existing methods can be derived as instances of parametric JKO schemes for f-divergences objectives, and we establish equivalences between several recently proposed algorithms. We extend this framework beyond f-divergences to integral probability metrics, deriving new JKO-based generative algorithms for objectives such as Maximum Mean Discrepancy. We also clarify their connections with GANs. Finally, we analyze parametric Wasserstein flows, where the evolution is restricted to distributions generated by parameterized maps. We characterize the resulting dynamics as projected or preconditioned Wasserstein gradient flows, highlighting the role of the Wasserstein geometry in shaping the learning dynamics of generative models.