Anderson Mixing in Bures Wasserstein Space of Gaussian Measures

TL;DR

This paper introduces a Riemannian Anderson Mixing method for Gaussian measures in Bures Wasserstein space, significantly accelerating fixed-point computations like Wasserstein barycenters with minimal additional costs.

Contribution

It develops and implements the Riemannian Anderson Mixing method tailored for Gaussian distributions in Bures Wasserstein space, enhancing convergence speed for fixed-point problems.

Findings

Significant acceleration over Picard iteration

Comparable performance to Riemannian Gradient Descent

Applicable convergence in certain Bures-Wasserstein regions

Abstract

Various statistical tasks, including sampling or computing Wasserstein barycenters, can be reformulated as fixed-point problems for operators on probability distributions. Accelerating standard fixed-point iteration schemes provides a promising novel approach to the design of efficient numerical methods for these problems. The Wasserstein geometry on the space of probability measures, although not precisely Riemannian, allows us to define various useful Riemannian notions, such as tangent spaces, exponential maps and parallel transport, motivating the adaptation of Riemannian numerical methods. We demonstrate this by developing and implementing the Riemannian Anderson Mixing (RAM) method for Gaussian distributions. The method reuses the history of the residuals and improves the iteration complexity, and we argue that the additional costs, compared to Picard method, are negligible. We show that certain open balls in the Bures-Wasserstein manifold satisfy the requirements for convergence of RAM. The numerical experiments show a significant acceleration compared to a Picard iteration, and performance on par with Riemannian Gradient Descent and Conjugate Gradient methods.