SIMPLE: A Gradient Estimator for k-subset Sampling

Kareem Ahmed * 1 Zhe Zeng * 1 Mathias Niepert 2 Guy Van den Broeck 1

Abstract

training. This provided the impetus for other, low-variance,
gradient estimators, chief among them are those based on
the reparameterization trick, which allows for biased, but
low-variance gradient estimates. The reparameterization
trick, however, does not allow for a direct application todiscrete distributions , thereby prompting continuous relaxations, such as the Gumbel-softmax trick (Jang et al., 2017),
or the Concrete distribution (Maddison et al., 2017): continuous relaxations that allow for reparameterized gradientsw.r.t. the parameters of a categorical distribution, precluding
the k-subset distribution. Reparameterizable subset sampling (Xie & Ermon, 2019) generalizes Gumbel-softmax
to k-subsets which renders k-subset sampling amenable tobackpropagation but introduces bias in the learning by using
relaxed samples.

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 thereparameterized samples to allow for backpropagation, with the risk of introducing high bias andhigh variance. In this work, we fall back to discrete k-subset sampling on the forward pass. Thisis coupled with using the gradient with respect
to the exact marginals, computed efficiently, as aproxy for the true gradient. We show that our gradient estimator, S IMPLE, exhibits lower bias andvariance compared to state-of-the-art estimators,
including the straight-through Gumbel estimatorwhen k = 1. Empirical results show improved
performance on learning to explain and sparselinear regression. We give an algorithm computing the exact ELBO for the k-subset distribution,
obtaining significantly lower loss than SOTA.

In this paper, we set out with the goal of avoiding all suchrelaxations. Instead, we fall back to discrete sampling on
the forward pass. On the backward pass, we reparameterize
the gradient of the loss function with respect to the samples as a function of the exact marginals of the k-subset
distribution. Computing the exact conditional marginals
is, in general, intractable (Roth, 1996). We give an efficient algorithm for computing the k-subset probability, andshow that the conditional marginals correspond to partial
derivatives, and are therefore tractable for the k-subset distribution. We show that our proposed gradient estimator for
the k-subset distribution, coined S IMPLE, is reminiscent of
the straight-through (ST) Gumbel estimator when k = 1,
with the gradients taken w.r.t. the unperturbed marginals.
We empirically demonstrate that S IMPLE exhibits lower bias
and variance compared to other known gradient estimators,
including the ST Gumbel estimator in the case k = 1.

1. Introduction
k-subset sampling, sampling a subset of size k of n variables, is omnipresent in machine learning. It lies at thecore of many fundamental problems that rely upon learning
sparse features representations of input data, All such tasks
involve optimizing an expectation of an objective function
with respect to a latent discrete distribution parameterized
by a neural network, which are often assumed intractable.
Score-function estimators offer a cloyingly simple solution:
rewrite the gradient of the expectation as an expectation of
the gradient, which can subsequently be estimated using
a finite number of samples offering an unbiased estimateof the gradient. Simple as it is, score-function estimators
suffer from very high variance which can interfere withWe include an experiment on the task of learning to explain(L2X) using the B EER A DVOCATE dataset (McAuley et al.,2012), where the goal is to select the subset of words thatbest explains the model’s classification of a user’s review.
We also include an experiment on the task of stochastic
sparse linear regression, where the goal is to learn the best
sparse model, and show that we are able to recover theKuramoto–Sivashinsky equation. Finally, we develop an
efficient computation for the calculation of the exact variational evidence lower bound (ELBO) for the k-subset distribution, which when used in conjunction with S IMPLE leads
to state-of-the-art discrete sparse VAE learning.

*
Equal contribution 1 Dept. of Computer Science, UCLA
Dept. of Computer Science, Stuttgart University. Correspondence to: Kareem Ahmed <ahmedk@cs.ucla.edu>, Zhe Zeng<zhezeng@cs.ucla.edu>.
2

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1SIMPLE: A Gradient Estimator for k-subset Sampling
0.6

Bias

0.15

Variance
0.4

0.10

0.2

0.05

0.0

0.00
Exact

SoftSub

IMLE

SFE

SIMPLE-FSIMPLE-B SIMPLE(a) The problem setting considered in our paper. On the forward
pass we sample exactly from a k-subset distribution parameterized
by a neural network. On the backward pass, we approximate thetrue gradient by the product of the derivative of marginals and theloss gradient of the sample-wise loss.

(b) A comparison of the bias and variance of the gradientestimators (left). 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.

2. Problem Statement and Motivation

general and subsumes many settings. Different choices ofmappings hv and fu , and sample-wise loss ℓ define various
tasks. Learning then requires computing the gradient of L
w.r.t. ω = (v, u). The gradient of L w.r.t. u is given by

We consider models described by the equationsPθ = hv (x), z ∼ pθ (z | i zi = k), ŷ = fu (z, x)(1)∇u L(x, y; ω) = Ez∼pθ (z|Pi zi =k) [∂u fu (z, x)⊤ ∇ŷ ℓ(ŷ, y)],(4)where ŷ = fu (z, x) is the decoding of a latent sample z.
Furthermore, the gradient of L w.r.t. v is given by

where x ∈ X and ŷ ∈ Y denote feature inputs and target
outputs, respectively, hv : X → Θ and fu : Z × X → Y
are smooth, parameterized maps and θ are logits inducing a
distribution over the latent binary vector z defined aspθ (z) =nYpθi (zi ), where pθi (zi ) = σ(θi ).

∇v L(x, y; ω) = ∂v hv (x)⊤ ∇θ L(x, y; ω)(2)where the loss’s gradient w.r.t. the encoder being∇θ L(x, y; ω) := ∇θ Ez∼pθ (z|Pi zi =k) [ℓ(fu (z, x), ŷ)].

i=1

One challenge lies in computing the expectation in (3) and(4), which has no known closed-form solutionP
requiring aMonte-Carlo estimate by sampling from pθ (z | i zi = k).

The goal of our stochastic latent layer is not to simply sample from pθ (z), which would yield samples with a Hamming weight between 0 and n (i.e., with an arbitrary numberof ones). Instead, we are interested in sampling from the
distribution restricted to samples with a Hamming weightof k, for any given k. That is, we are interested
in samplingP
from the conditional distribution pθ (z | i zi = k).

A second, more substantial hurdle lies in computing∇θ L(x, y; ω) in (5) due to the non-differentiable natureof discrete sampling. One could rewrite ∇θ L(x, y; ω) as
Ez∼pθ (z|P zi =k) [ℓ(fu (z, x), y) · ∇θ log pθ (z |iConditioning the distribution pθ (z) on this k-subset constraint introduces intricate dependencies between each of
the zi ’s. The probability of sampling a given k-subset z isPPPpθ (z | i zi = k) = pθ (z)/pθ ( i zi = k) · J i zi = kKPi zi = k)]which is known as the REINFORCE estimator (Williams,
1992), or the score function estimator (SFE). It is typically
avoided due to its notoriously high variance, despite its apparent simplicity. Instead, typical approaches (Xie & Ermon,
2019; Plötz & Roth, 2018) reparameterize the samples as adeterministic transformation of the parameters, and independent standard Gumbel noise, and relaxing the deterministictransformation, the top-k function, to allow for backprop.

