Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity

TL;DR

This paper proves that the Wasserstein proximal algorithm converges linearly without geodesic convexity assumptions, extending its applicability and demonstrating faster convergence than noisy gradient descent in neural network training.

Contribution

It provides a novel convergence analysis of the Wasserstein proximal algorithm under a Wasserstein Polyak-{\L}ojasiewicz inequality, beyond geodesic convexity.

Findings

Achieves unbiased, linear convergence rate under new conditions.

Improves convergence rates over existing methods for Wasserstein gradient flows.

Numerical experiments show faster convergence than noisy gradient descent in neural networks.

Abstract

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-{\L}ojasiewicz 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.