# A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning

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

## Key facts

- Cited by 82 papers in the map (data through 14 Oct 2026).
- Builds on 93 of its 172 references that are in the map.
- First posted on arXiv 4 Oct 2021.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [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)
- [Hierarchical Neural Story Generation](https://mlsanity.com/paper/1805.04833/) (2018)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Pixel Recurrent Neural Networks](https://mlsanity.com/paper/1601.06759/) (2016)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable Neural Architecture Search](https://mlsanity.com/paper/1812.03443/) (2018)
- [vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations](https://mlsanity.com/paper/1910.05453/) (2019)
- [Improved Semantic Representations From Tree-Structured Long Short-Term Memory Networks](https://mlsanity.com/paper/1503.00075/) (2015)
- [SNAS: Stochastic Neural Architecture Search](https://mlsanity.com/paper/1812.09926/) (2018)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Emergence of Grounded Compositional Language in Multi-Agent Populations](https://mlsanity.com/paper/1703.04908/) (2017)
- [Learning to Explain: An Information-Theoretic Perspective on Model Interpretation](https://mlsanity.com/paper/1802.07814/) (2018)
- [Concrete Dropout](https://mlsanity.com/paper/1705.07832/) (2017)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Interpreting Graph Neural Networks for NLP With Differentiable Edge Masking](https://mlsanity.com/paper/2010.00577/) (2020)
- [FBNetV2: Differentiable Neural Architecture Search for Spatial and Channel Dimensions](https://mlsanity.com/paper/2004.05565/) (2020)
- [A* Sampling](https://mlsanity.com/paper/1411.0030/) (2014)
- [Modeling Point Clouds with Self-Attention and Gumbel Subset Sampling](https://mlsanity.com/paper/1904.03375/) (2019)
- [Convolutional Networks with Adaptive Inference Graphs](https://mlsanity.com/paper/1711.11503/) (2017)
- [Learning to Assemble Neural Module Tree Networks for Visual Grounding](https://mlsanity.com/paper/1812.03299/) (2018)
- [Emergence of Language with Multi-agent Games: Learning to Communicate with Sequences of Symbols](https://mlsanity.com/paper/1705.11192/) (2017)
- [Neural Variational Inference and Learning in Belief Networks](https://mlsanity.com/paper/1402.0030/) (2014)
- [Counterfactual Off-Policy Evaluation with Gumbel-Max Structural Causal Models](https://mlsanity.com/paper/1905.05824/) (2019)

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

- [Self-Improvement for Neural Combinatorial Optimization: Sample without Replacement, but Improvement](https://mlsanity.com/paper/2403.15180/) (2024)
- [Efficient-DLM: From Autoregressive to Diffusion Language Models, and Beyond in Speed](https://mlsanity.com/paper/2512.14067/) (2025)
- [Advancing Sequential Numerical Prediction in Autoregressive Models](https://mlsanity.com/paper/2505.13077/) (2025)
- [Residual Quantization with Implicit Neural Codebooks](https://mlsanity.com/paper/2401.14732/) (2024)
- [Sparse Graph Learning from Spatiotemporal Time Series](https://mlsanity.com/paper/2205.13492/) (2022)
- [Differentiable Agent-based Epidemiology](https://mlsanity.com/paper/2207.09714/) (2022)
- [Counterfactual Temporal Point Processes](https://mlsanity.com/paper/2111.07603/) (2021)
- [A skeletonization algorithm for gradient-based optimization](https://mlsanity.com/paper/2309.02527/) (2023)
- [Differentiable Clustering with Perturbed Spanning Forests](https://mlsanity.com/paper/2305.16358/) (2023)
- [Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick](https://mlsanity.com/paper/2311.12569/) (2023)
- [MEGA: Masked Generative Autoencoder for Human Mesh Recovery](https://mlsanity.com/paper/2405.18839/) (2024)
- [Explain My Surprise: Learning Efficient Long-Term Memory by Predicting Uncertain Outcomes](https://mlsanity.com/paper/2207.13649/) (2022)
- [Conditional computation in neural networks: principles and research trends](https://mlsanity.com/paper/2403.07965/) (2024)
- [Prototype-based Dataset Comparison](https://mlsanity.com/paper/2309.02401/) (2023)
- [Evaluation of Large Language Models via Coupled Token Generation](https://mlsanity.com/paper/2502.01754/) (2025)
- [GRAPES: Learning to Sample Graphs for Scalable Graph Neural Networks](https://mlsanity.com/paper/2310.03399/) (2023)
- [Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient](https://mlsanity.com/paper/2406.10576/) (2024)
- [Coupling without Communication and Drafter-Invariant Speculative Decoding](https://mlsanity.com/paper/2408.07978/) (2024)
- [A Review of Bayesian Uncertainty Quantification in Deep Probabilistic Image Segmentation](https://mlsanity.com/paper/2411.16370/) (2024)
- [Learning Group Importance using the Differentiable Hypergeometric Distribution](https://mlsanity.com/paper/2203.01629/) (2022)
- [Predictive Querying for Autoregressive Neural Sequence Models](https://mlsanity.com/paper/2210.06464/) (2022)
- [Improving Dual-Encoder Training through Dynamic Indexes for Negative Mining](https://mlsanity.com/paper/2303.15311/) (2023)
- [Understanding What Affects the Generalization Gap in Visual Reinforcement Learning: Theory and Empirical Evidence](https://mlsanity.com/paper/2402.02701/) (2024)
- [NLGR: Utilizing Neighbor Lists for Generative Rerank in Personalized Recommendation Systems](https://mlsanity.com/paper/2502.06097/) (2025)
- [SketchDNN: Joint Continuous-Discrete Diffusion for CAD Sketch Generation](https://mlsanity.com/paper/2507.11579/) (2025)
- [Fast Offline Policy Optimization for Large Scale Recommendation](https://mlsanity.com/paper/2208.05327/) (2022)
- [FROG: Fair Removal on Graphs](https://mlsanity.com/paper/2503.18197/) (2025)
- [A Simple Contrastive Framework Of Item Tokenization For Generative Recommendation](https://mlsanity.com/paper/2506.16683/) (2025)
- [Neural Image Abstraction Using Long Smoothing B-Splines](https://mlsanity.com/paper/2511.05360/) (2025)
- [Causal Temporal Reasoning for Markov Decision Processes](https://mlsanity.com/paper/2212.08712/) (2022)

## Abstract (arXiv)

The Gumbel-max trick is a method to draw a sample from a categorical distribution, given by its unnormalized (log-)probabilities. Over the past years, the machine learning community has proposed several extensions of this trick to facilitate, e.g., drawing multiple samples, sampling from structured domains, or gradient estimation for error backpropagation in neural network optimization. The goal of this survey article is to present background about the Gumbel-max trick, and to provide a structured overview of its extensions to ease algorithm selection. Moreover, it presents a comprehensive outline of (machine learning) literature in which Gumbel-based algorithms have been leveraged, reviews commonly-made design choices, and sketches a future perspective.