where J·K denotes the indicator function. In other words,
the probabilityof sampling each k-subset is re-normalizedPby pθ ( i zi = k) – the probability of sampling exactly
k items from the unconstrained
induced by
P distribution Ph.
Thequantity
p
(
z
=
k)
=
encodervθiiz pθ (z) ·PJ i zi = kK appears to be intractable. We show that not tobe the case, providing a tractable algorithm for computing it.

3. SIMPLE: Subset Implicit Likelihood
Our goal is to build a gradient estimator for ∇θ L(x, y; ω).
We envision a hypothetical sampling-free architecture,
where the downstream neural network fuPis a function of
the marginals, µ := µ(θ) := {pθ (zj | i zi = k)}nj=1 ,instead of a discrete sample z, resulting in a loss Lm s.t.

Given a set of samples D, we are concerned with learning
the parameters ω = (v, u) of the architecture in (1) throughminimizing the training error L, which is the expected loss:

L(x, y; ω) = Ez∼pθ (z|Pi zi =k) [ℓ(fu (z, x), y)]

(5)

(3)

∇θ Lm (x, y; ω) = ∂θ µ(θ)⊤ ∇µ ℓm (fu (µ, x), y).

where θ = hv (x) and ℓ : Y × Y → R+ is a point-wiseloss function. This formulation, illustrated in Figure 1a, is

When the marginals µ(θ) can be efficiently computed anddifferentiated, such a hypothetical pipeline can be trained
2

SIMPLE: A Gradient Estimator for k-subset Sampling

end-to-end. Furthermore, Domke (2010) observed that,for an arbitrary loss function ℓm defined on the marginals,
the Jacobian of the marginals w.r.t. the logits is symmetric. Consequently, computing the gradient of the loss w.r.t.
the logits, ∇θ Lm (x, y; ω), reduces to computing the directional derivative, or the Jacobian-vector product, of themarginals w.r.t. the logits in the direction of the gradientof the loss. This offers an alluring opportunity: the conditional marginals characterize the probability of each zi inthe sample, and could be thought of as a differentiable proxy
for the samples. Specifically, by reparameterizing z as afunction of the conditional marginal µ under approximation∂µ z ≈ I as proposed by Niepert et al. (2021), and using
the straight-through estimator for the gradient of the sample
w.r.t. the marginals on the backward pass, we approximate∇θ L(x, y; ω) ≈ ∂θ µ(θ)∇z L(x, y; ω),

Discrete 10-Subset VAE

225

0

20

40

Gumbel Softmax
SIMPLE260

SoG I-MLE
Gumbel I-MLE
SIMPLE200
175

Discrete 1-Subset VAE
ELBO

250

ELBO

Algorithm 1 PrExactlyk(θ, n, k)
Algorithm 2 Sample(θ, n, k)
Input: The logits θ of the distribution, the num- Input: The logits θ of the distribution, the number of variables
ber of variables n, and the subset size
P n, and the subset size k
P k
Output: pθ ( i zi = k)
Output: z = (z1 , . . . , zn ) ∼ pθ (z | i zi = k)Pisample = [ ], j = k// a[i, j] = pθ ( m=1 zm = j) for all i, j
for i = n to 1 do
initialize a to be 0Peverywhere0// cf. proof of Prop. 2
a[0, 0] = 1 // pθ ( m=1 zm = 0) = 1
p = a[i − 1, j − 1]for i = 1 to n dozi ∼ Bernoulli(p · pθi (zi = 1)/a[i, j])
for j = 0 to k do// Pick next state based on value of sample// cf. constructive proof of Prop. 1
if
zi = 1 then j = j − 1
a[i, j] = a[i − 1, j] · pθi (zi = 0)
sample.append(zi )+ a[i − 1, j − 1] · pθi (zi = 1)
return sample
return a[n, k]

240
60

Epochs

80

100

0

20

40

60

Epochs80

100

Figure 2: ELBO against epochs. (Left) S IMPLE against
different flavors of IMLE on the 10-subset DVAE, and(Right) ST Gumbel Softmax on the 1-subset DVAE.

retains all the terms referencing that variable, and drops allother terms; this corresponds exactly to the unnormalized
conditional marginals. By taking the derivative of the logprobability, this introduces the k-subset probability in the
denominator, leading to conditional marginals. Intuitively,
the rate of change of the k-subset probability w.r.t. a variable
only depends on that variable through its length-k subsets.
P
Theorem 1. Let pθ ( j zj = k) be the probability ofexactly-k of the unconstrained distribution parameterized
by logits θ. For every Zi , its conditional marginal is


P
P
∂
p θ zi | j zj = k =log pθ ( j zj = k),∂αiwhere αi := log pθi (zi ) denotes the log marginals.

(6)

To establish the tractability of the above computation of theconditional marginals, we needP
to show that the probabilityof the exactly-k constraint pθ ( i zi = k) is tractable
P
Proposition 1. The probability pθ ( i zi = k) of sampling
exactly k items from the unconstrained distribution pθ (z)over n items can be computed exactly in time O(nk).

where the directional derivative of the marginals can betaken along any downstream gradient, rendering the whole
pipeline end-to-end learnable despite presence of sampling.
Now, estimating the gradient of the loss w.r.t. the parameters can be thought of as decomposing into two subproblems: (P1) Computing the derivatives of conditional
marginals ∂θ µ(θ), which requires the computation of theconditional marginals, and (P2) Computing the gradientof the loss w.r.t. the samples ∇z L(x, y; ω) using samplewise loss, which requires drawing exact samples. These
two problems are complicated by conditioning on the ksubset constraint, which introduces intricate dependencies
to the distribution, and is infeasible to solve naively, e.g. byenumeration. We will show simple, efficient, and exact solutions to each problem, at the heart of which is the insightthat we need not care about the variables’ order, only theirsum, introducing symmetries that simplify the problem.

Pn
Since we have a closed-form pθ ( i zi = k), which
P allows
us to compute conditional marginals pθ (zi | j zj = k)by Theorem 1 via auto-differentiation. This further allows
the computation of the P
derivatives of conditional marginals∂θ µ(θ)i = ∂θ pθ (zi | j zj = k) to be amenable to autodifferentiation, solving problem (P1) exactly and efficiently.
3.2. Gradients of Loss w.r.t. Samples
What remains is to estimate the value of the loss, which
necessitates faithful sampling from the k-subset distribution.
Exact k-subset Sampling Next we show how
Pto sample
exactly from the k-subset distribution pθ (z | i zi = k).
We start by sampling the variables in reverse order, thatis, we sample zn through z1 . The main intuition being
that, having sampled (zn , zn−1 , · · · , zi+1 ) with a Hamming
weight of k −j, we sample Zi with a probability of choosingk − j of n − 1 variables and the n-th variable given that wechoose k − j + 1 of n variables, formalized in Proposition 2.

