Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Zhu, Shuailong, Chen, Xiaohui
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866918509628358656
author Zhu, Shuailong
Chen, Xiaohui
author_facet Zhu, Shuailong
Chen, Xiaohui
contents 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-Ł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.
format Preprint
id arxiv_https___arxiv_org_abs_2501_14993
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity
Zhu, Shuailong
Chen, Xiaohui
Optimization and Control
Machine Learning
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-Ł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.
title Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity
topic Optimization and Control
Machine Learning
url https://arxiv.org/abs/2501.14993