Improved Gradient-Based Optimization Over Discrete Distributions

Abstract

In many applications we seek to maximize an expectation with respect to a distribution over discrete variables. Estimating gradients of such objectives with respect to the distribution parameters is a challenging problem. We analyze existing solutions including finite-difference (FD) estimators and continuous relaxation (CR) estimators in terms of bias and variance. We show that the commonly used Gumbel-Softmax estimator is biased and propose a simple method to reduce it. We also derive a simpler piece-wise linear continuous relaxation that also possesses reduced bias. We demonstrate empirically that reduced bias leads to a better performance in variational inference and on binary optimization tasks.

ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables

“…Here we consider both categorical latent variable models and policy optimization for discrete actions, which arise in a wide array of real-world applications. A number of unbiased estimators for backpropagating the gradient through discrete latent variables have been recently proposed Grathwohl et al, 2018;Yin & Zhou, 2019;Andriyash et al, 2018). However, they all mainly, if not exclusively, focus on the binary case (i.e., C = 2).…”