3.1. Derivatives of Conditional Marginals
In many probabilistic models, marginal inference is a #Phard problem (Roth, 1996). That is not the case for the
k-subset distribution. We notice the conditional marginals
correspond to the partial derivatives of the log-probability of
the k-subset constraint. To see this, note that the derivativeof a multi-linear function with respect to a single variable
3

SIMPLE: A Gradient Estimator for k-subset Sampling

Table 1: Results for the three different aspects with k = 10(top) and for the aroma aspect for different k (bottom).

4. Experiments
4.1. Synthetic Experiments
We carried out a series of experiments with a 5-subset distribution, and a latent space of dimension 10. We set theloss to L(θ) = Ez∼pθ (z|P zi =k) [∥z − b∥2 ], where b is

MethodS IMPLE (Ours)
L2X (t = 0.1)
SoftSub (t = 0.5)
I-MLE (τ = 30)ithe groundtruth logits sampled from N(0, I). Such a distribution is tractable: we only have 10
5 = 252 k-subsets,allowing the exact gradient to be computed in closed form.

MethodS IMPLE (Ours)
L2X (t = 0.1)
SoftSub (t = 0.5)
I-MLE (τ = 30)

We compare against: exact, which denotes the exact gradient; SoftSub (Xie & Ermon, 2019), which uses an extensionof the Gumbel-Softmax trick to sample relaxed k-subsets
on the forward pass; I-MLE, which denotes the IMLE gradient estimator (Niepert et al., 2021), where approximate
samples are obtained using perturb-and-map (PAM) on theforward pass, approximating the marginals using samples
on the backward pass; and score function estimator SFE.

Appearance

PalateTaste

Test MSE

Precision

Test MSE

Precision

Test MSE

Precision

2.35 ± 0.28
10.70 ± 4.82
2.48 ± 0.10
2.51 ± 0.05

66.81 ± 7.56
30.02 ± 15.82
52.86 ± 7.08
65.47 ± 4.95

2.68 ± 0.06
6.70 ± 0.63
2.94 ± 0.08
2.96 ± 0.04

44.78 ± 2.75
50.39 ± 13.58
39.17 ± 3.17
40.73 ± 3.15

2.11 ± 0.02
6.92 ± 1.61
2.18 ± 0.10
2.38 ± 0.04

42.31 ± 0.61
32.23 ± 4.92
41.98 ± 1.42
41.38 ± 1.55

k=5

k = 10k = 15

Test MSE

Precision

Test MSE

Precision

Test MSE

Precision

2.27 ± 0.05
5.75 ± 0.30
2.57 ± 0.12
2.62 ± 0.05

57.30 ± 3.04
33.63 ± 6.91
54.06 ± 6.29
54.76 ± 2.50

2.23 ± 0.03
6.68 ± 1.08
2.67 ± 0.14
2.71 ± 0.10

47.17 ± 2.11
26.65 ± 9.39
44.44 ± 2.27
47.98 ± 2.26

3.20 ± 0.04
7.71 ± 0.64
2.52 ± 0.07
2.91 ± 0.18

53.18 ± 1.09
23.49 ± 10.93
37.78 ± 1.71
39.56 ± 2.07

variables model a probability distribution over k-subsets.
In prior work, the KL-divergence was approximated using the unconditional marginals, obtained simply througha Softmax layer. Instead we show that the KL-divergence
between the k-subset distribution and the uniform distribution can be computed exactly. First note that, throughsimple algebraic manipulations, the KL-divergence between
the k-subset distribution and the constrained uniform distribution can be rewritten asPthe sum of negative entropy,−H(z), where z ∼pθ (z | i zi = k) and log the numberof k-subsets, log nk (see appendix for details), reducing thehardness of computing the KL-divergence, to computing theentropy of a k-subset distribution, for which Algorithm 3gives a tractable algorithm. Intuitively, the uncertainty in the
distribution over a sequence of length n, k of which are true,
decomposes as the uncertainty over Zn , and the average ofthe uncertainties over the remainder of the sequence. We
defer the algorithm and its correctness to the appendix.

We tease apart S IMPLE’s improvements by comparing threedifferent flavors: S IMPLE - F, which only uses exact sampling, falling back to estimating the marginals using exact
samples; S IMPLE - B, which uses exact marginals on the
backward pass with approximate PAM samples on the forward pass; S IMPLE, using exact samples on the forward
pass with exact marginals on the backward pass.
Our results are shown in Figure 1b. As expected, we observethat SFE exhibits no bias, but high variance whereas SoftSub suffers from both bias and variance, due to the Gumbel
noise injection into the samples to make them differentiable.
We observe that I-MLE exhibits very high bias, as wellas very low variance. This can be attributed to the PAM
sampling, which in the case of k-subset distribution does
not sample faithfully from the distribution, but is instead
biased to sampling only the mode of the distribution. This
also means that, by approximating the marginals using PAM
samples, there is a lot less variance to our gradients. On to
our S IMPLE gradient estimator, we see that it exhibits lessbias as well as less variance compared to all the other gradient estimators. We also see that each estimated gradientis, on average, much more aligned with the exact gradient.
To understand why that is, we compare S IMPLE, S IMPLE - F,
and S IMPLE - B. As hypothesized, we observe that exact
sampling, S IMPLE - F, reduces the bias, but increases thevariance compared to I-MLE, this is since, unlike the PAM
samples, our exact sample span the entire sample space.We
also observe that, even compared to I-MLE, S IMPLE - B, reduces variances by marginalizing over all possible samples.

We plot the ELBO against the number of epochs, as seenin Figure 3. We compared against I-MLE using sum-ofgamma noise as well as Gumbel noise for PAM sampling,on the 10-subset DVAE, and against ST Gumbel Softmax
on the 1-subset DVAE. We observe a significantly lower
loss on the test set on the 10-subset DVAE, but also on the
1-subset DVAE compared to ST Gumbel Softmax.
4.3. Learning to Explain
Next, we address the problem introduced by the L2X paper (Chen et al., 2018) of learning a k-subset distribution
over words that best explain a given rating. We compare torelaxation-based baselines L2X (Chen et al., 2018) and SoftSub (Xie & Ermon, 2019) as well as to I-MLE which uses
perturb-and-MAP to both compute an approximate sample
in the forward pass and to estimate the marginals. Table 1
lists results for k ∈ {5, 10, 15} for the A ROMA aspect. The
mean-squared error (MSE) of S IMPLE is almost always
lower and its subset precision never significantly exceededby the baselines. Table 1 shows results on the remaining
aspects Appearance, Palate, and Taste for k = 10.

4.2. Discrete Variational Auto-Encoder
Next, we test our S IMPLE gradient estimator on the taskof learning a generative model of MNIST digits using a
discrete variational auto-encoder (DVAE), where the latent

4

SIMPLE: A Gradient Estimator for k-subset Sampling

Referencesmodel interpretation. In Dy, J. and Krause, A. (eds.),
Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine
Learning Research, pp. 883–892. PMLR, 10–15 Jul 2018.

Ahmed, K., Wang, E., Chang, K.-W., and Van den Broeck,
G. Leveraging unlabeled data for entity-relation extraction through probabilistic constraint satisfaction, mar
2021.

