Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator
arXiv:2010.04838
Abstract
Gradient estimation in models with discrete latent variables is a challenging problem, because the simplest unbiased estimators tend to have high variance. To counteract this, modern estimators either introduce bias, rely on multiple function evaluations, or use learned, input-dependent baselines. Thus, there is a need for estimators that require minimal tuning, are computationally cheap, and have low mean squared error. In this paper, we show that the variance of the straight-through variant of the popular Gumbel-Softmax estimator can be reduced through Rao-Blackwellization without increasing the number of function evaluations. This provably reduces the mean squared error. We empirically demonstrate that this leads to variance reduction, faster convergence, and generally improved performance in two unsupervised latent variable models.
References in corpus (9)
- Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
- The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
- Neural Variational Inference and Learning in Belief Networks
- Gradient Estimation Using Stochastic Computation Graphs
- REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
- Variational inference for Monte Carlo objectives
- Rao-Blackwellized Stochastic Gradients for Discrete Distributions
- Estimating Gradients for Discrete Random Variables by Sampling without Replacement
- Credit Assignment Techniques in Stochastic Computation Graphs