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

Abstract

To address the challenge of backpropagating the gradient through categorical variables, we propose the augment-REINFORCE-swap-merge (ARSM) gradient estimator that is unbiased and has low variance. ARSM first uses variable augmentation, REINFORCE, and Rao-Blackwellization to re-express the gradient as an expectation under the Dirichlet distribution, then uses variable swapping to construct differently expressed but equivalent expectations, and finally shares common random numbers between these expectations to achieve significant variance reduction. Experimental results show ARSM closely resembles the performance of the true gradient for optimization in univariate settings; outperforms existing estimators by a large margin when applied to categorical variational auto-encoders; and provides a "try-and-see self-critic" variance reduction method for discrete-action policy gradient, which removes the need of estimating baselines by generating a random number of pseudo actions and estimating their action-value functions.

Storchastic: A Framework for General Stochastic Automatic Differentiation

“…2. Storchastic should be flexible enough to allow implementing a wide range of reviewed gradient estimation methods, including score-function-based methods with complex sampling techniques [23,25,50] or control variates [14,46], and other methods such as measurevalued derivatives [16,38] and SPSA [45] which are missing AD implementations [34]. 3.…”

“…We show these conditions hold for a wide variety of gradient estimation methods for first order differentiation. For score-function related methods such as RELAX [14], ARM [50,51], MAPO [25] and unordered set esimator [23], the conditions also hold for any-order differentiation. Storchastic enables us to only have to prove these conditions locally, meaning that we can combine different gradient estimation methods for different stochastic nodes.…”