Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Anderson Mixing in Bures Wasserstein Space of Gaussian Measures

Vitalii Aksenov, Martin Eigel, Mathias Oster

TL;DR

This work extends RAM to the Bures-Wasserstein space of Gaussian measures (BW Gaussians), enabling accelerated fixed-point iterations for tasks such as sampling and Wasserstein barycenter computation. By exploiting BW geometry—tangent spaces, the exponential map, and vector transport—the authors formulate a Riemannian RAM (BWRAM) and prove a local convergence result within a BW ball around a nondegenerate fixed point, with a residual contraction that is competitive with Picard iteration. The paper provides explicit operator examples, including Ornstein–Uhlenbeck dynamics and KL-divergence minimization, and analyzes computational complexity of RAM versus standard methods. Numerical experiments across OU dynamics, KL minimization, and distribution averaging demonstrate significant speedups over Picard and performance on par with Riemannian Gradient Descent and Conjugate Gradient, while offering practical guidance on vector transport choices. Overall, this work delivers a computationally efficient, geometry-aware acceleration framework for a broad class of Wasserstein-space problems restricted to Gaussian measures.

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.

Anderson Mixing in Bures Wasserstein Space of Gaussian Measures

TL;DR

This work extends RAM to the Bures-Wasserstein space of Gaussian measures (BW Gaussians), enabling accelerated fixed-point iterations for tasks such as sampling and Wasserstein barycenter computation. By exploiting BW geometry—tangent spaces, the exponential map, and vector transport—the authors formulate a Riemannian RAM (BWRAM) and prove a local convergence result within a BW ball around a nondegenerate fixed point, with a residual contraction that is competitive with Picard iteration. The paper provides explicit operator examples, including Ornstein–Uhlenbeck dynamics and KL-divergence minimization, and analyzes computational complexity of RAM versus standard methods. Numerical experiments across OU dynamics, KL minimization, and distribution averaging demonstrate significant speedups over Picard and performance on par with Riemannian Gradient Descent and Conjugate Gradient, while offering practical guidance on vector transport choices. Overall, this work delivers a computationally efficient, geometry-aware acceleration framework for a broad class of Wasserstein-space problems restricted to Gaussian measures.

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.
Paper Structure (26 sections, 14 theorems, 127 equations, 4 figures, 6 tables, 2 algorithms)

This paper contains 26 sections, 14 theorems, 127 equations, 4 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Anderson mixing in Euclidean space and for Riemannian manifolds with bounded sectional curvature
  3. The Bures-Wasserstein space of Gaussians
  4. Geometric notions
  5. Properties in the neighborhood of a nondegenerate distribution
  6. Construction of the vector transport map
  7. Riemannian Anderson Mixing on Bures-Wasserstein space
  8. Convergence
  9. Examples of operators
  10. Ornstein-Uhlenbeck process
  11. Minimization of functionals
  12. Kullback-Leibler divergence
  13. Complexity analysis
  14. Numerical experiments
  15. Implementation details
  16. ...and 11 more sections

Key Result

Theorem 1

Let $G(\Sigma) = \operatorname{Exp}_\Sigma(-F(\Sigma))$ be a contractive mapping on the Bures-Wasserstein manifold and $\Sigma_* \succeq \lambda \operatorname{Id}$ its fixed point with $\lambda>0$. If the initial iterate $\Sigma_0 \in B_{W_2}(\Sigma_*, r)$ for a sufficiently small $r$, and under add where $\operatorname{Exp}$ denotes the Riemannian exponential map, $r_k = -F(x_k)$ the fixed-point

Figures (4)

  • Figure 1: Approximate vector transport along Bures-Wasserstein geodesic
  • Figure 2: Convergence of BWRAM on model problems, in comparison to Picard (blue), RGD (orange) and RCG (red) methods
  • Figure 3: Convergence of the cost. BWRAM (greens) with different history length on the averaging problems, in comparison to Picard (blue), RGD (orange) and RCG (red) methods
  • Figure 4: Performance of methods on the $\operatorname{KL}$ problem with varying condition number $\sqrt{\frac{\lambda^*_1}{\lambda^*_d}}$. Exact parallel transport, one-step and trivial approximation.

Theorems & Definitions (29)

  • Theorem 1: Convergence of BWRAM (informal statement)
  • Remark 3.1
  • Theorem 2
  • Remark 3.2
  • Proposition 3.1
  • Remark 3.3
  • Theorem 3
  • Remark 4.1
  • Theorem 4
  • proof
  • ...and 19 more