Variance control for black box variational inference using the James–Stein estimator

Black box variational inference is a promising framework in a succession of recent efforts to make Variational Inference more “black box”. However, in its basic version it either fails to converge due to instability or requires some fine-tuning of the update steps prior to execution that hinders it from being completely general purpose. We propose a method for regulating its parameter updates by re-framing stochastic optimization as a multivariate estimation problem. Borrowing from estimation theory, we examine the properties of the James–Stein estimator as a replacement for the arithmetic mean of Monte Carlo estimates of the gradient of the evidence lower bound. Theoretical guarantees for its variance reduction properties are also given. We show through simulations that the proposed method provides relatively weaker variance reduction than Rao-Blackwellization, but offers a tradeoff of being simpler and requiring no prior analysis on the part of the user. Comparisons on benchmark datasets also demonstrate a consistent performance at par or better than the Rao-Blackwellized approach in terms of resulting model fit.