An Accelerated Proximal Gradient Method for Bures–Wasserstein Space
Jakob Nylöf ⋅ Sloan Nietert ⋅ Julien Pallage ⋅ Giancarlo Ferrari-Trecate ⋅ Daniel Kuhn ⋅ John Lygeros
Abstract
We study composite optimization over the Bures–Wasserstein (BW) space of Gaussian measures. Our main observation is that a useful class of BW problems can be lifted into the Euclidean space of affine transport maps, such that the lifted problems can be solved efficiently using *accelerated* proximal gradient methods. For objectives whose smooth and nonsmooth components admit convex smooth and convex proximable affine lifts, respectively, our algorithm BW-FISTA achieves minimax-optimal objective gap $O(k^{-2})$ after $k$ iterations. We verify these regularity conditions for potential and interaction energies together with common spectral regularizers, and we treat stochastic gradients under a suitable noise model. As a principal application, we consider Gaussian variational inference with a log-concave, log-smooth target. Here, the resulting entropy term is handled by an efficient, exact proximal step, and we improve the previous rate of $O(k^{-1})$ to $O(k^{-2})$. We numerically verify this predicted boost with synthetic experiments.
Chat is not available.
Successful Page Load