Variational inference via Gaussian interacting particles in the Bures-Wasserstein geometry

TL;DR

This paper introduces a novel Gaussian particle-based optimization algorithm in the Bures-Wasserstein geometry for variational inference, demonstrating robustness and improved performance over gradient methods.

Contribution

It proposes a new LBW space for Gaussian measures and a consensus-based optimization algorithm with theoretical convergence analysis.

Findings

Algorithm outperforms gradient-based methods on non log-concave targets.

LBW space enables efficient computation while preserving geometric features.

Theoretical analysis confirms well-posedness and convergence of the particle system.

Abstract

Motivated by variational inference methods, we propose a zeroth-order algorithm for solving optimization problems in the space of Gaussian probability measures. The algorithm is based on an interacting system of Gaussian particles that stochastically explore the search space and self-organize around global minima via a consensus-based optimization (CBO) mechanism. Its construction relies on the Linearized Bures-Wasserstein (LBW) space, a novel parametrization of Gaussian measures we introduce for efficient computations. LBW is inspired by linearized optimal transport and preserves key geometric features while enabling computational tractability. We establish well-posedness and study the convergence properties of the particle dynamics via a mean-field approximation. Numerical experiments on variational inference tasks demonstrate the algorithm's robustness and superior performance with respect to deterministic gradient-based method in presence of low-dimensional non log-concave targets.