Videos

Wasserstein proximal coordinate gradient algorithms

Presenter
May 7, 2024
Abstract
This talk concerns composite (geodesically) convex optimization over multiple distributions. The objective functional under consideration is composed of a convex potential energy, defined on a product of Wasserstein spaces (the space of all distributions with a finite second moment), and a sum of convex self-interaction and internal energies associated with each distribution. To efficiently solve this problem, we introduce the Wasserstein Proximal Coordinate Gradient (WPCG) algorithm. Under a Quadratic Growth (QG) condition on the objective functional, a condition more relaxed than the typical strongly convex requirement, WPCG is proven to attain exponential convergence to the unique global optimum. Implications regarding the choice of step size and update schemes (parallel, sequential and random) are also discussed. In the absence of the QG condition, WPCG is still demonstrated to converge to the global optimal solution, albeit at a slower polynomial rate. The algorithm and theoretical framework are applied to two representative examples: approximation Bayesian computation using mean-field variational approximation, and the computation of equilibrium in multi-species systems with cross-interaction. Numerical results for both examples are consistent with our theoretical findings.