A Unifying View of Variational Generative Wasserstein Flows
This work provides a unifying theoretical perspective for researchers and practitioners working with various generative models, clarifying connections and enabling the derivation of new algorithms.
This paper proposes Generative Wasserstein Flows (GWF), a unified theoretical framework for generative modeling based on Wasserstein gradient flows. It demonstrates that many existing methods can be derived as instances of parametric JKO schemes for f-divergence objectives, and extends this framework to Integral Probability Metrics and squared Maximum Mean Discrepancy, deriving new JKO-based algorithms.
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 (GWF). We show that a broad class of existing methods can be derived as instances of parametric JKO schemes for $f$-divergence objectives, and we establish equivalences between several recently proposed algorithms. We extend this framework beyond f-divergence to Integral Probability Metrics and squared Maximum Mean Discrepancy, deriving new JKO-based generative algorithms, and clarifying their connections with GANs. We study empirically the impact of the JKO regularization for a wide set of objectives. Finally, we analyze parametric Wasserstein flows, where the dynamics are restricted to distributions induced by parametrized maps.