mlsanityData to 14 Oct 2026

Gradient Estimation with Stochastic Softmax Tricks

arXiv 2006.08063 ↗First posted 15 Jun 2020stat.ML cs.LGcited by 67

Not on the map. Why?

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.

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

Builds on

No reference list yet.

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2021: 13; 2022: 17; 2023: 15; 2024: 8; 2025: 5; 2026: 9.

67 papers.

Sort:

By area:

  1. 2023Decision-Focused Learning: Foundations, State of the Art, Benchmark and Future Opportunitiescited by 130
  2. 2021DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learningcited by 115
  3. 2021Implicit MLE: Backpropagating Through Discrete Exponential Family Distributionscited by 98
  4. 2021A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learningcited by 82
  5. 2021GrooMeD-NMS: Grouped Mathematically Differentiable NMS for Monocular 3D Object Detectioncited by 66
  6. 2022Ordered Subgraph Aggregation Networkscited by 57
  7. 2022Differentiable DAG Samplingcited by 47
  8. 2022SIMPLE: A Gradient Estimator for k-Subset Samplingcited by 34
  9. 2022Unsupervised Learning for Combinatorial Optimization with Principled Objective Relaxationcited by 32
  10. 2022Backpropagation through Combinatorial Algorithms: Identity with Projection Workscited by 29
  11. 2023Graph Deep Learning for Time Series Forecastingcited by 26
  12. 2022Sparse Graph Learning from Spatiotemporal Time Seriescited by 22
38 more
  1. 2021Learning with Algorithmic Supervision via Continuous Relaxationscited by 21
  2. 2023Compact Neural Graphics Primitives with Learned Hash Probingcited by 19
  3. 2022Neural Set Function Extensions: Learning with Discrete Functions in High Dimensionscited by 14
  4. 2023Differentiable Clustering with Perturbed Spanning Forestscited by 14
  5. 2023Backpropagation of Unrolled Solvers with Folded Optimizationcited by 13
  6. 2021End-to-end learnable EEG channel selection for deep neural networks with Gumbel-softmaxcited by 12
  7. 2022Adaptive Perturbation-Based Gradient Estimation for Discrete Latent Variable Modelscited by 12
  8. 2024Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Modelcited by 11
  9. 2021Stochastic Iterative Graph Matchingcited by 10
  10. 2021Combinatorial Optimization for Panoptic Segmentation: A Fully Differentiable Approachcited by 10
  11. 2023DAG Learning on the Permutahedroncited by 9
  12. 2022Training Discrete Deep Generative Models via Gapped Straight-Through Estimatorcited by 8
  13. 2021Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spacescited by 7
  14. 2022Gradient Estimation for Binary Latent Variables via Gradient Variance Clippingcited by 7
  15. 2022Mirror Sinkhorn: Fast Online Optimization on Transport Polytopescited by 6
  16. 2024LPGD: A General Framework for Backpropagation through Embedded Optimization Layerscited by 6
  17. 2025Large (Vision) Language Models are Unsupervised In-Context Learnerscited by 6
  18. 2022Learning Group Importance using the Differentiable Hypergeometric Distributioncited by 5
  19. 2023Improving Dual-Encoder Training through Dynamic Indexes for Negative Miningcited by 5
  20. 2026Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusioncited by 5
  21. 2023Structured Dialogue Discourse Parsingcited by 4
  22. 2025Geometric Algorithms for Neural Combinatorial Optimization with Constraintscited by 4
  23. 2021Learning to Extend Program Graphs to Work-in-Progress Codecited by 3
  24. 2022Learning to Drop Out: An Adversarial Approach to Training Sequence VAEscited by 3
  25. 2023Contrastive Learning for Non-Local Graphs with Multi-Resolution Structural Viewscited by 3
  26. 2024Soft Condorcet Optimization for Ranking of General Agentscited by 3
  27. 2021Sparse Communication via Mixed Distributionscited by 2
  28. 2021Scaling Structured Inference with Randomizationcited by 2
  29. 2022Learning Discrete Structured Variational Auto-Encoder using Natural Evolution Strategiescited by 2
  30. 2022Dual Decomposition of Convex Optimization Layers for Consistent Attention in Medical Imagescited by 2
  31. 2023Efficient Learning of Discrete-Continuous Computation Graphscited by 2
  32. 2023SynJax: Structured Probability Distributions for JAXcited by 2
  33. 2024End-to-End Learning for Fair Multiobjective Optimization Under Uncertaintycited by 2
  34. 2025Information Geometry of Variational Bayescited by 2
  35. 2026SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradientscited by 2
  36. 2021Reconciling the Discrete-Continuous Divide: Towards a Mathematical Theory of Sparse Communicationcited by 1
  37. 2022Structured Recognition for Generative Models with Explaining Awaycited by 1
  38. 2023A Model-Agnostic Graph Neural Network for Integrating Local and Global Informationcited by 1

50 most cited of 67 All 67 (JSON)