Correia, G. M., Niculae, V., and Martins, A. F. Adaptively sparse transformers. In Proceedings of the 2019
Conference on Empirical Methods in Natural Language
Processing and the 9th International Joint Conference
on Natural Language Processing (EMNLP-IJCNLP), pp.
2174–2184, 2019.

Ahmed, K., Li, T., Ton, T., Guo, Q., Chang, K.-W., Kordjamshidi, P., Srikumar, V., Van den Broeck, G., andSingh, S. Pylon: A pytorch framework for learningwith constraints. In Proceedings of the 36th AAAI
Conference on Artificial Intelligence (Demo Track),
feb 2022a. URL http://starai.cs.ucla.edu/
papers/AhmedAAAI22.pdf.

Correia, G. M., Niculae, V., Aziz, W., and Martins, A. F.
Efficient marginalization of discrete and structured latentvariables via sparsity. Advances in Neural Information
Processing Systems, 2020.

Ahmed, K., Teso, S., Chang, K.-W., Van den Broeck, G.,
and Vergari, A. Semantic probabilistic layers for neurosymbolic learning. In Advances in Neural Information
Processing Systems 35 (NeurIPS), dec 2022b.

Corro, C. and Titov, I. Differentiable perturb-and-parse:
Semi-supervised parsing with a structured variational
auto-encoder. In International Conference on Learning
Representations, 2019.

Ahmed, K., Wang, E., Chang, K.-W., and den Broeck, G. V.
Neuro-symbolic entropy regularization. In The 38th Conference on Uncertainty in Artificial Intelligence, 2022c.

Cuturi, M., Teboul, O., and Vert, J.-P. Differentiable ranking
and sorting using optimal transport. Advances in neural
information processing systems, 32, 2019.

Ahmed, K., Chang, K.-W., and Van den Broeck, G. Semantic strengthening of neuro-symbolic learning. In Proceedings of the 26th International Conference on Artificial
Intelligence and Statistics (AISTATS), apr 2023.

de Silva, B., Champion, K., Quade, M., Loiseau, J.-C., Kutz,
J., and Brunton, S. Pysindy: A python package for thesparse identification of nonlinear dynamical systems from
data. Journal of Open Source Software, 5(49):2104, 2020.
doi: 10.21105/joss.02104. URL https://doi.org/
10.21105/joss.02104.

Belle, V., Passerini, A., and Van den Broeck, G. Probabilistic inference in hybrid domains by weighted model
integration. In Proceedings of IJCAI, pp. 2770–2776,
2015a.

Domke, J. Implicit differentiation by perturbation. In Advances in Neural Information Processing Systems 23, pp.
523–531. 2010.

Belle, V., Van den Broeck, G., and Passerini, A. Hashingbased approximate probabilistic inference in hybrid domains. In Proceedings of the 31st Conference on Uncertainty in Artificial Intelligence (UAI), 2015b.

Grathwohl, W., Choi, D., Wu, Y., Roeder, G., and Duvenaud,
D. Backpropagation through the void: Optimizing control
variates for black-box gradient estimation. ICLR, 2018.

Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.,
and Bach, F. R. Learning with differentiable pertubed
optimizers. In NeurIPS, 2020.

Grover, A., Wang, E., Zweig, A., and Ermon, S. Stochastic
optimization of sorting networks via continuous relaxations. In International Conference on Learning Representations, 2018.

Blondel, M. Structured prediction with projection oracles.
Advances in neural information processing systems, 32,
2019.

Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon,
France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.

Blondel, M., Martins, A. F., and Niculae, V. Learning withfenchel-young losses. J. Mach. Learn. Res., 21(35):1–69,
2020a.
Blondel, M., Teboul, O., Berthet, Q., and Djolonga, J.
Fast differentiable sorting and ranking. In International
Conference on Machine Learning, pp. 950–959. PMLR,
2020b.

Kaptanoglu, A. A., de Silva, B. M., Fasel, U., Kaheman,K., Goldschmidt, A. J., Callaham, J., Delahunt, C. B.,
Nicolaou, Z. G., Champion, K., Loiseau, J.-C., Kutz,
J. N., and Brunton, S. L. Pysindy: A comprehensive
python package for robust sparse system identification.
Journal of Open Source Software, 7(69):3994, 2022. doi:

Chen, J., Song, L., Wainwright, M., and Jordan, M. Learning to explain: An information-theoretic perspective on
5

SIMPLE: A Gradient Estimator for k-subset Sampling

10.21105/joss.03994. URL https://doi.org/10.
21105/joss.03994.

Minervini, P., Franceschi, L., and Niepert, M. Adaptive
perturbation-based gradient estimation for discrete latentvariable models. In AAAI, 2023.

Kim, C., Sabharwal, A., and Ermon, S. Exact samplingwith integer linear programs and random perturbations.
In Thirtieth AAAI Conference on Artificial Intelligence,
2016.

Niculae, V. and Martins, A. F. T. Lp-sparsemap: Differentiable relaxed optimization for sparse structured prediction. In ICML, 2020.

Kolb, S., Morettin, P., Zuidberg Dos Martires, P., Sommavilla, F., Passerini, A., Sebastiani, R., and De Raedt, L.
The pywmi framework and toolbox for probabilistic inference using weighted model integration. In Proceedingsof the Twenty-Eighth International Joint Conference on
Artificial Intelligence, IJCAI, pp. 6530–6532, 7 2019. doi:
10.24963/ijcai.2019/946.

Niculae, V., Martins, A. F. T., Blondel, M., and Cardie, C.
Sparsemap: Differentiable sparse structured inference. In
ICML, 2018.

Kool, W., Van Hoof, H., and Welling, M. Stochastic beams
and where to find them: The gumbel-top-k trick forsampling sequences without replacement. In International Conference on Machine Learning, pp. 3499–3508.
PMLR, 2019.

Paulus, M. B., Choi, D., Tarlow, D., Krause, A., and Maddison, C. J. Gradient estimation with stochastic softmax
tricks. arXiv preprint arXiv:2006.08063, 2020.

Niepert, M., Minervini, P., and Franceschi, L. Implicit mle:
Backpropagating through discrete exponential family distributions. Advances in Neural Information Processing
Systems, 34, 2021.

Peters, B., Niculae, V., and Martins, A. F. Sparse sequenceto-sequence models. In Proceedings of the 57th Annual
Meeting of the Association for Computational Linguistics,
pp. 1504–1519, 2019.

Lei, T., Barzilay, R., and Jaakkola, T. Rationalizing neuralpredictions. In EMNLP, 2016.

Petersen, F., Borgelt, C., Kuehne, H., and Deussen, O.
Learning with algorithmic supervision via continuous
relaxations. Advances in Neural Information Processing
Systems, 34, 2021.

Maddison, C. J., Tarlow, D., and Minka, T. A∗ sampling. In
Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N.,
and Weinberger, K. (eds.), Advances in Neural Information Processing Systems, volume 27. Curran Associates,
Inc., 2014.

