Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference

TL;DR

This paper demonstrates that mean-field black-box variational inference converges nearly independently of dimension for certain target distributions, significantly improving efficiency over traditional methods.

Contribution

The authors prove dimension-independent convergence rates for BBVI with mean-field families, especially for strongly log-concave targets, and establish fundamental limits on gradient variance bounds.

Findings

Convergence rate is nearly independent of dimension for strongly log-concave targets.

For heavy-tailed families, convergence depends polynomially on dimension.

Gradient variance bounds are shown to be optimal under spectral Hessian bounds.

Abstract

We prove that, given a mean-field location-scale variational family, black-box variational inference (BBVI) with the reparametrization gradient converges at a rate that is nearly independent of explicit dimension dependence. Specifically, for a $d$-dimensional strongly log-concave and log-smooth target, the number of iterations for BBVI with a sub-Gaussian family to obtain a solution $\epsilon$-close to the global optimum has a dimension dependence of $\mathrm{O}(\log d)$. This is a significant improvement over the $\mathrm{O}(d)$ dependence of full-rank location-scale families. For heavy-tailed families, we prove a weaker $\mathrm{O}(d^{2/k})$ dependence, where $k$ is the number of finite moments of the family. Additionally, if the Hessian of the target log-density is constant, the complexity is free of any explicit dimension dependence. We also prove that our bound on the gradient variance, which is key to our result, cannot be improved using only spectral bounds on the Hessian of the target log-density.