# Reparameterizable Subset Sampling via Continuous Relaxations

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

## Key facts

- Cited by 70 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 29 Jan 2019.

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

- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (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)
- [Differentiable Patch Selection for Image Recognition](https://mlsanity.com/paper/2104.03059/) (2021)
- [Diffusion Models With Learned Adaptive Noise](https://mlsanity.com/paper/2312.13236/) (2023)
- [Differentiable Subset Pruning of Transformer Heads](https://mlsanity.com/paper/2108.04657/) (2021)
- [Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws](https://mlsanity.com/paper/2410.11820/) (2024)
- [SIMPLE: A Gradient Estimator for k-Subset Sampling](https://mlsanity.com/paper/2210.01941/) (2022)
- [Backpropagation through Combinatorial Algorithms: Identity with Projection Works](https://mlsanity.com/paper/2205.15213/) (2022)
- [AdaProp: Learning Adaptive Propagation for Graph Neural Network based Knowledge Graph Reasoning](https://mlsanity.com/paper/2205.15319/) (2022)
- [Probabilistically Rewired Message-Passing Neural Networks](https://mlsanity.com/paper/2310.02156/) (2023)
- [DSKReG: Differentiable Sampling on Knowledge Graph for Recommendation with Relational GNN](https://mlsanity.com/paper/2108.11883/) (2021)
- [End-to-End Video Question Answering with Frame Scoring Mechanisms and Adaptive Sampling](https://mlsanity.com/paper/2407.15047/) (2024)
- [A Comprehensive Survey on Self-Interpretable Neural Networks](https://mlsanity.com/paper/2501.15638/) (2025)
- [On Strengthening and Defending Graph Reconstruction Attack with Markov Chain Approximation](https://mlsanity.com/paper/2306.09104/) (2023)
- [Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs](https://mlsanity.com/paper/2503.16870/) (2025)
- [Differentiable Multi-Target Causal Bayesian Experimental Design](https://mlsanity.com/paper/2302.10607/) (2023)
- [Probabilistic Graph Rewiring via Virtual Nodes](https://mlsanity.com/paper/2405.17311/) (2024)
- [Adaptive Perturbation-Based Gradient Estimation for Discrete Latent Variable Models](https://mlsanity.com/paper/2209.04862/) (2022)
- [Calibration of Machine Reading Systems at Scale](https://mlsanity.com/paper/2203.10623/) (2022)
- [Relational Conformal Prediction for Correlated Time Series](https://mlsanity.com/paper/2502.09443/) (2025)
- [Stochastic smoothing of the top-K calibrated hinge loss for deep imbalanced classification](https://mlsanity.com/paper/2202.02193/) (2022)
- [LLaVA-ReID: Selective Multi-image Questioner for Interactive Person Re-Identification](https://mlsanity.com/paper/2504.10174/) (2025)
- [Adversarial Contrastive Pre-training for Protein Sequences](https://mlsanity.com/paper/2102.00466/) (2021)
- [Leveraging Recursive Gumbel-Max Trick for Approximate Inference in Combinatorial Spaces](https://mlsanity.com/paper/2110.15072/) (2021)
- [Dynamic Encoder Size Based on Data-Driven Layer-wise Pruning for Speech Recognition](https://mlsanity.com/paper/2407.18930/) (2024)
- [Differentiable Mathematical Programming for Object-Centric Representation Learning](https://mlsanity.com/paper/2210.02159/) (2022)
- [Dynamic Probabilistic Pruning: A general framework for hardware-constrained pruning at different granularities](https://mlsanity.com/paper/2105.12686/) (2021)
- [Learning Group Importance using the Differentiable Hypergeometric Distribution](https://mlsanity.com/paper/2203.01629/) (2022)

## Abstract (arXiv)

Many machine learning tasks require sampling a subset of items from a collection based on a parameterized distribution. The Gumbel-softmax trick can be used to sample a single item, and allows for low-variance reparameterized gradients with respect to the parameters of the underlying distribution. However, stochastic optimization involving subset sampling is typically not reparameterizable. To overcome this limitation, we define a continuous relaxation of subset sampling that provides reparameterization gradients by generalizing the Gumbel-max trick. We use this approach to sample subsets of features in an instance-wise feature selection task for model interpretability, subsets of neighbors to implement a deep stochastic k-nearest neighbors model, and sub-sequences of neighbors to implement parametric t-SNE by directly comparing the identities of local neighbors. We improve performance in all these tasks by incorporating subset sampling in end-to-end training.