# Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 45 references that are in the map.
- First posted on arXiv 25 Jan 2025.

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

- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices](https://mlsanity.com/paper/1903.08568/) (2019)
- [Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit](https://mlsanity.com/paper/1902.06015/) (2019)
- [Analysis of Langevin Monte Carlo via convex optimization](https://mlsanity.com/paper/1802.09188/) (2018)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Improved analysis for a proximal algorithm for sampling](https://mlsanity.com/paper/2202.06386/) (2022)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (2022)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Mean-Field Langevin Dynamics and Energy Landscape of Neural Networks](https://mlsanity.com/paper/1905.07769/) (2019)
- [Minimax Mixing Time of the Metropolis-Adjusted Langevin Algorithm for Log-Concave Sampling](https://mlsanity.com/paper/2109.13055/) (2021)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Proximal Langevin Algorithm: Rapid Convergence Under Isoperimetry](https://mlsanity.com/paper/1911.01469/) (2019)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Improved dimension dependence of a proximal algorithm for sampling](https://mlsanity.com/paper/2302.10081/) (2023)
- [Improved Discretization Analysis for Underdamped Langevin Monte Carlo](https://mlsanity.com/paper/2302.08049/) (2023)
- [Uniform-in-time propagation of chaos for mean field Langevin dynamics](https://mlsanity.com/paper/2212.03050/) (2022)
- [Mean-field Variational Inference via Wasserstein Gradient Flow](https://mlsanity.com/paper/2207.08074/) (2022)
- [First-order Conditions for Optimization in the Wasserstein Space](https://mlsanity.com/paper/2209.12197/) (2022)
- [Sampling from the Mean-Field Stationary Distribution](https://mlsanity.com/paper/2402.07355/) (2024)
- [Wasserstein Proximal Coordinate Gradient Algorithms](https://mlsanity.com/paper/2405.04628/) (2024)
- [Proximal Oracles for Optimization and Sampling](https://mlsanity.com/paper/2404.02239/) (2024)
- [Non-geodesically-convex optimization in the Wasserstein space](https://mlsanity.com/paper/2406.00502/) (2024)
- [Weak topology and Opial property in Wasserstein spaces, with applications to Gradient Flows and Proximal Point Algorithms of geodesically convex functionals](https://mlsanity.com/paper/2104.06121/) (2021)
- [Linear convergence of proximal descent schemes on the Wasserstein space](https://mlsanity.com/paper/2411.15067/) (2024)
- [On Representations of Mean-Field Variational Inference](https://mlsanity.com/paper/2210.11385/) (2022)
- [Mean-field underdamped Langevin dynamics and its spacetime discretization](https://mlsanity.com/paper/2312.16360/) (2023)

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

- [Wasserstein Saddle-Free Newton: Saddle Escape and Local Spectral Convergence](https://mlsanity.com/paper/2605.17963/) (2026)
- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)

## Abstract (arXiv)

The proximal algorithm is a powerful tool to minimize nonlinear and nonsmooth functionals in a general metric space. Motivated by the recent progress in studying the training dynamics of the noisy gradient descent algorithm on two-layer neural networks in the mean-field regime, we provide in this paper a simple and self-contained analysis for the convergence of the general-purpose Wasserstein proximal algorithm without assuming geodesic convexity of the objective functional. Under a natural Wasserstein analog of the Euclidean Polyak-Lojasiewicz inequality, we establish that the proximal algorithm achieves an unbiased and linear convergence rate. Our convergence rate improves upon existing rates of the proximal algorithm for solving Wasserstein gradient flows under strong geodesic convexity. We also extend our analysis to the inexact proximal algorithm for geodesically semiconvex objectives. In our numerical experiments, proximal training demonstrates a faster convergence rate than the noisy gradient descent algorithm on mean-field neural networks.