Variational inference via Gaussian interacting particles in the Bures-Wasserstein geometry
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. 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.
Lay Summary
Many scientific and statistical problems require estimating an unknown probability distribution, for example the range of plausible explanations for observed data. A common shortcut is to approximate this distribution by a Gaussian, or bell-shaped, distribution, because Gaussians are easy to store and compute with. However, choosing the best Gaussian approximation can be difficult when the target distribution has several peaks or an irregular shape. We propose a new method that searches over Gaussian distributions without needing gradient information. The method uses a group of Gaussian “particles” that explore different possible approximations and gradually move toward a shared consensus, giving more influence to particles that better match the target distribution. To make this efficient, we introduce a new way of representing Gaussian distributions that preserves useful geometric structure while simplifying the computations. Our theoretical results show that the resulting particle system is well-defined and explain its convergence behavior through a large-population approximation. In experiments on variational inference tasks, the method is more robust than standard gradient-based approaches in low-dimensional non-convex settings.