Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space

TL;DR

This paper introduces a novel stochastic forward-backward algorithm for Gaussian variational inference that leverages the Bures-Wasserstein space structure, providing convergence guarantees under various conditions.

Contribution

It develops the first forward-backward Gaussian VI algorithm utilizing the Bures-Wasserstein space with proven convergence guarantees.

Findings

Achieves state-of-the-art convergence for log-smooth, log-concave targets.

Provides first convergence guarantees to stationary points for broader target classes.

Utilizes the composite structure of KL divergence in Wasserstein space.

Abstract

Variational inference (VI) seeks to approximate a target distribution $\pi$ by an element of a tractable family of distributions. Of key interest in statistics and machine learning is Gaussian VI, which approximates $\pi$ by minimizing the Kullback-Leibler (KL) divergence to $\pi$ over the space of Gaussians. In this work, we develop the (Stochastic) Forward-Backward Gaussian Variational Inference (FB-GVI) algorithm to solve Gaussian VI. Our approach exploits the composite structure of the KL divergence, which can be written as the sum of a smooth term (the potential) and a non-smooth term (the entropy) over the Bures-Wasserstein (BW) space of Gaussians endowed with the Wasserstein distance. For our proposed algorithm, we obtain state-of-the-art convergence guarantees when $\pi$ is log-smooth and log-concave, as well as the first convergence guarantees to first-order stationary solutions when $\pi$ is only log-smooth.