Convergence Analysis of Wasserstein Proximal Algorithm beyond Geodesic Convexity

TMLR Paper9014 Authors

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 on the objective functional. Under a natural Wasserstein analog of the Euclidean Polyak-{\L}ojasiewicz inequality, we show that the proximal algorithm achieves an unbiased and dimension-free linear convergence rate. Our convergence rate improves upon existing rates of the proximal algorithm for solving Wasserstein gradient flows when specialized to strong geodesic convex functionals. 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 two-layer mean-field neural networks.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission: None

Assigned Action Editor: ~Pan_Xu1

Submission Number: 9014