Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space
This addresses the problem of approximating complex distributions for computational statistics, offering a novel theoretical and algorithmic approach that is incremental in building on existing variational inference methods.
The paper tackles mean-field variational inference by developing a polyhedral optimization framework in the Wasserstein space, providing approximation rates and an accelerated gradient descent algorithm for strongly log-concave and log-smooth distributions, with end-to-end analysis for gradient-based methods.
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 $π$ over $\mathbb{R}^d$ by a product measure $π^\star$. When $π$ is strongly log-concave and log-smooth, we provide (1) approximation rates certifying that $π^\star$ is close to the minimizer $π^\star_\diamond$ of the KL divergence over a \emph{polyhedral} set $\mathcal{P}_\diamond$, and (2) an algorithm for minimizing $\text{KL}(\cdot\|π)$ over $\mathcal{P}_\diamond$ based on accelerated gradient descent over $\R^d$. As a byproduct of our analysis, we obtain the first end-to-end analysis for gradient-based algorithms for MFVI.