# Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator

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

## Key facts

- Cited by 37 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Oct 2020.

## Built on by (most cited first, 30 of 37)

- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [Model Agnostic Sample Reweighting for Out-of-Distribution Learning](https://mlsanity.com/paper/2301.09819/) (2023)
- [Bridging Discrete and Backpropagation: Straight-Through and Beyond](https://mlsanity.com/paper/2304.08612/) (2023)
- [Stochastic RAG: End-to-End Retrieval-Augmented Generation through Expected Utility Maximization](https://mlsanity.com/paper/2405.02816/) (2024)
- [Variational Information Pursuit for Interpretable Predictions](https://mlsanity.com/paper/2302.02876/) (2023)
- [Structured Sparse Transition Matrices to Enable State Tracking in State-Space Models](https://mlsanity.com/paper/2509.22284/) (2025)
- [Neural Set Function Extensions: Learning with Discrete Functions in High Dimensions](https://mlsanity.com/paper/2208.04055/) (2022)
- [Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick](https://mlsanity.com/paper/2311.12569/) (2023)
- [UDC: Unified DNAS for Compressible TinyML Models](https://mlsanity.com/paper/2201.05842/) (2022)
- [Coupled Gradient Estimators for Discrete Latent Variables](https://mlsanity.com/paper/2106.08056/) (2021)
- [Training Discrete Deep Generative Models via Gapped Straight-Through Estimator](https://mlsanity.com/paper/2206.07235/) (2022)
- [Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces](https://mlsanity.com/paper/2110.15072/) (2021)
- [Variational Sparse Coding with Learned Thresholding](https://mlsanity.com/paper/2205.03665/) (2022)
- [Bayesian Nonparametrics for Offline Skill Discovery](https://mlsanity.com/paper/2202.04675/) (2022)
- [Multi-Sample Training for Neural Image Compression](https://mlsanity.com/paper/2209.13834/) (2022)
- [Scalable Out-of-distribution Robustness in the Presence of Unobserved Confounders](https://mlsanity.com/paper/2411.19923/) (2024)
- [CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator](https://mlsanity.com/paper/2110.14002/) (2021)
- [High-Dimensional Learning Dynamics of Quantized Models with Straight-Through Estimator](https://mlsanity.com/paper/2510.10693/) (2025)
- [Soft Actor-Critic With Integer Actions](https://mlsanity.com/paper/2109.08512/) (2021)
- [VQ-AR: Vector Quantized Autoregressive Probabilistic Time Series Forecasting](https://mlsanity.com/paper/2205.15894/) (2022)
- [Learning to Drop Out: An Adversarial Approach to Training Sequence VAEs](https://mlsanity.com/paper/2209.12590/) (2022)
- [Revisiting the Gumbel-Softmax in MADDPG](https://mlsanity.com/paper/2302.11793/) (2023)
- [Implicit Bilevel Optimization: Differentiating through Bilevel Optimization Programming](https://mlsanity.com/paper/2302.14473/) (2023)
- [VQSynery: Robust Drug Synergy Prediction With Vector Quantization Mechanism](https://mlsanity.com/paper/2403.03089/) (2024)
- [Beyond Discreteness: Sample Complexity Analysis of Straight-Through Estimator for 1-bit Quantization](https://mlsanity.com/paper/2505.18113/) (2025)
- [TRIDENT: Breaking the Hybrid-Safety-Physics Coupling for Provably Safe Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2606.18308/) (2026)
- [Bayesian Unsupervised Disentanglement of Anatomy and Geometry for Deep Groupwise Image Registration](https://mlsanity.com/paper/2401.02141/) (2024)
- [Rao-Blackwellised Reparameterisation Gradients](https://mlsanity.com/paper/2506.07687/) (2025)
- [Gradient estimators for parameter inference in discrete stochastic kinetic models](https://mlsanity.com/paper/2604.02121/) (2026)
- [VGAS: Variance-Reduced Guidance and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion](https://mlsanity.com/paper/2608.26585/) (2026)

## Abstract (arXiv)

Gradient estimation in models with discrete latent variables is a challenging problem, because the simplest unbiased estimators tend to have high variance. To counteract this, modern estimators either introduce bias, rely on multiple function evaluations, or use learned, input-dependent baselines. Thus, there is a need for estimators that require minimal tuning, are computationally cheap, and have low mean squared error. In this paper, we show that the variance of the straight-through variant of the popular Gumbel-Softmax estimator can be reduced through Rao-Blackwellization without increasing the number of function evaluations. This provably reduces the mean squared error. We empirically demonstrate that this leads to variance reduction, faster convergence, and generally improved performance in two unsupervised latent variable models.