Variance Control for Black Box Variational Inference Using The James-Stein Estimator
This addresses the problem of making variational inference more robust and general-purpose for practitioners, though it is incremental as it builds on existing methods with a tradeoff in variance reduction.
The paper tackles the instability and fine-tuning issues in Black Box Variational Inference by proposing a method using the James-Stein estimator to regulate parameter updates, resulting in consistent performance at par or better than existing approaches in model fit and convergence time on benchmark datasets.
Black Box Variational Inference is a promising framework in a succession of recent efforts to make Variational Inference more ``black box". However, in basic version it either fails to converge due to instability or requires some fine-tuning of the update steps prior to execution that hinder it from being completely general purpose. We propose a method for regulating its parameter updates by reframing stochastic gradient ascent as a multivariate estimation problem. 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. The proposed method provides relatively weaker variance reduction than Rao-Blackwellization, but offers a tradeoff of being simpler and requiring no fine tuning on the part of the analyst. Performance on benchmark datasets also demonstrate a consistent performance at par or better than the Rao-Blackwellized approach in terms of model fit and time to convergence.