Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient\n Estimator

Gradient estimation in models with discrete latent variables is a challenging\nproblem, because the simplest unbiased estimators tend to have high variance.\nTo counteract this, modern estimators either introduce bias, rely on multiple\nfunction evaluations, or use learned, input-dependent baselines. Thus, there is\na need for estimators that require minimal tuning, are computationally cheap,\nand have low mean squared error. In this paper, we show that the variance of\nthe straight-through variant of the popular Gumbel-Softmax estimator can be\nreduced through Rao-Blackwellization without increasing the number of function\nevaluations. This provably reduces the mean squared error. We empirically\ndemonstrate that this leads to variance reduction, faster convergence, and\ngenerally improved performance in two unsupervised latent variable models.\n