Reparameterizable Subset Sampling via Continuous Relaxations

By Sang Michael Xie, Stefano Ermon. 2019

115 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Optimization & Training, Probabilistic & Bayesian Methods

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Optimization & Training

In Probabilistic & Bayesian Methods

Citations in the corpus

DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning (2021, neutral): “A related line of work focuses on stochastic k-subset selection in neural networks, e.g., see [26, 5, 39] and the references therein.”
Differentiable Top-k with Optimal Transport (2020, methodology): “Specifically, Kool et al. (2019) adapted such a trick to sample k elements from n choices, and Xie and Ermon (2019) further applied the trick to stochastic kNN, where neural networks are used to approximating the sorting operator.”
RetrieveGAN: Image Synthesis via Differentiable Patch Retrieval (2020, methodology): “Therefore, we adopt the weighted reservoir sampling strategy [39] to perform the subset sampling from the candidate patch sets.”
Differentiable Patch Selection for Image Recognition (2021, methodology): “Differentiable Top-K : The subset sampling operation can be implemented by expanding the Gumbel-Softmax trick [43] or based on optimal-transport formulations for ranking and sorting [44, 10], the latter of which was recently made significantly faster by Blondel et.”
Patch-based Knowledge Distillation for Lifelong Person Re-Identification (2022, methodology): “To draw K patches from an image, a popular solution is to sample K times from the same distribution without replacement [59, 65].”
Data Selection for Language Models via Importance Resampling (2023, methodology): “Sampling without replacement avoids choosing the same example multiple times, is more statistically efficient for importance resampling [23], and can be implemented efficiently with the Gumbel top-k trick [81; 38; 87; 39].”
Gradient Estimation with Stochastic Softmax Tricks (2020, background): “L2X [17] and SoftSub [84] are bespoke relaxations.”