# Beyond ReinMax: Low-Variance Gradient Estimators for Discrete Latent Variables

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 15 references that are in the map.
- First posted on arXiv 9 Mar 2026.

## Builds on (most cited first, 9 of 9)

- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [A* Sampling](https://mlsanity.com/paper/1411.0030/) (2014)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator](https://mlsanity.com/paper/2010.04838/) (2020)
- [Bridging Discrete and Backpropagation: Straight-Through and Beyond](https://mlsanity.com/paper/2304.08612/) (2023)
- [Training Discrete Deep Generative Models via Gapped Straight-Through Estimator](https://mlsanity.com/paper/2206.07235/) (2022)

## Abstract (arXiv)

Machine learning models involving discrete latent variables require gradient estimators to facilitate backpropagation in a computationally efficient manner. The most recent addition to the Straight-Through family of estimators, ReinMax, can be viewed from a numerical ODE perspective as incorporating an approximation via Heun's method to reduce bias, but at the cost of high variance. In this work, we introduce the ReinMax-Rao and ReinMax-CV estimators which incorporate Rao-Blackwellisation and control variate techniques into ReinMax to reduce its variance. Our estimators demonstrate superior performance on training variational autoencoders with discrete latent spaces. Furthermore, we investigate the possibility of leveraging alternative numerical methods for constructing more accurate gradient approximations and present an alternative view of ReinMax from a simpler numerical integration perspective.