ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables
arXiv:1905.01413
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.
arXiv commentPublished in ICML 2019. We have updated Section 4.2 and the Appendix to reflect the improvements brought by fixing some bugs hidden in our original code. Please find the Errata in the authors' websites and check the updated code in Github
Cited by in corpus (10)
- Estimating Gradients for Discrete Random Variables by Sampling without Replacement
- DisARM: An Antithetic Gradient Estimator for Binary Latent Variables
- Storchastic: A Framework for General Stochastic Automatic Differentiation
- Coupled Gradient Estimators for Discrete Latent Variables
- Joint Stochastic Approximation and Its Application to Learning Discrete Latent Variable Models
- Adaptive Correlated Monte Carlo for Contextual Categorical Sequence Generation
- CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator
- Learning Discrete Energy-based Models via Auxiliary-variable Local Exploration
- Discrete Action On-Policy Learning with Action-Value Critic
- ARSM Gradient Estimator for Supervised Learning to Rank