Fast Black-box Variational Inference through Stochastic Trust-Region Optimization

TL;DR

TrustVI is a new second-order black-box variational inference algorithm that converges faster and finds better solutions than existing methods by leveraging trust-region optimization and stochastic second-order information.

Contribution

We introduce TrustVI, a novel second-order trust-region algorithm for black-box variational inference with proven convergence and practical efficiency improvements.

Findings

TrustVI converges at least ten times faster than ADVI.

TrustVI finds better variational distributions than HFSGVI.

The algorithm demonstrates practical benefits of stochastic second-order optimization.

Abstract

We introduce TrustVI, a fast second-order algorithm for black-box variational inference based on trust-region optimization and the reparameterization trick. At each iteration, TrustVI proposes and assesses a step based on minibatches of draws from the variational distribution. The algorithm provably converges to a stationary point. We implemented TrustVI in the Stan framework and compared it to two alternatives: Automatic Differentiation Variational Inference (ADVI) and Hessian-free Stochastic Gradient Variational Inference (HFSGVI). The former is based on stochastic first-order optimization. The latter uses second-order information, but lacks convergence guarantees. TrustVI typically converged at least one order of magnitude faster than ADVI, demonstrating the value of stochastic second-order information. TrustVI often found substantially better variational distributions than HFSGVI, demonstrating that our convergence theory can matter in practice.