Combining Wasserstein-1 and Wasserstein-2 proximals: robust manifold learning via well-posed generative flows

TL;DR

This paper introduces a novel approach combining Wasserstein-1 and Wasserstein-2 proximals to create well-posed, robust generative flows capable of learning distributions on low-dimensional manifolds, with theoretical guarantees and practical efficacy.

Contribution

It formulates a new mean-field game framework integrating Wasserstein-1 and Wasserstein-2 proximals for stable, well-posed generative flows with provable uniqueness.

Findings

The combined proximals ensure well-posedness and uniqueness of the generative flow system.

The approach effectively generates high-dimensional images without autoencoders.

The method demonstrates robustness and theoretical soundness in learning low-dimensional manifold distributions.

Abstract

We formulate well-posed continuous-time generative flows for learning distributions that are supported on low-dimensional manifolds through Wasserstein proximal regularizations of $f$-divergences. 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, we show that the 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 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 adversarial training of continuous-time flows, which bypasses the need for reverse simulation. We demonstrate the efficacy of our approach for generating high-dimensional images without the need to resort to autoencoders or specialized architectures.