Plötz, T. and Roth, S. Neural nearest neighbors networks.
In Bengio, S., Wallach, H. M., Larochelle, H., Grauman,
K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advancesin Neural Information Processing Systems 31: Annual
Conference on Neural Information Processing Systems
2018, NeurIPS 2018, December 3-8, 2018, Montréal,
Canada, pp. 1095–1106, 2018.

Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete
distribution: A continuous relaxation of discrete randomvariables. In 5th International Conference on Learning
Representations, ICLR 2017, Toulon, France, April 24-26,
2017, Conference Track Proceedings. OpenReview.net,
2017.

Rolinek, M., Musil, V., Paulus, A., Vlastelica, M., Michaelis,C., and Martius, G. Optimizing rank-based metrics withblackbox differentiation. In Proceedings of the IEEE/CVF
Conference on Computer Vision and Pattern Recognition,
pp. 7620–7630, 2020.

Manhaeve, R., Dumancic, S., Kimmig, A., Demeester, T.,
and De Raedt, L. Deepproblog: Neural probabilistic logic
programming. In NeurIPS, 2018.
Martins, A. and Astudillo, R. From softmax to sparsemax: Asparse model of attention and multi-label classification. In
International conference on machine learning, pp. 1614–
1623. PMLR, 2016.

Rolínek, M., Swoboda, P., Zietlow, D., Paulus, A., Musil,V., and Martius, G. Deep graph matching via blackbox
differentiation of combinatorial solvers. In ECCV, 2020.

McAuley, J., Leskovec, J., and Jurafsky, D. Learning attitudes and attributes from multi-aspect reviews. 2012
IEEE 12th International Conference on Data Mining, pp.
1020–1025, 2012.

Roth, D. On the hardness of approximate reasoning.
Artif. Intell., 82(1-2):273–302, 1996. doi: 10.1016/
0004-3702(94)00092-1. URL https://doi.org/
10.1016/0004-3702(94)00092-1.

Mensch, A. and Blondel, M. Differentiable dynamic programming for structured prediction and attention. In
International Conference on Machine Learning, pp. 3462–
3471. PMLR, 2018.

Shirobokov, S., Belavin, V., Kagan, M., Ustyuzhanin, A.,
and Baydin, A. G. Black-box optimization with local
generative surrogates. Advances in Neural Information
Processing Systems, 33:14650–14662, 2020.
6

SIMPLE: A Gradient Estimator for k-subset Samplingwhere αi := log pθi (zi ) denotes the log marginals.

Tucker, G., Mnih, A., Maddison, C. J., Lawson, D., andSohl-Dickstein, J. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 2017.

P
Proof. We first rewrite the marginal pθ ( i zi = k) intoa summation as the probability for all possible events bydefinition as follows.
Ppθ ( j zj = k)YXY
(1 − exp(ᾱj ))exp(αj ) ·=Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach.
Learn., 8:229–256, 1992. doi: 10.1007/BF00992696.
Xie, S. M. and Ermon, S. Reparameterizable subset sampling via continuous relaxations. In Proceedings of theTwenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI-19, pp. 3919–3925. International
Joint Conferences on Artificial Intelligence Organization,
7 2019. doi: 10.24963/ijcai.2019/544.

z:P

j zj =kj:zj =0

j:zj =1

Here we assume that the probability of zj = 0 is a constant∂term w.r.t. parameter αj , i.e., ∂α(1 − exp(ᾱj )) = 0.*
P jFurther, the derivative of pθ ( j zj = k) w.r.t. θi is asfollows,

Xie, Y., Dai, H., Chen, M., Dai, B., Zhao, T., Zha, H.,
Wei, W., and Pfister, T. Differentiable top-k with optimaltransport. Advances in Neural Information Processing
Systems, 33:20520–20531, 2020.

P∂pθ ( j zj = k)
∂αi
X
∂
=
∂αi P
z:

Xu, J., Zhang, Z., Friedman, T., Liang, Y., and Van den
Broeck, G. A semantic loss function for deep learningwith symbolic knowledge. In International conferenceon machine learning, pp. 5502–5511. PMLR, 2018.

=

j zj =k∧zi =1∂
exp(αi )∂αi= exp(θi )z:= pθ (

Zeng, Z. and Broeck, G. V. d. Collapsed inference forbayesian deep learning. arXiv preprint arXiv:2306.09686,
2023.

P

P

z:

P

Y

exp(αj )j zj =k∧zi =1

X

j zj =k∧zi =1(1 − exp(ᾱj ))j:zj =0

j:zj =1

X

YY

exp(αj )j:zj =1,j̸=iexp(αj )(1 − exp(ᾱj ))j:zj =0

j:zj =1,j̸=i

Y

Y

Y

(1 − exp(ᾱj ))j:zj =0j zj = k ∧ zi = 1),

where the first equality holds since terms corresponding tozi ̸= 1 has their derivative to be zero w.r.t. θi . It further
holds thatP∂P∂
∂αi pθ (j zj = k)P
log pθ ( j zj = k) =
z
pθ (
∂αi
j j = k)Ppθ ( j zj = k ∧ zi = 1)
P=pθ ( j zj = k)P
= p θ zi | j z j = kZeng, Z. and Van den Broeck, G. Efficient search-based
weighted model integration. Proceedings of UAI, 2019.
Zeng, Z., Morettin, P., Yan, F., Vergari, A., and Broeck,
G. V. d. Scaling up hybrid probabilistic inference withlogical and arithmetic constraints via message passing. In
Proceedings of the International Conference of Machine
Learning (ICML), 2020a.
Zeng, Z., Morettin, P., Yan, F., Vergari, A., and Van den
Broeck, G. Probabilistic inference with algebraic constraints: Theoretical limits and practical approximations.
Advances in Neural Information Processing Systems, 33:
11564–11575, 2020b.

which finishes our proof.
P
Proposition 1. The probability pθ ( i zi = k) of sampling
exactly k items from the unconstrained distribution pθ (z)over n items can be computed exactly in time O(nk).

Zeng, Z., Morettin, P., Yan, F., Vergari, A., and Van den
Broeck, G. Is parameter learning via weighted model integration tractable? In Proceedings of the UAI Workshop
on Tractable Probabilistic Modeling (TPM), jul 2021.

Proof. Our proof is constructive. As a base case, consider
the probability of sampling k = −1 out of n = 0 items.
We can see that the probability of such an event is 0. As
a second base case, consider the probability of samplingk = 0 out of n = 0 items. We can see that the probably ofsuch an eventthat we are given the probPis 1. Now assumen−1′
ability pθ
zi = k , for k ′ = 0, . . . , k, and we areiPninterested in computing pθ ( i zi = k). By the partition

A. Proofs
P
Theorem 1. Let pθ ( j zj = k) be the probability ofexactly-k of the unconstrained distribution parameterized
by logits θ. For every Zi , its conditional marginal is


P
P
∂
p θ z i | j zj = k =log pθ ( j zj = k),∂αi* In practice, this can be easily implemented. For example, in framework Tensorflow, it can be done by setting
tf.stop_gradients.

7

SIMPLE: A Gradient Estimator for k-subset Sampling

