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

arXiv: https://arxiv.org/abs/1905.01413 · HTML: https://mlsanity.com/paper/1905.01413/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 4 May 2019.

## Built on by (most cited first, 10 of 10)

- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [Bayesian Optimization over Discrete and Mixed Spaces via Probabilistic Reparameterization](https://mlsanity.com/paper/2210.10199/) (2022)
- [Fantastic Rewards and How to Tame Them: A Case Study on Reward Learning for Task-oriented Dialogue Systems](https://mlsanity.com/paper/2302.10342/) (2023)
- [Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick](https://mlsanity.com/paper/2311.12569/) (2023)
- [Storchastic: A Framework for General Stochastic Automatic Differentiation](https://mlsanity.com/paper/2104.00428/) (2021)
- [ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables](https://mlsanity.com/paper/2105.14141/) (2021)
- [Coupled Gradient Estimators for Discrete Latent Variables](https://mlsanity.com/paper/2106.08056/) (2021)
- [Gradient Estimation for Binary Latent Variables via Gradient Variance Clipping](https://mlsanity.com/paper/2208.06124/) (2022)
- [CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator](https://mlsanity.com/paper/2110.14002/) (2021)
- [Accelerating Evolutionary Strategy via Rao-Blackwellizing Realization of Uncertain Input](https://mlsanity.com/paper/2608.02073/) (2026)

## Abstract (arXiv)

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.