Variance Control for Black Box Variational Inference Using The James-Stein Estimator

TL;DR

This paper introduces a variance control method for Black Box Variational Inference using the James-Stein estimator, improving stability and reducing the need for fine-tuning in stochastic gradient updates.

Contribution

It proposes reframing gradient estimation as a multivariate problem and applies the James-Stein estimator for variance reduction, offering a simpler alternative to Rao-Blackwellization.

Findings

Achieves comparable or better model fit than Rao-Blackwellization.

Demonstrates stable convergence without extensive fine-tuning.

Provides a computationally efficient variance reduction method.

Abstract

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.