SIMPLE: A Gradient Estimator for $k$-Subset Sampling

TL;DR

The paper tackles end-to-end learning with discrete $k$-subset sampling by avoiding relaxations and introducing SIMPLE, a gradient estimator that uses exact conditional marginals on the backward pass while performing discrete forward sampling. It shows that the conditional marginals equal derivatives of the log of the exactly-$k$ subset probability and provides an $O(nk)$ algorithm to compute these probabilities, plus an $O(n)$ exact sampler. Empirically, SIMPLE exhibits lower bias and variance than state-of-the-art estimators across synthetic tests, discrete VAEs with exact ELBO computations, learning-to-explain tasks, and sparse regression, demonstrating practical improvements without relaxing the forward pass. The approach enables scalable, exact inference for $k$-subset distributions and yields a competitive, principled alternative to relaxation-based methods with significant impact on discrete latent variable modeling.

Abstract

$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.

Figures (5)

• Figure 1: A comparison of the bias and variance of the gradient estimators (left) and the average and standard deviation of the cosine distance of a single-sample gradient estimate to the exact gradient. We used the cosine distance, defined as ($1 -$ cosine similarity), in place of the euclidean distance as we only care about the direction of the gradient, not magnitude. The bias, variance and error were estimated using a sample of size 10,000. The details of this experiment are provided in \ref{['exp:synthetic']}.
• Figure 2: The problem setting considered in our paper. On the forward pass, a neural network $h_v$ outputs $\bm{\theta}$ parameterizing a discrete distribution over subsets of size $k$ of $n$ items, i.e., the $k$-subset distribution. We sample exactly, and efficiently, from this distribution, and feed the samples to a downstream neural network. On the backward pass, we approximate the true gradient by the product of the derivative of marginals and the gradient of the sample-wise loss.
• Figure 3: Bias and variance of Simple and Gumbel Softmax over $10$k samples
• Figure 4: ELBO against # of epochs. (Left) Comparison of Simple against different flavors of IMLE on the 10-subset DVAE, and (Right) against ST Gumbel Softmax on the 1-subset DVAE.
• Figure : $\mathop{\mathrm{\mathtt{PrExactlyk}}}\nolimits(\bm{\theta}, n, k)$

Theorems & Definitions (10)

• Theorem 1
• Proposition 1
• proof
• Proposition 2
• proof
• Proposition 3
• proof
• proof
• Proposition 4
• proof