A Unifying View of Variational Generative Wasserstein Flows

2026-05-29

AI Summary

This work proposes a unified generative modeling framework grounded in Wasserstein gradient flows, interpreting diverse mainstream methods as parameterized instances of the Jordan–Kinderlehrer–Otto (JKO) scheme under various divergences or integral probability metrics—such as f-divergences and maximum mean discrepancy (MMD). While existing generative models aim to minimize distributional discrepancies, they lack a cohesive theoretical account of their algorithmic and geometric diversity. The proposed framework systematically reveals the intrinsic equivalences among distinct generative algorithms and extends beyond f-divergences, establishing novel connections to generative adversarial networks (GANs). Building on this foundation, the authors derive new JKO-based generative algorithms, empirically demonstrate the efficacy of JKO regularization across multiple objectives, and elucidate the dynamic evolution of parameterized Wasserstein flows in terms of the induced pushforward distributions.

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-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.

Problem

• Generative Modeling
• Wasserstein Gradient Flows
• f-divergence
• Integral Probability Metrics
• Maximum Mean Discrepancy

Innovation

• Wasserstein gradient flows
• JKO scheme
• generative modeling
• f-divergence
• Integral Probability Metrics