Source-linked AI summary
Reparameterizable Subset Sampling via Continuous Relaxations
Sang Michael Xie, Stefano Ermon
TL;DR
Subset sampling is difficult to optimize with low-variance reparameterized gradients, unlike single-item categorical sampling. The paper introduces a continuous top-k relaxation extending the Gumbel-max trick, and reports improved performance across feature selection, stochastic k-nearest neighbors, and parametric t-SNE.
Problem
Subset sampling without replacement typically lacks the low-variance reparameterization gradients available for single-item sampling.
Method
The paper uses a continuous top-k relaxation and an extension of the Gumbel-max trick to obtain approximate reparameterized gradients for subset distributions.
Results
The method improves performance in feature subset selection, stochastic neural k-nearest neighbors, and parametric t-SNE.
Takeaways & Limitations
The approach serves as a drop-in subset-sampling procedure for end-to-end training across multiple machine-learning tasks.
Takeaways & Limitations
The formulation reduces the number of free parameters to n −1, which restricts the model but makes the algorithm tractable.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper addresses subset sampling without replacement, where low-variance reparameterized gradients are generally unavailable. It introduces a continuous top-k relaxation based on an extension of the Gumbel-max trick and reports improvements across three subset-selection applications.
- Problem: The paper targets subset sampling without replacement, which lacks the low-variance reparameterization gradients available for single-item sampling.The Gumbel-max trick is not directly applicable because treating every possible subset as a category creates combinatorially many categories.
- Method: A continuous relaxation combines Gumbel perturbations with differentiable top-k selection to produce reparameterized subset samples for backpropagation.The method extends the Gumbel-max trick and relaxes top-k selection while preserving the sampling structure.
- Applications: The relaxation is presented as a drop-in replacement for subset-selection routines in feature explanations, stochastic neural k-nearest neighbors, and parametric t-SNE.The applications use subset sampling for instance-wise feature selection, neighbor selection, and direct comparison of local-neighbor identities.
- Results: The method improves performance on all three tasks using the original architectures and metrics.The paper frames this as end-to-end training with the subset procedure incorporated into the model.
- Preliminaries: Weighted reservoir sampling generates subset samples by assigning random keys to items and returning the top-k keys.The listed algorithm samples uniform noise, transforms it using item weights, and applies TopK.
2 Preliminaries
Weighted reservoir sampling defines ordered samples without replacement with probabilities proportional to item weights. Its random-key construction connects this sampling procedure to the Gumbel-max trick and motivates a differentiable top-k relaxation.
- Weighted Reservoir Sampling: A weighted reservoir sample is a sequence of k distinct one-hot selections drawn without replacement, with probabilities proportional to item weights.For k = 1, the distribution reduces to a softmax with logits log(w_i).
- Weighted Reservoir Sampling: The weighted reservoir algorithm assigns each item a random key and returns the k items with the largest keys.Efraimidis and Spirakis proved that this procedure has the desired weighted-reservoir distribution.
- Gumbel-max Trick: The Gumbel-max trick samples one item by adding Gumbel noise to log-weight logits and taking the maximum.The resulting categorical sample has probability proportional to the item weight.
- Gumbel-max Trick: The paper extends this reparameterization idea from max to top-k selection and relaxes top-k so gradients can flow through subset sampling.The deterministic parameter transformation and independent randomness are separated before the top-k operation is relaxed.
3 Reparameterizable Continuous Relaxation for Subset Sampling
The method generalizes Gumbel-max sampling to subsets by combining weighted-reservoir equivalence with differentiable top-k relaxations, yielding reparameterizable relaxed k-hot samples. The relaxation converges to exact subset samples as temperature approaches zero and preserves top-k ordering for temperatures t ≥ 1, while restricting the family of subset distributions for tractability.
- Setup: The subset model restricts the full distribution over k-element subsets from many free parameters to n − 1, gaining tractability at the cost of expressiveness.Subset probabilities are obtained by summing the probabilities of all ordered samples containing the same elements.
- Gumbel-max Extension: A Gumbel-max extension samples subsets by taking the top-k perturbed keys, because these keys are monotonic transformations of weighted-reservoir keys.This preserves the weighted-reservoir sampling distribution while enabling a reparameterized construction.
- Differentiable Top-k Procedures: Algorithm 2 replaces discrete top-k selection with a differentiable top-k relaxation that returns a relaxed k-hot vector with entries in [0,1] summing to k.The relaxation uses Gumbel-perturbed log weights, subset size k, and temperature t as inputs.
- Gumbel-max Extension: As t → 0, the relaxed procedure converges to exact samples from p(S|w), and intermediate relaxed outputs converge to weighted-reservoir samples when they represent ordered selections.This provides approximate sampling for both unordered subsets and ordered samples without replacement.
- Differentiable Top-k Procedures: The top-k relaxation preserves input-key ordering for any temperature t ≥ 1, although consistency is not guaranteed for all 0 < t < 1.For example, at t = 0.4 an input with equal top-k structure can produce an inconsistent output ordering.
- Differentiable Top-k Procedures: Higher temperatures reduce gradient variance but can introduce bias, so ordering consistency helps retain top-k information while balancing optimization trade-offs.Alternative top-k relaxations, including relaxed permutation matrices, can also be used in Algorithm 2.
4 Experiments
Experiments show that the proposed subset-sampling relaxation produces accurate samples and improves or matches established methods across feature explanation, stochastic k-nearest neighbors, and neighbor-preserving embeddings.
- Synthetic Experiments: The relaxation samples subsets with empirical distributions within 0.016 total variation distance of the target distribution across temperatures 0.1, 1, and 10.This agrees with Theorem 1, which predicts correct top-k samples even at higher temperatures.
- Model Explanations: RelaxSubSample improves L2X by up to 1% in post-hoc accuracy on IMDB sentiment explanations while changing only the sampling procedure.The evaluation uses k = 10 words or k = 1 sentence and deterministic highest-weight subsets at test time.
- Stochastic K-Nearest Neighbors: RelaxSubSample reaches 99.3% versus 99.4% for a cross-entropy CNN on MNIST and improves NeuralSort by 0.6% on CIFAR-10.The stochastic k-nearest-neighbor model uses k = 9 neighbors selected from m = 100 candidates.
- Stochastic K-Nearest Neighbors: The top-k implementation scales better than NeuralSort for larger candidate sets because it uses O(km) time and storage instead of O(m^2).Both methods have comparable runtimes for small m, but the difference grows with larger candidate sets.
- Stochastic Neighbor Embeddings: RSS-SNE replaces Student-t distributions with neighbor-ranking preservation and evaluates embeddings using trustworthiness and 1-NN classification on MNIST and 20 Newsgroups.The method directly compares sampled neighbor distributions, including ordered subsequences, in the original and embedded spaces.
5 Conclusion
The paper presents a relaxation for sampling from subset distributions that can be integrated into deep models trained with backpropagation, enabling end-to-end tuning of subset-dependent tasks.
- The proposed subset-sampling relaxation enables deep models to train with backpropagation while incorporating subset sampling.