Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space

Abstract

We develop a theory of finite-dimensional polyhedral subsets over the Wasserstein space and optimization of functionals over them via first-order methods. Our main application is to the problem of mean-field variational inference, which seeks to approximate a distribution $\ddot{I}€$ over ${\mathbb{R}}^d$ by a product measure $\ddot{I}{€}^{\star }$. When $\ddot{I}€$ is strongly log-concave and log-smooth, we provide (1) approximation rates certifying that $\ddot{I}{€}^{\star }$ is close to the minimizer $\ddot{I}{€}_{\diamond }^{\star }$ of the KL divergence over a \emph{polyhedral} set ${\mathcal{P}}_{\diamond }$, and (2) an algorithm for minimizing $\text{KL}(\cdot \Vert \ddot{I}€)$ over ${\mathcal{P}}_{\diamond }$ based on accelerated gradient descent over ${\mathbb{R}}^d$. As a byproduct of our analysis, we obtain the first end-to-end analysis for gradient-based algorithms for MFVI.

49 pages