[Paper Review] ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables

Mingzhang Yin, Yuguang Yue+1 more|arXiv (Cornell University)|May 4, 2019

Adversarial Robustness in Machine Learning17 citations

TL;DR

This paper proposes ARSM, an unbiased, low-variance gradient estimator for backpropagation through categorical variables by combining variable augmentation, REINFORCE, Rao-Blackwellization, and a novel variable-swapping mechanism to construct equivalent gradient expectations under the Dirichlet distribution. The method enables significant variance reduction via shared random numbers across these expectations, outperforming existing estimators in variational autoencoders and providing a 'try-and-see self-critic' baseline-free policy gradient method for discrete reinforcement learning.

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.

Motivation & Objective

• To address the challenge of high-variance gradient estimation in backpropagation through categorical variables in deep learning and reinforcement learning.
• To develop an unbiased gradient estimator that maintains low variance for multivariate categorical variables with C ≥ 2 categories.
• To eliminate the need for baseline estimation in discrete-action policy gradients by introducing a self-critic mechanism via pseudo-action rollouts.
• To unify and extend the ARM estimator, originally for binary variables, to the general categorical case (C ≥ 2).
• To provide a practical, reproducible method for training models with discrete latent variables and discrete control policies.

Proposed method

• The ARSM estimator uses variable augmentation to re-express the gradient as an expectation under a Dirichlet distribution.
• It applies the REINFORCE gradient estimator in the augmented space to obtain an initial unbiased gradient estimate.
• Rao-Blackwellization is applied to reduce variance by conditioning on sufficient statistics.
• Variable swapping is used to generate multiple equivalent but differently expressed gradient expectations under the Dirichlet distribution.
• The method performs a merge step by sharing common random numbers across these swapped expectations to further reduce variance.
• The final estimator is computed as a weighted average of these shared, equivalent expectations, achieving significant variance reduction.

Experimental results

Research questions

• RQ1 Can a low-variance, unbiased gradient estimator be constructed for categorical variables with C ≥ 2 using augmentation and variance reduction techniques?
• RQ2 How can variable swapping under the Dirichlet distribution generate equivalent gradient expectations that enable effective variance reduction?
• RQ3 Can ARSM eliminate the need for baseline estimation in discrete-action policy gradients by generating pseudo-actions and estimating their value functions?
• RQ4 How does ARSM compare to existing estimators in training categorical variational autoencoders and discrete reinforcement learning agents?
• RQ5 Does ARSM achieve performance close to the true gradient in univariate settings and state-of-the-art results in multivariate and complex settings?

Key findings

• ARSM closely approximates the true gradient in univariate settings, demonstrating high accuracy in gradient estimation.
• In categorical variational autoencoders, ARSM achieves state-of-the-art training and out-of-sample prediction performance compared to existing estimators.
• For discrete reinforcement learning, ARSM provides a 'try-and-see self-critic' method that removes the need for baseline estimation by generating and evaluating pseudo-actions.
• The method significantly reduces gradient variance through shared random numbers across swapped expectations, outperforming REINFORCE and other baselines in both synthetic and real-world tasks.
• ARSM with C ≥ 2 cannot be realized via antithetic sampling alone, necessitating the novel variable-swapping mechanism, which is essential for variance reduction.
• The proposed algorithm is reproducible, with code publicly available on GitHub for the ARSM gradient estimator and its application in RL and VAEs.