Proof. By the definition of KL divergence, it holds thatXzi = k) || U(z))
D(pθ (z |theorem, we can see thatPnpθ ( i zi = k) =

P
n−1
pθ
z
=
k
· pθn (zn = 0)+
i
i

P
n−1
pθ
zi = k − 1 · pθn (zn = 1)ii=pθ (z |X−X

pθ (z |

zi = k) · log

i

i

z∈C

pθ (z |pθ (z |

Pi zi = k)

U (z)Xzi = k))zi = k) log pθ (z |iXzi = k) log U (z)

i

z∈C

 
n
= −H(z) + log.
k
The last equality holds since U (z) ≡ 1/ nk .

B. Correctness of Sampling
Proposition 2. Let Sample be defined as in Algorithm 2.
Given n random variables Z1 , · P· · , Zn , a subset size k, anda k-subset distribution pθ (z | i zi = k) parameterized
by logits θ, Algorithm 2 draws exact samples from pθ (z |Pi zi = k) in time O(n).

Algorithm 3 Entropy(θ, n, k)
Input: The logits θ of the distribution, the number of variables n, and the subset size k
Output: H(z) = −Ez∼pθ (z|Pi zi =k) [log p(z)]h = zeros(n, k)
for i = k to n do
for j = 0 to
Pki do// p(zi | m=1 zm = j)p = a[i − 1, j − 1] ∗ pθi (zi = 1)/a[i, j]// cf. proof of Prop. 4 in Appendix
h[i, j] = Hb (p) + p ∗ h[i − 1, j]+(1 − p) ∗ h[i − 1, j + 1]
return h

Proof. Assume that variables Zn , · · · , Zi+1 are sampled
Pn and have their values to be zn , · · · , zi+1 with
m=i+1 zm = k − j. By Algorithm 2 we have that theprobability of sampling Zi is
pSample (zi = 1 | zn , · · · , zi+1 )PnPpθ ( m=i zm = k − j + 1 | m zm = k) pθi (zi = 1)PnP=pθ ( m=i+1 zm = k − j | m zm = k)PnPpθ ( m=i+1 zm = k − j | zi = 1, m zm = k) pθi (zi = 1)PnP=pθ ( m=i+1 zm = k − j | m zm = k)PnP= pθ (zi = 1 | m=i+1 zm = k − j, m zm = k)

Proposition 4. Let Entropy be defined as in Algorithm 3.
Given variables,
Z1 , · · · , Zn , and a k-subset distributionPpθ (z | i zi = k) parameterizedby θ, Algorithm 3 comP
putes entropy of pθ (z | i zi = k).

The last line holds by Bayes’ theorem. It follows that samples drawn
P from Algorithm 2 are distributed according topθ (z | i zi = k).

Proof. In a slight abuse of notation, let zn denote zn = 1,
denote by
and let z̄n denote zn = 0. Furthermore,
Pn
Pwe
n−1
k−1
k
σnk , σn−1
and σn−1
the events i=0 = k, i=0 = k and
Pn−1
i=0 = k − 1, respectively.

C. Exact ELBO Computation
We will now present a formal proof on how to compute
the KL-divergence between the k-subset distribution and a
uniform distribution tractably and exactly.
P
Proposition 3. Let pθ (z | i zi = k) be a k-subset distribution parameterized by θ and U (z) be aPuniform distribution on the constrained space C = {z | i zi =Pk}. Thenthe KL-divergence between distribution pθ (z | i zi = k)and U (z) can be computed by XnD(pθ (z |zi = k) || U(z)) = −H(z) + log
kiwhere H denotes the entropy of pθ (z |

X

=(

Pn−1
Pn−1
as events i zi = k and i zi = k − 1 are disjoint
and, for any k, partition the sample space. Intuitively, for any
k and n, we can sample k out of n items by choosing k ofn−1 items, and not the n-th item, or choosing k −1 of n−1
items, and the n-th item.
P The above process gives Algorithm
1, which returns pθ ( i zi = k) in time O(nk).

P

z∈C

X

X

The entropy of the k-subset distribution is
H(Z) = −Ez∼pθ (z|σnk ) [log p(z)]X=−pθ (z | σnk ) log pθ (z | σnk )kz:σnWe start by simplifying the expression for pθ (z | σnk ),where, by the chain rule , the above isXkpθ (z̄n | σnk ) · pθ (σn−1| σnk , z̄n )
k
z:σn

k−1
+ pθ (zn | σnk ) · pθ (σn−1| σnk , zn )i zi = k).

8SIMPLE: A Gradient Estimator for k-subset Sampling

Algorithm 4 PrExactlyk(θ, l, u, k)Input: The logits θ of the distribution, range of variableindices [l, u], and the subset size k
Output: The exact P
marginal probability of variables sumu
ming up to k, P ( i=l Xi = k)if l > u then return 0
if l = u then return pθ (Xl = k)
for m = 0 to k do
pm = PrExactlyk(θ, l, ⌊u/2⌋, m)∗
PrExactlyk(θ, ⌊u/2⌋ + 1, u, k − m)Pkreturn m=0 pm

Plugging the above in the expression for the entropy anddistributing the sum over the product, we get

H(Z) = −Xkpθ (z̄n | σnk ) · pθ (σn−1| σnk , z̄n )

k
z:σn



k−1
k
| σnk , zn )· log pθ (z̄n | σnk ) · pθ (σn−1| σnk , z̄n ) + pθ (zn | σnk ) · pθ (σn−1k−1+ pθ (zn | σnk ) · pθ (σn−1| σnk , zn )


k−1
k
k
| σnk , zn ) ,· log pθ (z̄n | σn ) · pθ (σn−1 | σnk , z̄n ) + pθ (zn | σnk ) · pθ (σn−1where, since the two events z̄n and zn are mutually exclusive, we can simplify the above to

H(Z) = −X

Algorithm 5 Sample(θ, l, u, k)Input: The logits θ of the distribution, range of variableindices [l, u], and the subset size k
POutput: A sample z = (z1 , . . . , zn ) from pθ (z | i zi =k)define p(x = m) = pm , m = 0, · · · , k// with pm as defined in Algorithm 4
sample m∗ from pzl:⌊u/2⌋ = Sample(θ, l, ⌊u/2⌋, m∗ )z⌊u/2⌋+1:u = Sample(θ, ⌊u/2⌋ + 1, u, k − m∗ )return Concat(zl:⌊u/2⌋ , z⌊u/2⌋+1:u )kpθ (z̄n | σnk ) · pθ (σn−1| σnk , z̄n )

k
z:σn

k
· log pθ (z̄n | σnk ) · pθ (σn−1| σnk , z̄n )





k−1
| σnk , zn )+ pθ (zn | σnk ) · pθ (σn−1k−1· log pθ (zn | σnk ) · pθ (σn−1| σnk , zn ) .

Expanding the logarithms, rearranging terms, and using thatconditional probabilities sum to 1 we get

H(Z) = pθ (z̄n | σnk ) log pθ (z̄n | σnk )E. Experimental Details+ pθ (zn | σnk ) log pθ (zn | σnk )kk+ pθ (z̄n | σnk ) · pθ (σn−1| σnk , z̄n ) log pθ (σn−1| σnk , z̄n )E.1. Synthetic Experiments

k−1
k−1+ pθ (zn | σnk ) · pθ (σn−1| σnk , zn ) log pθ (σn−1| σnk , zn )k= −Ezn ∼pθ (zn |σnk ) − log pθ (zn | σn ) +
Ezn ∼pθ (zn |σnk ) H(Z:n−1 |σnk , zn )= Hb (Zn | σnk ) + Ezn ∼pθ (zn |σnk ) H(Z:n−1 |σnk , zn ) .

In this experiment we analyzed the behavior of various
discrete gradient estimators for the k-subset distribution. We
were interested in three different metrics: the bias of the thegradients estimators, the variance of the gradient estimators,
as well as the average deviation of each estimated gradientfrom the exact gradient. We used cosine distance, definedas 1− cosine similarity as our measure of distance, as we
typically care about the direction, not the magnitude of thegradient; the latter can be recovered using an appropriate
learning rate. Following Niepert et al. (2021), we chosea tractable 5-subset distribution,
where n = 10, and were
therefore limited to 10
5 = 252 possible subsets. We set
the loss to L(θ) = Ez∼pθ (z|P zi =k) [∥z − b∥2 ], where bThat is, simply stated, the entropy of the k-subset distribution decomposes as the entropy of the constrained distribution over Zn , and average entropy of the distribution on theremaining variables.
As the base case, the entropy of the k-subset distribution
when k = n is 0; there is only one way in which to pickto choose n of n variables, and the k-subset distribution istherefore deterministic.

iis the groundtruth logits sampled from N (0, I). We used asample size of 10000 to estimate each of our metrics.

D. Optimized Algorithms

E.2. Discrete Variational Auto-Encoder

Algorithm 4 is the optimized version of Algorithm 1, bothof which compute the marginal probability of the exactly-k
constraint. Algorithm 5 is the optimized version of Algorithm 2, both of which sample faithfully from the k-subset
distribution.

We tested our S IMPLE gradient estimator in the discrete
k-subset Variational Auto-Encoder (VAE) setting, where
the latent variables model a probability distribution overk-subsets, and has a dimensionality of 20. The experimental
setup is similar to those used in prior work on the Gumbel
9

SIMPLE: A Gradient Estimator for k-subset Sampling

Algorithm 6 The proposed algorithm for the k-subset distributionfunction BACKWARD PASS(∇z ℓ(fu (z, x), y))load θ from the forwardP
pass// Derivatives of pθ (z | i zi = k)µ = ∇θ log a[n, k] // by auto-diff// Return the directional derivative of the// marginals along the downstream gradients
return JVP(µ, ∇z ℓ(fu (z, x))functionPF ORWARD PASS(θ)
i// pθ ( m=1 zm = j) for all i, j
a = PrExactlyk(θ, n,Pk)// Sample from pθ (z | i zi = k)z = Sample(θ, n, k)save a for the backward passreturn z

softmax tricks (Jang et al., 2017) and IMLE (Niepert et al.,2021) The encoding and decoding functions of the VAE
consist of three dense layers (encoding: 512-256-20x20;
decoding: 256-512-784). As is commonplace in discrete
VAEs, the loss is the sum of the reconstruction loss (binary
cross-entropy loss on output pixels) and KL divergenceof the k−subset distribution and the uniform distribution,
known as the evidence lower bound, or the ELBO. The
task being to learn a sparse generative model of MNIST.
As in prior work, we use a batch size of 100 and train for
100 epochs, plotting the test loss after each epoch. We use
the standard Adam settings in Tensorflow 2.x, and do not
employ any learning rate scheduling. The encoder network
consists of an input layer with dimension 784 (we flatten
the images), a dense layer with dimension 512 and ReLu
activation, a dense layer with dimension 256 and ReLu
activation, and a dense layer with dimension 400(20 × 20)
which outputs θ and no non-linearity S IMPLE takes θ as
input and outputs a discrete latent code of size 20 × 20. The
decoder network, which takes this discrete latent code as
input, consists of a dense layer with dimension 256 andReLu activation, a dense layer with dimension 512 andReLu activation, and finally a dense layer with dimension
784 returning the logits for the output pixels. Sigmoids areapplied to these logits and the binary cross-entropy loss iscomputed. To obtain the best performing model of each of
the compared methods, we performed a grid search overthe learning rate in the range [1 × 10−3 , 5 × 10−4 ], λ in therange [1 × 10−3 , 1 × 10−2 , 1 × 10−1 , 1 × 100 , 1 × 101 , 1 ×
102 , 1 × 103 ], and for SoG I-MLE, the temparature τ in therange [1 × 10−1 , 1 × 100 , 1 × 101 , 1 × 102 ]annotations of the words that best describe the review scorewith respect to the various aspects. Following the experimental setup of recent work (Paulus et al., 2020; Niepert
et al., 2021), we address the problem introduced by the L2X
paper (Chen et al., 2018) of learning a k-subset distribution
over words that best explain a given aspect rating. Subset
precision was computed using a set of 993 annotated reviews. We use pre-trained word embeddings from (Lei et al.,2016)† We use the standard neural network architecturefrom prior work Chen et al. (2018); Paulus et al. (2020) with
4 convolutional and one dense layer. This neural network
outputs the parameters θ of the k-subset distribution over
k-hot binary latent masks with k ∈ {5, 10, 15}. We train for
20 epochs using the standard Adam settings in Tensorflow
2.x, and no learning rate schedule. We always evaluate themodel with the best validation MSE among the 20 epochs.

E.3. Learning to Explain

One might wonder if our gradient estimator reduces to
the Straight-Through (ST) Gumbel-Softmax estimator, or
relates to it in any way when k = 1. On the forward
pass, the ST Gumbel Softmax estimator makes use of
the Gumbel-Max trick (Maddison et al., 2014), whichstates that we can efficiently sample from a categorical
distribution by perturbing each of the logits with standard Gumbel noise, and taking the MAP, or more formallyz = OneHot(arg maxi∈{1,...,k} θi + gi ) ∼ pθ where thegi ’s are i.i.d Gumbel(0, 1) samples, and OneHot encodes

F. Connection to Straight-ThroughGumbel-Softmax
0.06
0.04

Bias
Variance

0.02
0.00

Exact ST Gumbel Softmax SIMPLE

0.125
0.100
0.075
0.050
0.025
0.000

Figure 3: Bias and variance of S IMPLE and Gumbel Softmax

The B EER A DVOCATE dataset (McAuley et al., 2012) consists of free-text reviews and ratings for 4 different aspectsof beer: appearance, aroma, palate, and taste. The training
set has 80k reviews for the aspect A PPEARANCE and 70k
reviews for all other aspects. The maximum review lengthis 350 tokens. We follow Niepert et al. (2021) in computing
10 different evenly sized validation/test splits of the 10k
held out set and compute mean and standard deviation over
10 models, each trained on one split. In addition to theratings for all reviews, each sentence in the test set contains†

10

http://people.csail.mit.edu/taolei/beer/.

SIMPLE: A Gradient Estimator for k-subset Sampling

the sample as a binary vector.

assuming perfect parallelization. The optimization is possible by computing the marginal probability in a divide-andconquer way: it partitions the variables into two subsets andcompute their marginals respectively such that the complexity O(n) is reduced to O(log n); the summation over the kterms also has its complexity reduced to O(log k) in a similar manner. We refer the readers to Algorithm 4 in Appendix
for the optimized algorithm. We further modify Algorithm 2
to perform divide-and-conquer such that sampling k-subsets
achieves a vectorized complexity being O(log n), shownas Algorithm 5 in the Appendix. As a comparison, SoftSub (Xie & Ermon, 2019) has a complexity O(nk) due tothe relaxed top-k operation and its vectorized complexity tobe O(k log n) stemming from the fact that softmax layers
need O(log n) rounds of communication for normalization.

Since arg max is non-differentiable, Gumbel-Softmax uses
the perturbed relaxed samples, y = Softmax(θ + gi )as a proxy for discrete samples z on the backward
pass, using differentiable Softmax in place of the nondifferentiable arg max, with the entire function returning(z−y). detach()+y where detach ensures that the gradient
flows only through the relaxed samples on the backward.
That is, just like S IMPLE, ST Gumbel-Softmax returns exact,
discrete samples. However, whereas S IMPLE backpropagates through the exact marginals, ST Gumbel Softmax
backpropagates through the perturbed marginals that result
from applying the Gumbel-max trick. As can be seen in
Figure 3, such a minor difference means that, empirically,
S IMPLE exhibits lower bias and variance compared to ST
Gumbel Softmax while being exactly as efficient.

H. Related Work
There is a large body of work on gradient estimation for categorical random variables. Maddison et al. (2017); Jang et al.
(2017) propose the Gumbel-softmax distribution (named theconcrete distribution by the former) to relax categorical random variables. For more complex distributions, such as the
k-subset distribution which we are concerned with in this paper, existing approaches either use the straight-through andscore function estimators or propose tailor-made relaxations(see for instance (Kim et al., 2016; Chen et al., 2018; Grover
et al., 2018)). We directly compare to the score function andstraight-through estimator as well as the tailored relaxationsof (Chen et al., 2018; Grover et al., 2018) and show that weare competitive and obtain a lower bias and/or variance thanthese other estimators. Tucker et al. (2017) and Grathwohl
et al. (2018) develop parameterized control variates basedon continuous relaxations for the score-function estimator.
Lastly, Paulus et al. (2020) offers a comprehensible work on
relaxed gradient estimators, deriving several extensions of
the softmax trick. All of the above works, ours included, assume the independence of the selected items, beyond there
being k of them. That is with the exception of (Paulus et al.,2020) which make use of a relaxation using pairwise embeddings, but do not make their code available. We leavethat to future work.

F.1. Sparse Linear Regression
Given a library of feature functions, the task of sparse linear
regression aims to learn from data which feature subset best
describes the nonlinear partial differential equation (PDE)that the data are sampled from. We propose to tackle thistask by learning a k-subset distribution over the featurefunctions. During learning, we first sample from the ksubset distribution to decide which feature function subset
to choose. With k chosen features, we perform linear regression to learn the coefficients of the features from data,
and then update the k-subset distribution logit parameters
by minimizing RMSE.
To test our proposed approach, we follow the experimental
setting in PySINDy (de Silva et al., 2020; Kaptanoglu et al.,2022) and use the dataset collected by PySINDy where thesamples are collected from the Kuramoto–Sivashinsky (KS)equation, a fourth-order nonlinear PDE known for its chaotic
behavior. This PDE takes the form vt = −vxx −vxxxx −vvx ,which can be seen as a linear combination of feature functions V = {vxx , vxxxx , vvx } with the coefficients all set toa value of −1. At test time, we use the MAP estimationof the learned k-subset distribution to choose the k featurefunctions. For k = 3, our proposed method achieves thesame performance as the state-of-the-art solver on this task,
PySINDy. It identifies the KS PDE from data by choosing
exactly the ground truth feature function subset V, with
RMSE 0.00622 after applying linear regression on V.

A related line of work has developed and analyzed sparse
variants of the softmax function, motivated by their potential computational and statistical advantages. Representative examples are (Blondel et al., 2020a; Peters et al.,2019; Correia et al., 2019; Martins & Astudillo, 2016).
SparseMAP (Niculae et al., 2018) has been proposed in
the context of structured prediction and latent variable models, also replacing the softmax with a sparser distribution.
LP-SparseMAP (Niculae & Martins, 2020) is an extensionthat uses a relaxation of the optimization problem rather
than a MAP solver. Sparsity can also be exploited for effi-

G. Complexity Analysis
In Proposition 1, we prove that computing the marginal probability of the exactly-k constraint can be done tractably intime O(nk). In the context of deep learning, we often careabout vectorized complexity. We demonstrate an optimized
algorithm achieving a vectorized complexity O(log k log n),
11

SIMPLE: A Gradient Estimator for k-subset Sampling

cient marginal inference in latent variable models (Correia
et al., 2020). Contrary to our work, they cannot control
the sparsity level exactly through a k-subset constraint or
guarantee a sparse output. Also, we aim at cases where
samples in the forward pass are required.
Integrating specialized discrete algorithms into neural networks is a line of research with increasing popularity. Examples are sorting algorithms (Cuturi et al., 2019; Blondel
et al., 2020b; Grover et al., 2018), ranking (Rolinek et al.,2020; Kool et al., 2019), dynamic programming (Mensch& Blondel, 2018; Corro & Titov, 2019), and solvers for
combinatorial optimization problems (Berthet et al., 2020;
Rolínek et al., 2020; Shirobokov et al., 2020; Niepert et al.,2021; Minervini et al., 2023; Zeng et al., 2021) or even probabilistic circuits over structured output spaces, either as hardconstraints (Ahmed et al., 2022b; Blondel, 2019), or soft
constraints (Xu et al., 2018; Manhaeve et al., 2018; Ahmed
et al., 2021; 2022a;c; 2023). There has also been work on
making common programming language expression such as
conditional statements, loops, and indexing differentiablethrough relaxations (Petersen et al., 2021). Xie et al. (2020)
propose optimal transport as a way to obtain differentiablesorting methods for top-k classification. In contrast, we
focus on the k-subset sampling problem and provide exact discrete sampling and marginal inference algorithms,
obtaining a gradient estimator for the k-subset distribution
with a favorable bias-variance trade-off. Besides the discrete
constraints we consider in this work, probabilistic inference
over hybrid constraints that involve both discrete and continuous variables remains an active topic (Belle et al., 2015b;a;Zeng & Van den Broeck, 2019; Kolb et al., 2019; Zenget al., 2020a) whose tractability is studied and leveraged forbuilding inference algorithms with applications in Bayesian
deep learning (Zeng et al., 2021; 2020b; Zeng & Broeck,2023).

I. Conclusion
We introduced a gradient estimator for the k-subset distribution which replaces relaxed and approximate samplingon the forward pass with exact sampling. It sidesteps thenon-differentiable nature of discrete sampling by estimating
the gradients as a function of our distribution’s marginals,
for which we prove a simple characterization, showing that
we can compute them exactly and efficiently. We demonstrated improved empirical results on a number of tasks:
L2X, DVAEs, and sparse regression.

12