Rao-Blackwellized Stochastic Gradients for Discrete Distributions
This incremental improvement addresses variance reduction for researchers and practitioners using stochastic gradient methods in machine learning with discrete distributions.
The paper tackles the problem of high variance in stochastic gradient estimators for expectations over large or infinite discrete sample spaces by introducing a Rao-Blackwellization technique that reduces variance without affecting bias, demonstrating improvements on semi-supervised classification and pixel attention tasks.
We wish to compute the gradient of an expectation over a finite or countably infinite sample space having $K \leq \infty$ categories. When $K$ is indeed infinite, or finite but very large, the relevant summation is intractable. Accordingly, various stochastic gradient estimators have been proposed. In this paper, we describe a technique that can be applied to reduce the variance of any such estimator, without changing its bias---in particular, unbiasedness is retained. We show that our technique is an instance of Rao-Blackwellization, and we demonstrate the improvement it yields on a semi-supervised classification problem and a pixel attention task.