Wasserstein accelerated proximal gradient algorithms

Abstract

We propose accelerated proximal gradient methods for composite optimization problems on the quadratic Wasserstein space of probability measures, including Kullback--Leibler divergence minimization. The objective function consists of a smooth geodesically convex functional and a possibly nonsmooth regularizer that is convex along generalized geodesics. The convexity and smoothness conditions on the smooth component are imposed only on the feasible set. As a preliminary result, we establish the equivalence between geodesic convexity of the smooth functional on the feasible set and convexity of its Hilbert-space lift on the corresponding feasible set. Assuming exact proximal steps, we prove an $O(k^{-2})$ convergence rate for the objective gap when the smooth component is geodesically convex and accelerated linear convergence when it is geodesically strongly convex. We also describe particle approximations of these algorithms.