# Gradient Estimation with Stochastic Softmax Tricks

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

## Key facts

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

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

- [Decision-Focused Learning: Foundations, State of the Art, Benchmark and Future Opportunities](https://mlsanity.com/paper/2307.13565/) (2023)
- [DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning](https://mlsanity.com/paper/2106.03760/) (2021)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [GrooMeD-NMS: Grouped Mathematically Differentiable NMS for Monocular 3D Object Detection](https://mlsanity.com/paper/2103.17202/) (2021)
- [Ordered Subgraph Aggregation Networks](https://mlsanity.com/paper/2206.11168/) (2022)
- [Differentiable DAG Sampling](https://mlsanity.com/paper/2203.08509/) (2022)
- [SIMPLE: A Gradient Estimator for k-Subset Sampling](https://mlsanity.com/paper/2210.01941/) (2022)
- [Unsupervised Learning for Combinatorial Optimization with Principled Objective Relaxation](https://mlsanity.com/paper/2207.05984/) (2022)
- [Backpropagation through Combinatorial Algorithms: Identity with Projection Works](https://mlsanity.com/paper/2205.15213/) (2022)
- [Graph Deep Learning for Time Series Forecasting](https://mlsanity.com/paper/2310.15978/) (2023)
- [Sparse Graph Learning from Spatiotemporal Time Series](https://mlsanity.com/paper/2205.13492/) (2022)
- [Learning with Algorithmic Supervision via Continuous Relaxations](https://mlsanity.com/paper/2110.05651/) (2021)
- [Compact Neural Graphics Primitives with Learned Hash Probing](https://mlsanity.com/paper/2312.17241/) (2023)
- [Neural Set Function Extensions: Learning with Discrete Functions in High Dimensions](https://mlsanity.com/paper/2208.04055/) (2022)
- [Differentiable Clustering with Perturbed Spanning Forests](https://mlsanity.com/paper/2305.16358/) (2023)
- [Backpropagation of Unrolled Solvers with Folded Optimization](https://mlsanity.com/paper/2301.12047/) (2023)
- [End-to-end learnable EEG channel selection for deep neural networks with Gumbel-softmax](https://mlsanity.com/paper/2102.09050/) (2021)
- [Adaptive Perturbation-Based Gradient Estimation for Discrete Latent Variable Models](https://mlsanity.com/paper/2209.04862/) (2022)
- [Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model](https://mlsanity.com/paper/2406.19905/) (2024)
- [Stochastic Iterative Graph Matching](https://mlsanity.com/paper/2106.02206/) (2021)
- [Combinatorial Optimization for Panoptic Segmentation: A Fully Differentiable Approach](https://mlsanity.com/paper/2106.03188/) (2021)
- [DAG Learning on the Permutahedron](https://mlsanity.com/paper/2301.11898/) (2023)
- [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)
- [Gradient Estimation for Binary Latent Variables via Gradient Variance Clipping](https://mlsanity.com/paper/2208.06124/) (2022)
- [Mirror Sinkhorn: Fast Online Optimization on Transport Polytopes](https://mlsanity.com/paper/2211.10420/) (2022)
- [LPGD: A General Framework for Backpropagation through Embedded Optimization Layers](https://mlsanity.com/paper/2407.05920/) (2024)
- [Large (Vision) Language Models are Unsupervised In-Context Learners](https://mlsanity.com/paper/2504.02349/) (2025)
- [Learning Group Importance using the Differentiable Hypergeometric Distribution](https://mlsanity.com/paper/2203.01629/) (2022)

## Abstract (arXiv)

The Gumbel-Max trick is the basis of many relaxed gradient estimators. These estimators are easy to implement and low variance, but the goal of scaling them comprehensively to large combinatorial distributions is still outstanding. Working within the perturbation model framework, we introduce stochastic softmax tricks, which generalize the Gumbel-Softmax trick to combinatorial spaces. Our framework is a unified perspective on existing relaxed estimators for perturbation models, and it contains many novel relaxations. We design structured relaxations for subset selection, spanning trees, arborescences, and others. When compared to less structured baselines, we find that stochastic softmax tricks can be used to train latent variable models that perform better and discover more latent structure.