Algorithms for Mean-Field Variational Inference Via Polyhedral Optimization in the Wasserstein Space

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 (MFVI), which seeks to approximate a distribution \pi $\pi$ over \mathbb {R}^d$\mathbb {R}^d$ by a product measure \pi ^\star $\pi ^\star$. When \pi $\pi$ is strongly log-concave and log-smooth, we provide (1) approximation rates certifying that \pi ^\star $\pi ^\star$ is close to the minimizer \pi ^\star _\diamond $\pi ^\star _\diamond$ of the KL divergence over a polyhedral set \mathcal {P}_\diamond $\mathcal {P}_\diamond$, and (2) an algorithm for minimizing \mathop {\textrm{KL}}\limits (\cdot \!\;\Vert \; \!\pi )$\mathop {\textrm{KL}}\limits (\cdot \!\;\Vert \; \!\pi )$ over \mathcal {P}_\diamond $\mathcal {P}_\diamond$ based on accelerated gradient descent over \mathbb {R}^d$\mathbb {R}^d$. As a byproduct of our analysis, we obtain the first end-to-end analysis for gradient-based algorithms for MFVI.