# Gradient Estimation with Discrete Stein Operators

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

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 73 references that are in the map.
- First posted on arXiv 19 Feb 2022.

## Builds on (most cited first, 30 of 33)

- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [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)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Neural Variational Inference and Learning in Belief Networks](https://mlsanity.com/paper/1402.0030/) (2014)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [Measuring Sample Quality with Stein's Method](https://mlsanity.com/paper/1506.03039/) (2015)
- [Oops I Took A Gradient: Scalable Sampling for Discrete Distributions](https://mlsanity.com/paper/2102.04509/) (2021)
- [The Optimal Reward Baseline for Gradient-Based Reinforcement Learning](https://mlsanity.com/paper/1301.2315/) (2013)
- [Variational inference for Monte Carlo objectives](https://mlsanity.com/paper/1602.06725/) (2016)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [VarGrad: A Low-Variance Gradient Estimator for Variational Inference](https://mlsanity.com/paper/2010.10436/) (2020)
- [Control functionals for Monte Carlo integration](https://mlsanity.com/paper/1410.2392/) (2014)
- [Minimum Stein Discrepancy Estimators](https://mlsanity.com/paper/1906.08283/) (2019)
- [MuProp: Unbiased Backpropagation for Stochastic Neural Networks](https://mlsanity.com/paper/1511.05176/) (2015)
- [Pathwise Derivatives Beyond the Reparameterization Trick](https://mlsanity.com/paper/1806.01851/) (2018)
- [Estimating Gradients for Discrete Random Variables by Sampling without Replacement](https://mlsanity.com/paper/2002.06043/) (2020)
- [DisARM: An Antithetic Gradient Estimator for Binary Latent Variables](https://mlsanity.com/paper/2006.10680/) (2020)
- [Probabilistic Best Subset Selection via Gradient-Based Optimization](https://mlsanity.com/paper/2006.06448/) (2020)
- [A unified view of likelihood ratio and reparameterization gradients](https://mlsanity.com/paper/2105.14900/) (2021)
- [The reproducing Stein kernel approach for post-hoc corrected sampling](https://mlsanity.com/paper/2001.09266/) (2020)
- [Scalable Control Variates for Monte Carlo Methods via Stochastic Optimization](https://mlsanity.com/paper/2006.07487/) (2020)
- [ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables](https://mlsanity.com/paper/2105.14141/) (2021)

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

- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Informed Correctors for Discrete Diffusion Models](https://mlsanity.com/paper/2407.21243/) (2024)
- [Bridging Discrete and Backpropagation: Straight-Through and Beyond](https://mlsanity.com/paper/2304.08612/) (2023)
- [Learning-Order Autoregressive Models with Application to Molecular Graph Generation](https://mlsanity.com/paper/2503.05979/) (2025)
- [Sparse Backpropagation for MoE Training](https://mlsanity.com/paper/2310.00811/) (2023)
- [Generalised Bayesian Inference for Discrete Intractable Likelihood](https://mlsanity.com/paper/2206.08420/) (2022)
- [Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick](https://mlsanity.com/paper/2311.12569/) (2023)
- [Discrete Neural Flow Samplers with Locally Equivariant Transformer](https://mlsanity.com/paper/2505.17741/) (2025)
- [Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model](https://mlsanity.com/paper/2406.19905/) (2024)
- [A kernel Stein test of goodness of fit for sequential models](https://mlsanity.com/paper/2210.10741/) (2022)
- [Gradient Estimation for Binary Latent Variables via Gradient Variance Clipping](https://mlsanity.com/paper/2208.06124/) (2022)
- [Meta-learning Control Variates: Variance Reduction with Limited Data](https://mlsanity.com/paper/2303.04756/) (2023)
- [Adaptive Instrument Design for Indirect Experiments](https://mlsanity.com/paper/2312.02438/) (2023)
- [On Divergence Measures for Training GFlowNets](https://mlsanity.com/paper/2410.09355/) (2024)
- [When can Regression-Adjusted Control Variates Help? Rare Events, Sobolev Embedding and Minimax Optimality](https://mlsanity.com/paper/2305.16527/) (2023)
- [Bayesian Numerical Integration with Neural Networks](https://mlsanity.com/paper/2305.13248/) (2023)
- [Secrets of GFlowNets' Learning Behavior: A Theoretical Study](https://mlsanity.com/paper/2505.02035/) (2025)
- [DBsurf: A Discrepancy Based Method for Discrete Stochastic Gradient Estimation](https://mlsanity.com/paper/2309.03974/) (2023)
- [What to Do When Your Discrete Optimization Is the Size of a Neural Network?](https://mlsanity.com/paper/2402.10339/) (2024)
- [Scalable Monte Carlo for Bayesian Learning](https://mlsanity.com/paper/2407.12751/) (2024)
- [Variational phylogenetic inference with products over bipartitions](https://mlsanity.com/paper/2502.15110/) (2025)
- [BayesSum: Bayesian Quadrature in Discrete Spaces](https://mlsanity.com/paper/2512.16105/) (2025)
- [Variance-Optimal Control Variates for Learning with Black-box Feedback](https://mlsanity.com/paper/2610.04766/) (2026)

## Abstract (arXiv)

Gradient estimation -- approximating the gradient of an expectation with respect to the parameters of a distribution -- is central to the solution of many machine learning problems. However, when the distribution is discrete, most common gradient estimators suffer from excessive variance. To improve the quality of gradient estimation, we introduce a variance reduction technique based on Stein operators for discrete distributions. We then use this technique to build flexible control variates for the REINFORCE leave-one-out estimator. Our control variates can be adapted online to minimize variance and do not require extra evaluations of the target function. In benchmark generative modeling tasks such as training binary variational autoencoders, our gradient estimator achieves substantially lower variance than state-of-the-art estimators with the same number of function evaluations.