Gradient Estimation with Discrete Stein Operators

TL;DR

The paper tackles high-variance gradient estimation for discrete distributions by introducing discrete Stein operators to build variance-reducing control variates for REINFORCE leave-one-out. It presents RODEO, a framework that augments RLOO with learnable Stein-based CVs and surrogate functions, enabling online variance minimization without extra evaluations of the target function $f$. The approach yields substantially reduced gradient variance and improved training objectives on binary and hierarchical Bernoulli VAEs, often outperforming state-of-the-art estimators at the same function-evaluation budget. This technique leverages neighboring-state information via Stein operators and online surrogate learning to provide practical, scalable variance reduction for discrete latent-variable models.

Abstract

Gradient estimation -- approximating the gradient of an expectation with respect to the parameters of a distribution -- is central to the solution of many machine learning problems. However, when the distribution is discrete, most common gradient estimators suffer from excessive variance. To improve the quality of gradient estimation, we introduce a variance reduction technique based on Stein operators for discrete distributions. We then use this technique to build flexible control variates for the REINFORCE leave-one-out estimator. Our control variates can be adapted online to minimize variance and do not require extra evaluations of the target function. In benchmark generative modeling tasks such as training binary variational autoencoders, our gradient estimator achieves substantially lower variance than state-of-the-art estimators with the same number of function evaluations.

Figures (12)

• Figure 1: Training binary latent VAEs with 2 or 3 $f$ evaluations per step on binarized MNIST.
• Figure 2: Training binary latent VAEs with Gaussian likelihoods, $K=2$, and non-binarized datasets.
• Figure 3: Training hierarchical binary latent VAEs with four stochastic layers on Fashion-MNIST. In this experiment, the estimators have very different behaviors towards the beginning and the end of training. We show this on the right by zooming into the first 50K steps of the gradient variance plot.
• Figure 4: Ablation study of impact of RODEO components: (a) Stein operators, (b) LOO baseline and global and local CVs, (c) surrogate functions on binary VAE training performance.
• Figure 5: Comparing the performance of RODEO and RLOO on more expensive ResNet VAE models trained on binarized MNIST with $K=2$. The middle plot shows the average wall clock performance over 5 trials.