Wasserstein Proximal Coordinate Gradient Algorithms

TL;DR

This paper introduces Wasserstein Proximal Coordinate Gradient algorithms for efficiently solving complex geodesically convex optimization problems over multiple distributions, with proven exponential and polynomial convergence rates.

Contribution

The paper develops novel WPCG algorithms with various update schemes and establishes convergence under weaker conditions than strong convexity.

Findings

WPCG converges exponentially under quadratic growth conditions.

WPCG converges polynomially without quadratic growth.

Numerical results support theoretical convergence claims.

Abstract

Motivated by approximation Bayesian computation using mean-field variational approximation and the computation of equilibrium in multi-species systems with cross-interaction, this paper investigates the composite geodesically convex optimization problem over multiple distributions. The objective functional under consideration is composed of a convex potential energy on a product of Wasserstein spaces 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) algorithms with parallel, sequential, and random update schemes. Under a quadratic growth (QG) condition that is weaker than the usual strong convexity requirement on the objective functional, we show that WPCG converges exponentially fast to the unique global optimum. In the absence of the QG condition, WPCG is still demonstrated to converge to the global optimal solution, albeit at a slower polynomial rate. Numerical results for both motivating examples are consistent with our theoretical findings.