# SIMPLE: A Gradient Estimator for k-Subset Sampling

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 66 references that are in the map.
- First posted on arXiv 4 Oct 2022.

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

- [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)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Neural Probabilistic Logic Programming in DeepProbLog](https://mlsanity.com/paper/1907.08194/) (2019)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Rationalizing Neural Predictions](https://mlsanity.com/paper/1606.04155/) (2016)
- [A Semantic Loss Function for Deep Learning with Symbolic Knowledge](https://mlsanity.com/paper/1711.11157/) (2017)
- [Learning to Explain: An Information-Theoretic Perspective on Model Interpretation](https://mlsanity.com/paper/1802.07814/) (2018)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [Sparse Sequence-to-Sequence Models](https://mlsanity.com/paper/1905.05702/) (2019)
- [Adaptively Sparse Transformers](https://mlsanity.com/paper/1909.00015/) (2019)
- [A* Sampling](https://mlsanity.com/paper/1411.0030/) (2014)
- [Discrete Variational Autoencoders](https://mlsanity.com/paper/1609.02200/) (2016)
- [Learning with Fenchel-Young Losses](https://mlsanity.com/paper/1901.02324/) (2019)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [Semantic Probabilistic Layers for Neuro-Symbolic Learning](https://mlsanity.com/paper/2206.00426/) (2022)
- [Neural Nearest Neighbors Networks](https://mlsanity.com/paper/1810.12575/) (2018)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [PySINDy: A comprehensive Python package for robust sparse system identification](https://mlsanity.com/paper/2111.08481/) (2021)
- [Deep Graph Matching via Blackbox Differentiation of Combinatorial Solvers](https://mlsanity.com/paper/2003.11657/) (2020)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [Reparameterizable Subset Sampling via Continuous Relaxations](https://mlsanity.com/paper/1901.10517/) (2019)
- [Gradient Estimation with Stochastic Softmax Tricks](https://mlsanity.com/paper/2006.08063/) (2020)
- [SparseMAP: Differentiable Sparse Structured Inference](https://mlsanity.com/paper/1802.04223/) (2018)
- [Differentiable Dynamic Programming for Structured Prediction and Attention](https://mlsanity.com/paper/1802.03676/) (2018)
- [Optimizing Rank-based Metrics with Blackbox Differentiation](https://mlsanity.com/paper/1912.03500/) (2019)
- [Neuro-Symbolic Entropy Regularization](https://mlsanity.com/paper/2201.11250/) (2022)

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

- [Probabilistically Rewired Message-Passing Neural Networks](https://mlsanity.com/paper/2310.02156/) (2023)
- [Graph Deep Learning for Time Series Forecasting](https://mlsanity.com/paper/2310.15978/) (2023)
- [Over-Squashing in Graph Neural Networks: A Comprehensive survey](https://mlsanity.com/paper/2308.15568/) (2023)
- [Scaling Tractable Probabilistic Circuits: A Systems Perspective](https://mlsanity.com/paper/2406.00766/) (2024)
- [A Pseudo-Semantic Loss for Autoregressive Models with Logical Constraints](https://mlsanity.com/paper/2312.03905/) (2023)
- [Probabilistic Graph Rewiring via Virtual Nodes](https://mlsanity.com/paper/2405.17311/) (2024)
- [Semantic Strengthening of Neuro-Symbolic Learning](https://mlsanity.com/paper/2302.14207/) (2023)
- [Controllable Generation via Locally Constrained Resampling](https://mlsanity.com/paper/2410.13111/) (2024)
- [A Unified Approach to Count-Based Weakly-Supervised Learning](https://mlsanity.com/paper/2311.13718/) (2023)
- [GRAPES: Learning to Sample Graphs for Scalable Graph Neural Networks](https://mlsanity.com/paper/2310.03399/) (2023)
- [Collapsed Inference for Bayesian Deep Learning](https://mlsanity.com/paper/2306.09686/) (2023)
- [Adaptive Computation Modules: Granular Conditional Computation For Efficient Inference](https://mlsanity.com/paper/2312.10193/) (2023)
- [Graph State-Space Models and Latent Relational Inference](https://mlsanity.com/paper/2301.01741/) (2023)
- [Tractable Representation Learning with Probabilistic Circuits](https://mlsanity.com/paper/2507.04385/) (2025)
- [Differentiable Random Partition Models](https://mlsanity.com/paper/2305.16841/) (2023)
- [Semantic Loss Functions for Neuro-Symbolic Structured Prediction](https://mlsanity.com/paper/2405.07387/) (2024)
- [SPHINX: Structural Prediction using Hypergraph Inference Network](https://mlsanity.com/paper/2410.03208/) (2024)
- [Learning with Local Search MCMC Layers](https://mlsanity.com/paper/2505.14240/) (2025)
- [Revisiting Score Function Estimators for k-Subset Sampling](https://mlsanity.com/paper/2407.16058/) (2024)
- [FedGrAINS: Personalized SubGraph Federated Learning with Adaptive Neighbor Sampling](https://mlsanity.com/paper/2501.12592/) (2025)
- [Deep Generative Models with Hard Linear Equality Constraints](https://mlsanity.com/paper/2502.05416/) (2025)
- [Zero-Variance Gradients for Variational Autoencoders](https://mlsanity.com/paper/2508.03587/) (2025)
- [Differentiable Knapsack and Top-k Operators via Dynamic Programming](https://mlsanity.com/paper/2601.21775/) (2026)
- [ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts](https://mlsanity.com/paper/2606.01509/) (2026)
- [Learning to Simplify Spatial-Temporal Graphs in Gait Analysis](https://mlsanity.com/paper/2310.03396/) (2023)
- [A Complexity Map of Probabilistic Reasoning for Neurosymbolic Classification Techniques](https://mlsanity.com/paper/2404.08404/) (2024)
- [Edge Probability Graph Models Beyond Edge Independency: Concepts, Analyses, and Algorithms](https://mlsanity.com/paper/2405.16726/) (2024)
- [Learning Predictive Checklists with Probabilistic Logic Programming](https://mlsanity.com/paper/2411.16790/) (2024)
- [Discrete Subgraph Sampling for Interpretable Graph based Visual Question Answering](https://mlsanity.com/paper/2412.08263/) (2024)
- [Exploiting the Potential Supervision Information of Clean Samples in Partial Label Learning](https://mlsanity.com/paper/2505.09354/) (2025)

## Abstract (arXiv)

k-subset sampling is ubiquitous in machine learning, enabling regularization and interpretability through sparsity. The challenge lies in rendering k-subset sampling amenable to end-to-end learning. This has typically involved relaxing the reparameterized samples to allow for backpropagation, with the risk of introducing high bias and high variance. In this work, we fall back to discrete k-subset sampling on the forward pass. This is coupled with using the gradient with respect to the exact marginals, computed efficiently, as a proxy for the true gradient. We show that our gradient estimator, SIMPLE, exhibits lower bias and variance compared to state-of-the-art estimators, including the straight-through Gumbel estimator when k = 1. Empirical results show improved performance on learning to explain and sparse linear regression. We provide an algorithm for computing the exact ELBO for the k-subset distribution, obtaining significantly lower loss compared to SOTA.