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

arXiv: https://arxiv.org/abs/2407.11901 · HTML: https://mlsanity.com/paper/2407.11901/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 23 references that are in the map.
- First posted on arXiv 16 Jul 2024.

## Builds on (most cited first, 17 of 17)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit](https://mlsanity.com/paper/1905.09883/) (2019)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Distributionally Robust Optimization: A Review](https://mlsanity.com/paper/1908.05659/) (2019)
- [Potential Flow Generator with L2 Optimal Transport Regularity for Generative Models](https://mlsanity.com/paper/1908.11462/) (2019)
- [A mean-field games laboratory for generative modeling](https://mlsanity.com/paper/2304.13534/) (2023)
- [ResNet After All? Neural ODEs and Their Numerical Solution](https://mlsanity.com/paper/2007.15386/) (2020)
- [Structure-preserving GANs](https://mlsanity.com/paper/2202.01129/) (2022)
- [Score-based generative models are provably robust: an uncertainty quantification perspective](https://mlsanity.com/paper/2405.15754/) (2024)
- [Wasserstein proximal operators describe score-based generative models and resolve memorization](https://mlsanity.com/paper/2402.06162/) (2024)
- [A kernel formula for regularized Wasserstein proximal operators](https://mlsanity.com/paper/2301.10301/) (2023)
- [Lipschitz-regularized gradient flows and generative particle algorithms for high-dimensional scarce data](https://mlsanity.com/paper/2210.17230/) (2022)
- [Function-space regularized R'enyi divergences](https://mlsanity.com/paper/2210.04974/) (2022)
- [Moreau-Yosida f-divergences](https://mlsanity.com/paper/2102.13416/) (2021)

## Built on by (most cited first, 4 of 4)

- [OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization](https://mlsanity.com/paper/2501.18793/) (2025)
- [Proximal optimal transport divergences](https://mlsanity.com/paper/2505.12097/) (2025)
- [Gradient flow in parameter space is equivalent to linear interpolation in output space](https://mlsanity.com/paper/2408.01517/) (2024)
- [Policy Transfer for Continuous-Time Reinforcement Learning: A (Rough) Differential Equation Approach](https://mlsanity.com/paper/2510.15165/) (2025)

## Abstract (arXiv)

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.