A Unifying View of Variational Generative Wasserstein Flows

Paul Caucheteux, Clément Bonet, Anna Korba

TL;DR: We unify several generative modeling methods through Wasserstein gradient flows and the JKO scheme, revealing connections between existing algorithms and leading to new ones.

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 (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$-divergences 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.

Lay Summary:

Generative models learn to produce new samples, such as images, that look like real data. However, many popular generative models are based on different training objectives and algorithms, which makes it difficult to understand how they are related.

In this paper, we propose a unified way to view several generative modeling methods. The key idea is to see learning as a gradual evolution of probability distributions: starting from a simple distribution that we can sample from, the model progressively transforms it into one that is close to the data distribution. We use a mathematical tool called the JKO scheme, which can be understood as a principled way to take stable steps along this evolution path.

Our framework shows that several existing methods are different expressions of the same underlying scheme, and it also leads to new algorithms. Experiments show that this perspective can improve or stabilize training in several cases. We also provide a first theoretical analysis of why this kind of model performs well. This helps clarify the mathematical structure behind generative models and may guide the design of future methods.

Link To Code: https://github.com/Paulcauch/Generative_Wasserstein_Flows

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Wasserstein gradient flows, Jordan–Kinderlehrer–Otto (JKO) scheme, generative adversarial networks

Submission Number: 16954