SIMPLE: A Gradient Estimator for k-Subset Sampling
Authors: Kareem Ahmed
Zhe Zeng
Mathias Niepert
Guy Van den Broeck
Venue: ICLR 2023 poster (2023)

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, but introduces both bias and 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 exhibits lower bias and variance compared to state-of-the-art estimators. Empirical results show improved performance on learning to explain and sparse models benchmarks. We provide an algorithm for computing the exact ELBO for the $k$-subset distribution, obtaining significantly lower loss compared to state-of-the-art discrete sparse VAEs. All of our algorithms are exact and efficient.

Published conference paper at ICLR 2023
SIMPLE: A GRADIENT ESTIMATOR FOR K- SUBSET SAMPLING
Kareem Ahmed" Zhe Zeng" " Mathias Niepert² Guy Van den Broeck'
'Dept. of Computer Science, UCLA Deptt of Computer Science, Stuttgart University (ahmedk, zhezeng, guyvdb ) êcs. pela..du.
mathias. wwwpeettsiitteeeeee wtwpchhhhniiutuuuttteeeu wtl-sssuutggrrr.c.
ABSTRACT
k-subset sampling is ubiquitous in machine learning, enabling regularization and interpretability through sparsity. The challenge lies in rendering k-subset sampling amenable end-to-end learning. This has typically involved relaxing the reparam- eterized 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-ard 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 algorithm for computing the exact ELBO for the k-subset distribution, obtaining significantly lower loss compared to SOTA.
INTRODUCTION
k-subset sampling. sampling a subset of size k of n variables, is omnipresent in machine learning. It lies at the core of many fundamental problems that rely upon learning sparse features representations of input data, including stochastic high-dimensional data visualization (van der Maaten 2009).
parametric k-nearest neighbors (Grover et al. 2018). learning to explain (Chen et al.) 2018). discrete variational auto-encoders (Rolfe 2017). and sparse regression, to name a few. All such tasks involve optimizing an expectation of objective function with respect to a latent discrete distribution parameterized by a neural network, which are often assumed intractable. Score-functin estimators offer a cloyingly simple solution: rewrite the gradient of the expectation an expectation of the gradient, which can subsequently be estimated using finite number of samples offering an unbiased estimate of the gradient. Simple as is, score-function estimators suffer from very high variance which can interfere with 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 to discrete distributions thereby prompting continuous relaxations, e.g. Gumbel- softmax (Jang et al.] 2017: Maddison et al. (017). that allow for reparameterized gradients w.r.t the parameters of a categorical distribution. Reparameterizable subset sampling (Xie & Ermon 2019) generalizes the Gunbel-softmax trick to k-subsets which while rendering k-subset sampling amenable to backpropagation the cost of introducing bias in the learning by using relaxed samples.
this paper, we set out with the goal of avoiding all such relaxations. Instead, we back to discrete sampling on the forward pass. On the back pass, we reparameterize the gradient of the loss function with respect to the samples as function of the exact marginals of the A-subset distribution. Computing the exact conditional marginals is, in general, intractable (Roth) 1996). We give an efficient algorithm for computing the k-subset probability, and show 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 SIMPLE, reminiscent of the straight-through (ST) Gumbel estimator when k = 1, with the gradients taken with respect to the unperturbed marginals. We empirically demonstrate that SIMPLE exhibits lower bias and variance compared to other known gradient estimators, including the ST Gumbel estimator in the case k 1.
"These authors contributed equally to this work.

---

Published conference paper at ICLR 2023
Blas Variance
-0.15
1.25
0.4
-0.10
0.75
0.05
0.25
Exart
SottSub
SIMPLE-FSMPPEE SIOPLE
Exortt SontSuu
IMRLE
SHIPLET SIPPEC- SMALE
Figure 1: A comparison of the bias and variance of the gradient estimators (left) and the average 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 cuclidean distance 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 Section 5.1
We include an experiment on the task of learning to explain (L2X) using the BEERADVOCATE dataset (McAuley et al. 2012), where the goal is select the subset of words that best explains the model's classification of a user's review. We also include an experiment on the task stochastic sparse linear regression, where the goal is to learn the best sparse model, and show that we are able to recover the Auramoto-SSvashinkk equation. Finally, we develop an efficient computation for the calculation the exact variational evidence lower bound (ELBO) for the k-subset distribution, which when used in conjunction with SIMPLE leads to state-of-the-art discrete sparse VAE learning.
Contributions. In summary. we propose replacing relaxed sampling on the forward pass with discrete sampling. On the backward pass, marginals as a proxy for the true gradient, giving an algorithm for computing them efficiently. We empirically demonstrate that discrete samples on the forward pass, coupled with exact conditional marginals the backward pass leads to gradient estimator, SIMPLE, with lower bias and variance compared to other known gradient estimators. We also provide an efficient computation the exact ELBO for the susbbtet distribution, leading to state-of ofte--o-tteer the-art discrete sparse VAE learning.
the gradient with respect the exact conditional
2 PROBLEM STATEMENT AND MOTIVATION
We consider models described by the equations
0 = (x)),,
~ p>pp Σ, =b),
-= fu(z,z),
(1)
where X and ÿ denote feature inputs and target outputs, respectively, ha X - O and fu : x X' - y are smooth, parameterized maps and 0 are logits inducing a distribution over the latent binary vector z. The induced distribution pe(z) is defined as
(z) II po. with Po, (2, 1) sigmoid(0.) and Po, (2: = = sigmoid(0:). (2) 1=1
The goal of our stochastic latent layer is not to simply sample from p=(x), which would yield samples with Hamming weight between 0 and (i.e., with an arbitrary number of ones). Instead, we are interested in sampling from the distribution restricted to samples with Hamming weight for any given kt That is, we are interested sampling from the conditional distribution Po(= Σ, 24 k).
Conditioning the distribution pa (z) on this k-subset constraint introduces intricate dependencies between each of the 's. The probability of sampling any given k-subset vector therefore, becomes
p=(z Σ 2, = k)) = po(=) /Ppp (E: = = k) IΣ, =; = k:]
where [.] denotes the indicator function. In other words, the probability of sampling each k-subset is re-normalized by po (E, = k) - the prohability of sampling exactly A: items from the unconstruined distribution induced by encoder hy. The quantity Po(E. = k:) Σ, Po (z) 1Σ; z = &] appears to be intractable. We show that not to be the case, providing tractable algorithm for computing it.
Given set of samples D, architecture (T) through minimizing the training error which the expected loss:
concerned with learning the parameters
(v, u) of the
L(x, w) = (=>-no=]EE, k) (C/// x), with = (,.(),,,
(3)

---

Published conference paper at ICLR 2023
Ju
Pe (z
Σ
k)
(ff(z,x).y)
y; = O x)
Figure The problem setting considered in our paper. On the forward pass, neural network hhe outputs 0 parameterizing a discrete distribution over subsets of size &o 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.
where é V V R a pisnt-wise loss function. This formulation, illus- trated Figure 2 is general and sub- sumes many settings. Different choices of mappings lhe and fu. and sample- wise loss define various tasks. [ presents some example settings used our experimental evaluation.
TASK
MAP le
MAP Ju
Loss
Discrete VAR (Sec.
5.2
Fazeder Embedding Identity
Decnder Regression
ELBO
Learn To Explain (Sec.
5.3
RMSE
Sparse Regression (Sec.
3.4
Linear Regression
RMSE
Table Architectures of the three experiment settings.
Learning then requires computing the gradient of L w.r.t. (v, u). The gradient of L w.r.t. is
V (xr, (x. = 1e E, =, +, 1ul fu
a²) T VV v, y)l,
(4)
where ÿ = fu(z, x)) the decoding of latent sample z. Furthermore, the gradient w.r.t. v
V.l(a, y:w) v. 0.4.(a)) " Vo Lxx y:
(5)
where VoL(x,y::)) == VoB=~po (0 1)). the loss' gradient w.r.t the encoder.
One challenge lies computing expectation ( and 4 which has no known closed-form
solution. This necessitates Monte-Carlo estimate through sampling from po(z Σ = k).
A second, and perhaps more substantial hurdle lies in computing VoL(x, y; () in (5) due to the
nonddifferentiable nature of discrete sampling. One could rewrite VoL(x, y w) as
VaL(x y w) = B (2 s.= (f(m)) x), y) Va logps (z Σ, 24 =]]
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) reparameterive the samples as a deterministic transformation of the parameters, and some independent standard Gumbel noise, and relaxing the deterministic transformation, the top-k function in this case, allow for backpropagation.
3 SIMPLE: SUBSET IMPLICIT LIKELIHOOD ESTIMATION
Our goal is to build gradient estimator for (a, y: w) We start by envisioning a hypothetical sampling-Foe free architecture, where the downstream neural network fu is function of the marginals, := (0) = {po(z) Σ b = k) instead of a discrete sample z, resulting in a loss Ly S.L.
Ve (z, w) = Oap(() () TV mm (p. x).y).
(6)
When the marginals u(0) be efficiently computed and differentiated, such hypothetical pipeline can be trained end-to-end, Furthermore, Domkee (2010) observed that, for an arbitrary loss function bm defined on the marginals, the Jacobian of the marginals w.r.t. the logits is symmetric, i.e.
VoLn ( y w) (0) T Em x),, y) (sr(O)V yu (f,
Consequently, computing the gradient of the loss w.r.t. the logits, VoLm (x, y: w). reduces to computing the directional derivative, or the Jacobian-vector product, of the marginals w.r.t. the logits in the direction of the gradient of the loss. This offers alluring opportunity: the conditional

---

Published conference paper at ICLR 2023
Algorithm 1 PrExactlyk(O n,)))
Algorithm 2 Sample(0, (amppee(,nnkk) n),
Input: The logits 9 of the distribution, the num-
Input: The logits @ of the distribution, the num-
ber of variables n. and the subset size k
ber of variables nr, and the subset size &
Output: pa(2
Output: (,),..., z.) pa(= > 2, k)
I j[ Po( Po(Emm -m 5) for i;
sample = [ =
initialize to be 0 everywhere
for i n to
a[0, 0] = 1 // pa m Am = 0) =1
cf. proof of Prop
p = a/i 1,j
for i = 1 to n do
~ Bernoulli(p
- ())
for j = to do
// Pick next state based on value of sample
cf. constructive proof Prop. T
a[t.j] = ali 1. Po, (2+ = 0)
2, 1 then j =j
sample.append(::)
1 ali 1.j 1 Po, (2+ = 1)
return a[n,A)
return sample
marginals characterize the probability of each 2; in the sample, and could be thought of as a differentiable proxy for the samples. Specifically, by reparameterizing z as function of the conditional marginal js under approximation 0n se I proposed by Niepert al. (202T), and using
the straight- estimator for the gradient the sample w.r.t. the marginals on the backward pass, we approximate our true VoL(x, w) as
Val(æ, () x V y:),,
(8)
where the directional derivative of the marginals can be taken along any downstream gradient, render- ing the whole pipeline end-to-end learnable, even in the presence of non-differentiable sampling.
Now, estimating the gradient of the loss w.r.t. the parameters be thought of as decomposing into two sub-problems: (PI) Computing the derivatives conditional marginals (ofu(0), which requires the computation of the conditional marginals, and (P2) Computing the gradient of the loss w.r.t. the samples V, L(x, y:w) using sample- wise loss, which requires drawing exact samples. These problems are complicated by conditioning on the k-subset constraint, which introduces intricate dependencies to the distribution, and is infcasible to solve naively, c.g. by enumeration. We will show simple, efficient, and exact solutions each problem, at the heart of which is the insight that need not care about the variables' order, only their sum, introducing symmetries that simplify the problem.
3.1 DERIVATIVES OF CONDITIONAL MARGINALS
In many probabilistic models, marginal inference is #P-d (Roth 1996 Zeng al. 2020 However, we observe that it not the case for the k-subset distribution. We notice that the conditional marginals correspond to the partial derivatives of the log- pro-obbbbblltyy of the k-subset constraint. To see this, note that the derivative of multi-linear function with respect to single variable retains all the terms referencing that variable, and drops all other terms; this corresponds exactly to the unnormalized conditional marginals. By taking the derivative of the log-probability, this introduces the k-subset probability in the denominator, leading the conditional marginals. Intuitively, the rate of change of the k-subset probability w.r.t. variable only depends on that variable through its length-k subsets.
Theorem 1. Let pa Ej 2j = k) be the probability of exactly-k k of the unconstrained distribution parameterized by logits O. For every variable Zu, its conditional marginal is
à 00;
pa E3 2g =
log popooo( 23 kk)
(9)
We refer the reader to the appendix for detailed proof of the above theorem. establish the tractability of the above computation of the conditional marginals, we need to show that the probability of the exactly-k constraint po EE 2; it) can be obtained tractably, which we demonstrate next. Proposition 1. The probability pe (E: 24 of sampling exactly le items from the unconstrained distribution po(2) over n items as in Equation 2 can be computed exactly in time O(nk).
Proof. Our proof is constructive. As base case, consider the probability of sampling k = 1 out of 0 items. We can that the probability of such event is 0. As second base case, consider

---

Published conference paper at ICLR 2023
the probability of sampling out of items. We can see that the probably of such an event
is 1. Now assume that we are given the probability po (E: Sn-1 = k') for k' = -..... ki and we
are interested computing po (2, 2i = ki). By the partition theorem, we can see that
E 2, - A:) Po Σ = k) Po.. (=xe = 0) + po Σ n = k ---1 1) po. (2n 1)
Σn-1 2, = and EP-1 2i k are disjoint and, for any k, partition the sample space, Intuitively, for any & and n. we can sample is out of n items by choosing & of n 1 items, and not items, and the n-th item. The above process gives rise Algorithm which returns po (2, 2; = k) time ((nk).
as events E;
the 1-th item, or choosing ls 1 of
By the construction described above, we obtain closed-form po (E. Di:=; k). which allows us compute conditional marginals (z; Σ3 2j k) by Theorem via auto-diiferenniiatin. This further allows the computation of the derivatives of conditional marginals 0op( Og (z Σj Zj k) to be amenable to auto-differentiation, solving problem (PI) exactly and efficiently.
3.2 GRADIENTS OF LOss W.R.T. SAMPLES
As alluded to in Section 3 we approximate VoL(x, yy; w) by the directional derivative of the marginals along the gradient of the loss w.r.t. discrete samples L,,(x, w:w). where is drawn from the k-subset distribution p= Σ; 2 = k). What remains is estimate the value of the loss, necessitating faithful sampling from the k-subset distribution, which might initially appear daunting.
Exact k-subset Sampling Next we show how to sample exactly from the k-subset distribution [=(= Σ, = k). We start by sampling the variables in reverse order, that is, we sample 2, through 21. The main intuition being that, having sampled (Z,, 2n-1,
2x+1) with a Hamming weight of k j. j. we sample Z; with a probability of choosing k j of n - variables and the n-th variable given that we choose k - j + I of variables. We formalize our intuition below.
Proposition 2. Let Sample be defined as Algarithm 2
Given random variables Zı,... Znn a k) parameterized by log probabilities O, > Z k) in time Onn).
subset size k. and a k-subset distribution po(z
Ee z
Algorithm draws exact samples from po z
Proof. Assume that variables Zn. with Em ++1 Zm = j. By Algorithm we have that the probability with which to sample /4,
1 sampled and have their values be zn. 26+11
Zm le - 1 Σ 2m k)) (= = 1)
Psaapee Fm. F++1)
Po m=i=+ Im = k Em Zm k)
Po >
(Em=i++1
3m k
1, Em Zm k) po, = 1)
pel (m=141 2m k Σ 2m k:)
(o mm=++1 Zm k Σ Zm = k) (by Bayes' theorem)
follows that samples drawn from Algorithm
2 are distributed according to pa(z E. 24 = k))
4 CONNECTION TO STRAIGHT- GUMBEL-
One might wonder if our gradient estimator reduces to the Straight-Through Through (ST) Gumbel-Softmax estimator, or relates to it in any way when = 1. On the forward pass, the ST Gumbel Softmax estimator makes use of the Gumbel-Max trick (Maddison et al.] 2014), which states that we can efficiently sample from categorical distribution by perturbing each of the logits with standard Gumbel noise, and taking the MAP, or more formally = = OneHot (arg ma...()....)) 0 + g.) ~ Pe where the gi is are i.i.d (0, 1) samples, and OneHot encodes the sample as a binary vector.
Since arg max is non- -on-differentiable, Gumbel-Sone Sofbbl- uses the perturbed relaxed samples, y = Softmax @ (.)) as proxy for discrete samples on the backward pass, using differentiable Softmax place of the non-differentiable 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 pass.

---

Published conference paper at ICLR 2023
Algorithm 3 The proposed algorithm for the k-subsct distribution
function FORWARDPASS())
function BACKWARDPASS = (2,((z y))
I/ po ( m=1 Fm = 1) for all 1,
load from the forward pass // derivatives of po(= E,; k))
a PrExactlyx(0,11 k)
I/ Sample from (z L; 2 = k)
14 Vo log aln, kj II by auto-diff // Return the directional derivative of the #1 marginals along the downstream gradients
Sample(0
save for the backward pass return z
return (VP (=((f((z,)))
That is, just like SIMPLE, ST Gumbel-Softtax returns exact, discrete samples. However, whereas SIMPLE back- propagates through the exact marginals. ST Gumbel Soft- max backpropagates through the perturbed marginals that result from applying the Gumbel-max -max trick. As can be seen in Figure B] such minor difference means that, empiri- cally, SIMPLE exhibits lower bias and variance compared to ST Gumbel Softmax while being exactly as efficient.
0.06
Bias
0.125
Variance
0.100
0.04
0.075
0.02
0.050
0.025
0.00
0,000
Exat
Gumbel Softmex
SIMPLE
Figure 3: Bias and variance of SIMPLE and Gumbel Softmax over 10k samples
EXPERIMENTS
We conduct experiments on four different tasks: 1) A synthetic experiment designed to test the bias and variance, well as the average deviation of SIMPLE compared to a variety of well-established estimators in the literature. 2) A discrete k-subset Variational Auto-Enncoder (DVAE) setting, where the latent space models a probability distribution over k-subsets. We will show that we can compute the evidence lower bound (ELBO) exactly, and that, coupled with exact sampling and our SIMPLE gradient estimator, attain much lower loss compared to state of the in sparse DVAEs. 3) The learning to explain (L2X) setting, where the aim to select the k-subset of words that best describe the classifier's 's prediction, where we show an improved mean-squared error, as well as precision, across the board. 4) A novel, yet simple task, sparse linear regression, where, L2X, we wish to select k-subset features that give rise to a linear regression model, avoiding overfitting the spurious features present in the data. Table the architecture with the objective functions. Our code will be made publicly available at Swwph//hoooo/CCAAAASAAA/ASIII
vein similar
5.1 SYNTHETIC EXPERIMENTS
We carried out a series of experiments with 5-subset distribution, and a latent space of dimension 10. We set the loss L(0)= E (x)E, =x=,, || b||²], where b is the groundtruth logits sampled from (0, I). Such distribution is tractable: only have (20) 252 k-subsets, which easily enumerable and therefore, the exact gradient, the golden standard, can computed closed form.
In this experiment, we are interested in three metrics: bias, variance, and the average error of each gradient estimator, where the latter is measured by averaging the deviation of each single-sample gradient estimate from the exact gradient. We used the cosine distance, defined as 1- cosine similarity as the measure of deviation in our calculation of the metrics above, as we only care about direction.
We compare against four different baselines: exact, which denotes the exact gradient; SoftSub (Xie & Ermon 2019), which uses extension of the Gumbel-Softmar trick to sample relaxed k-subsets on the forward pass; I-MLEE, which denotes IMLE gradient estimator (Niepert al. (02T). where approximate samples are obtained using perturb-and-map (PAM) on the forward pass, approximating the marginals using PAM samples the backward pass; and score function estimator, denoted SFE.
We tease apart SIMPLE'S improvements by comparing three different flavors: SIMPLE-F, which only uses exact sampling, falling back to estimating the marginals using exact samples; SIMPLE-B, B, which uses exact marginals on the backward pass with approximate PAM samples on the forward pass; and SIMPLE, coupling exact samples the forward pass with exact marginals the backward pass.

---

Published a conference paper at ICLR 2023
Discrate 10 Subset VAE
Discreto 1-Subset VAE
250
Gumbel Suftmax SINPLE
Algorithm Eatropy(B.00 (B)
225
EL90
250
Input: The logits of the distriburioe, the num-
her of variables n, and the subset size à
200
Gumbell MLE SIIPPE
ScG SIGE
Oufput: H(x) -BBangull22 seay noayylo
240
h = zcrosfin R)
175
for = to
20
40
40
80
105
20
40
50
80
150
for &
Epochs
Epachs
p 1.3
pa. 12/of.5
Figure 4: ELBO against # of epochs. (Left) Comparison of SIM- PLE against different flavors of IMLE on the 10-subset DVAE, and (Right) against ST Gumbel Softmax on the I-subset DVAE.
off. pronf of Prop D Appendix
PJ
Our results are shown in Figure I As expected, we observe that 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 well very low variance. This be attributed to the PAM sampling, which in the case of k-subset distribution does not sample faithfully from the distribution, but instead biased to sampling only the mode of the distribution. This also means that, by approximating the marginals using PAM samples, there is lot less variance to our gradients. On to our SIMPLE gradient estimator, we that it exhibits less bias as well as less variance compared to all the other gradient estimators. We also see that each estimated gradient is. on average, much more aligned with the exact gradient. To understand why that is. compare SIMPLE, SIMPLE-F, and SIMPLE-B. B. As hypothesized. we observe that exact sampling. SIMPLE-F, reduces the bias, but increases the variance compared to I-MLE, this since, unlike the PAM samples, our exact sample span the entire sample space. We also observe that, even compared to I-MLE, SIMPLE-B, reduces the variance by marginalizing over all possible samples.
5.2 DISCRETE VARIATIONAL
Next. we test our SIMPLE gradient estimator in the k-subset discrete variational auto-encoder (DVAF) setting, where the latent variables model a probability distribution over k-subsets, and has a dimensionality of 20. Similar to prior work (Jang et al., 2017) decoding functions of the VAE consist of three dense layers (encoding: 512-256-20x20 decoding: 256-512-7844 The DVAE is trained to minimize the sum of reconstruction loss and Kl-divergence of the k-subset distribution and the constrained uniform distribution, known as the ELBO, on MNIST.
[ Niepert et al.] (02T). the encoding and
In prior work, the KL-divergence was approximated using the unconditional marginals, obtained simply through Softmax layer. Instead we show that the -L-diivegencce between the k-subset distri- bution and the uniform distribution can be computed exactly. First note that, through simple algebraic manipulations, the KL-divergence between the k-subset distribution and the constrained uniform dis- tribution can be rewritten as the sum of negative entropy, - (2), where z Po (= Σ. k) and log the number of k-subsets, log (+) (see appendix for details), reducing the hardness of computing the KL-divergence, to computing the entropy of a s-subset distribution, for which Algorithm gives a tractable algorithm. Intuitively, the uncertainty in the distribution over a sequence of length n. which are true, decomposes as the uncertainty over Zn, and the average of the uncertainties over the remainder of the sequence. We refer the reader to the appendix for the proof of the below proposition.
Proposition 3. Let Entropy be defined as in Algorithm + Given variables, Zi, Znn and a k-subset distribution po(z Σ, = k)), Algorithm computes entropy of po (z >; 2 = k).
We plot the loss ELBO against the number of epochs, seen in Figure 4 We compared against I-MLE using sum-of of--o-ff--maaaa-a 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, partly attributable to the exact ELBO computation, but also the --subset DVAE compared to ST Gumbel Softmax, where the sole difference is the backward pass.
5.3 LEARNING TO EXPLAIN
The BEERADVOCATE dataset (cAAley et al. 2012) consists of free-text reviews and ratings for 4 different aspects of beer: appearance, aroma, palate, and taste. The training set has 80k reviews for the aspect APPEARANCE and 70k reviews for all other aspects. In addition to the ratings for all reviews, each sentence in the test set contains annotations of the words that best describe the review

---

Published conference paper at ICLR 2023
Appearance
Palate
Taste
Method
Test MSE
Precision
Test MSE
Precision
Test MSE
Precision
SIMPLE (Ours)
2.35 +: 0.28
66.81 :: 7.56
2.68 0.06
44.78 2.75
2.11 : 0.02
42.31 1: 0.61
L2X (t= 0.1)
10.70 4.82
30.02 + 15.82
6.70 = 0.63
50.39
13.58
6.92 + 1.61
32.23 = 4.92
SoftSub (1 = 0.5)
2,48 :: 0.10
52.86 + 7.08
2.94 + 0.08
39.17 +: 3,17
2.18 + 0.10
41.98
= 1.42
1-MLE 30)
2.51 + 0.05
65.47 4.95
2.96 0.04
40.73 + 3.15
2.38 + 0.04
41.38 1.55
Table 2: Results for three aspects with k
10: test MSE and subset precision, both x x00
Is 5
k 10
= 15
Method
Test MSE
Precision
Test MSE
Precision
Test MSE
Precision
SIMPLE (Ours)
2.27 0.05
57.30 + 3.04
2.23 0.03
47.17
2.11
3.20 +: 0.04
53.18 1.09
L2X (t= 0.1)
5.75 + 0.30
33.63 = 6.91
6.68 = 1.08
26.65 + 9.39
7.71 + 0.64
23.49 = 10.93
SoftSub (1 0.5)
2.57 1: 0.12
54.06 1 6.29
2.67 + 0.14
44.44
1: 2.27
2.52 : 0.07
37.78 :: 1.71
I-MLE (7= 30)
2.62 0.05
54.76 + 2.50
2.71 = 0.10
47.98 +: 2.26
2.91 + 0.18
39.56 - 2.07
Table 3: Results for aspect Aroma: test MSE and subset precision, both x x100, for E (5, 10, 15).
score with respect to the various aspects. We address the problem introduced by the L2X paper (Chen al. 2018) of learning a k-subset distribution over words that best explain given rating. We follow the architecture suggested in the L2X paper, consisting of four convolutional and one dense layer.
We compare to relaxation-based baselines 1,2X (Chen et al. 2018) and SoftSub (Xie & Fammn 2019) well as I-MLE which uses perturb-and -attrb-and--AAL to both compute an approximate sample in the forward pass and to estimate the marginals. Prior work has shown that the straight-throl (STE) did not work well and we omit it here. We used the standard hyperparameter settings of Chen al. (2018) and choose the temperature parameter t (0.1, 0.5, 1.0, 2.0} )) for all methods. We used trained separate models for each aspect using MSE as point-wise loss Table lists results for k E {5, 10, 15) for the AROMA aspect. The mean-squared error (MSF) of SIMPLE is almost always lower and its subset precision never significantly exceeded by those of the baselines. Table [ shows results on remaining aspects Appearance, Palate, and Taste for 10.
estimator
the standard Adam settings
5.4 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 sampled from. We propose to tackle this task by learning a k-subset distribution over the feature functions. During learning, we first sample from the k-subset distribution to decide which feature function subset to choose. With & chosen features, perform linear regression learn the coefficients of the features from data, and then update the k-subset distribution logit parameters by minimizing RMST.
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 the samples are collected from the Kuramoto Sivasssiisvvvkyky (KS) equation, a fourth-order nonlinear PDE known for its chaotic behavior. This PDE takes the form V/ = Uvx Venca UU2, which can be seen as linear combination of feature functions V {Vzza- Uzazza ) with the coefficients all set to value of 11. At test time, we the MAP estimation of the learned k-subset distribution choose the & feature functions. For k = 3. our proposed method achieves the same performance as the state-of-the-ar solver on this task, PySINDy. It identifies the KS PDE from data by choosing exactly the ground truth feature function subset V, obtaining RMSE of 0.00622 after applying linear regression on V.
COMPLEXITY ANALYSIS
In Proposition] I we prove that computing the marginal probability of the exactly-k constraint can be done tractably in time O(nk). In the context of deep learning, ofien about vectorized com-

---

Published a conference paper at ICLR 2023
plexity. We demonstrate an optimized algorithm achieving vectorized complexity Olog log n). assuming perfect parallelization. The optimization is possible by computing the marginal probability in a divide-and-conquee way: it partitions the variables into two subsets and compute their marginals respectively such that the complexity (n) reduced to (log n): the summation over the terms also has its complexity reduced to Ollo k) in a similar manner. We refer the readers to Algorithm S in Appendix for the optimized algorithm. We further modify Algorithm 2 to perform divide-and- conquer such that sampling k-subsets achieves vectorized complexity being Olog n). shown as
Algorithm/p in the Appendix. As a comparison, SoftSub (Xie & Ermon] (219) has its complexity to
be (nk)) due to the relaxed top-k operation and its vectorized complexity to be O log n) stemming from the fact that softmax layers need O(log n) rounds of communication for normalization.
7 RELATED WORK
There is large body of work gradient estimation for categorical random variables. Maddison al. (2017); Jang et (011)) propose the jumbel-softmax distribution (named the concrete distribution by the former) to relax categorical random variables. For more complex distributions, such as the k-subset distribution which are concerned with in this paper, existing approaches either use the straight-through and score function estimators or propose mallor-aade relaxations (see for instance Kim et al. (2016): [Chnn et al. (2018): Grover et al. (2018). We directly compare to the score function and straight- through estimator as well as the tailored relaxations of [hen et al. (2018); Grover al. (2018) and show that we are competitive and obtain a lower bias and/or variance than these other estimators. [Tucker ct al. (2017); Grathwohl et al. (2018) develop parameterized control variates based continuous relaxations for the score-function estimator. Lastly, (Paulus al. (2000) 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 & of them. That is with the exception Paulus et al. (2020) which make use of a relaxation using pairwise embeddings, but do not make their code available, We leave that to future work.
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 al. (201)); Correia et al. (201)); Martins & Astudillo (2016). SparseMAP ( (Niculac et al. (0188) has been proposed in the context of structured prediction and latent variable models,
also replacing the softmax with sparser distribution. LP- SP--paareMMPP (Niculae & Martins 2020) is
an extension that uses relaxation of the optimization problem rather than a MAP solver. Sparsity can also be exploited for efficient 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 sparse output. Also, we aim cases where samples in the forward pass are required.
Integrating specialized discrete algorithms into neural networks is line of research with increasing popularity. Examples are sorting algorithms (Cuturi et al. 2019: Blondel et al. 2020b Grover
al. 2018), ranking (Rolinek et al. 2020) Kool et al. 2019), dynamic programming (Mensch & Bloudel| Corro& Titov [2019) and solvers for combinatorial optimization problems Berthet
et al. 22020g Rolínek et al.] (2020); Shirobokov et al. (2020); Niepert el al. (202) Iy Minervini et al. (2023) Zeng et al. or even probabilistic circuits over structured output spaces (Ahmed et al. [022] Blondell 20119). There has also been work on making common programming language expression such as conditional statements, loops, and indexing differrnniabbee through relaxations (Petersen et al., 2021). Xie et al. (2020) proposes optimal transport
way obtain differentiable sorting methods for top-k classification. In contrast, we focus on the k-subset sampling problem and provide exact discrete sampling and marginal inference al algorithms, obtaining a gradient estimator for the k-subset distribution with favorable bias variana trade-off.
8 CONCLUSION
We introduced gradient estimator for the k-subset distribution which replaces relaxed and approxi- mate sampling on the forward pass with exact sampling. It sidesteps the non-differentiable nature of discrete sampling by estimating the gradients as a function of our distribution's 's marginals, for which we prove a simple characterization, showing that we can compute them exactly and efficiently. We demonstrated improved empirical results on number of tasks: 1.2X, DVAEs, and sparse regression.

---

Published conference paper at ICLR 2023
ACKNOWLEDGMENTS
We would like to thank all the reviewers whose insightful comments helped improve this work. We would especially like to thank reviewers for pointing out that we can forgo the method of finite differences in favor of computing the directional derivative, alleviating the need for X, and enabling us to compute exact derivative along the downstream gradient. This work was funded in part by the DARPA Perceptually- 1220011 NSF grants #1IS-1943641 #1IS- 1956441, #CCF-1837129, Samsung, CISCO, a Sloan Fellowship, and by Deutsche Forschungsgemeinschaff under Germany's 's Excellence Strategy- EXC 2075. ZZ is supported by an Amazon Doctoral Student Fellowship.
Task Guidance (PTG) Program under contract number
REFERENCES
Kareem Ahmed, Stefano Teso, Kai-Wei Chang, Guy Van den Broeck, and Antonio Vergari. Semantic prohabilistic layers for neuro-symbolic learning. In Advances in Neural Information Processing Systems 35 (NeurlPSS, dee 2022.
Quentin Berthet, Mathieu Blondel, Olivier Teboul, Marco Cuturi, Phin-ppilippe Vert, and Francis R. Bach. Learning with differentiable pertubed optimizers. In NewrlPS, 2020.
Mathieu Blondel. Structured prediction with projection oracles. Advances in neural information processing systems, 32, 2019.
Mathieu Blondel, André FT Martins, and Vlad Niculae. Learning with fenchel-young losses. J. Mach. Learn. Res., 21(35):1-69 69, 2020a.
Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In International Conference on Machine Learning, pp. 950-959. PMLR, 2020b.
Jianbo Chen, Le Song, Martin Wainwright, and Michael Jordan. Learning to explain: An information- theoretic perspective on model interpretation. In Jennifer Dy and Andreas Krause (eds.). Proceed- ings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 883-892. PMLR, 10-15 Jul 2018.
Gonçalo M Correia, Vlad Niculae, and André FT Martins. 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.
Gonçalo M Correia, Vlad Niculae, Wilker Aziz, and André FT Martins. Efficient marginalization discrete and structured latent variables via sparsity. Advances in Neural Information Processing Systems, 2020.
Caio Corro and Ivan Titov. Differentiable perturb-and-parse: Semi-supervised parsing with a structured variational auto-encoder. In International Conference on Learning Representations, 2019.
Marco Cuturi, Olivier Teboul, and Jean-Philipppe Vert. Differentiable ranking and sorting using optimal transport. Advances in neural information processing systems, 32, 2019.
Brian de Silva, Kathleen Champion, Markus Quade, Jean-Christophe Loiseau, J. Kutz, and Steven Brunton. Pysindy: python package for the sparse identification of nonlinear dynamical systems from data. Journal of Open Source Software, 5(49):2104, 2020. doi: 10.21 105/joss. 02144. URL
https: //doi org. 10. 211 05/ joss. 02104
Justin Domke, Implicit differentiation by perturbation. In Advances in Neural Information Processing Systems 23, pp. 523-531. 2010.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-hox gradient estimation. ICI.R, 2018.
10

---

Published conference paper at ICLR 2023
Aditya Grover, Eric Wang, Aaron Zweig, and Stefano Ermon. Stochastic optimization of sorting networks via continuous relaxations. In International Conference on Learning Representations, 2018.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparamcterization with gumbel-softmax. In Sth International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Alan A. Kaptanoglu, Brian M. de Silva, Urban Fasel, Kadierdan Kaheman, Andy J. Goldschmidt, Jared Callaham, Charles B. Delahunt, Zachary G. Nicolaou, Kathleen Champion, Jean- SannChhrittophe Loiseau, Nathan Kutz, and Steven Brunton. Pysindy: comprehensive python package for robust sparse system identification. Journal of Open Source Software, 7(69):3994, 2022. doi: 10.21 105/joss.03394
/ttttt:s: org. 1ttp//1.2111000 21105//oss...... 03994
Carolyn Kim, Ashish Sabharwal, and Stefano Ermon. Exact sampling with integer linear programs and random perturbations. In Thirtieth AAAI Conference on Artificial Intelligence, 2016.
Wouter Kool, Herke Van Hoof, and Max Welling. Stochastic beams and where to find them: The gumbel-top-k trick for sampling sequences without replacement. In International Conference on Machine Learning, pp. 3499-35098 PMLR, 2019.
Tao Lei, R. Barzilay, and T. Jaakkola. Rationalizing neural predictions. In EMNLP, 2016.
Chris Maddison, Daniel Tarlow, and Tom Minka. A sampling. In Z. Ghahramani, M. Welling, C. Cortes, Lawrence, and K.Q. Weinherger (eds.), Advances Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous relax- ation discrete random variables. In Sth International Conference on Learning Representations, ICIR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Andre Martins and Ramon Astudillo. From softmax sparsemax: A sparse model of attention and multi-labe classification. In International conference machine learming. pp. 1614-1623, PMLR, 2016.
Julian McAuley, J. Leskovec, and Dan Jurafsky. Learning attitudes and attributes from multi-aspect reviews. 2012 IEEE 12th International Conference on Data Mining, pp. 1020-1025, 2012.
Arthur Mensch and Mathieu Blondel. Differentiable dynamic programming for structured prediction and attention. In International Conference on Machine Learning, pp. 3462-3471. PMLR, 2018.
Pasquale Minervini, Luca Franceschi, and Mathias Niepert. Adaptive perturbation-based gradient estimation for discrete latent variable models. In AAAI, 2023.
Vlad Niculae and André F. T. Martins. Lpssparsemap: Differentiable relaxed optimization for sparse structured prediction. In ICMI, 2020.
Vlad Niculae, André Martins, Mathieu Blondel, and Claire Cardie. Sparsemap: Differentiable sparse structured inference. In ICML, 2018.
Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit mle: Backpropagating through discrete exponential family distributions. Advances in Neural Information Processing Systems, 34, 2021.
Max B Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris Maddison. Gradient estimation with stochastic softmax tricks. arXiv preprint arXiv:2006.08663 2020.
Ben Peters, Vlad Niculae, and André FT Martins. Sparse sequence-to-seve -e-uueeece models. Proceed- ings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1504-1519, 2019.
Felix Petersen, Christian Borgelt, Hilde Kuchne, and Oliver Deussen. Learning with algorithmic supervision via continuous relaxations. Advances in Neural Information Processing Systems, 34, 2021.

---

Published conference paper at ICLR 2023
Tobias Plötz and Stefan Roth. Neural nearest neighbors networks. Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 1195-1106, 2018.
Jason Tyler Rolfe. Discrete variational autoencoders. International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Michal Rolinek, Vit Musil, Anselm Paulus, Marin Vlastelica, Claudio Michaelis, and Georg Martius. Optimizing rank-based metries with blackbox differentiation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7620-7630, 2020.
Michal Rolinek, Paul Swoboda, Dominik Zietlow, Anselm Paulus, Vít Musil, and Georg Martius. Deep graph matching via blackbox diffcrentiation of combinatorial solvers. In ECCV, 2020.
Dan Roth. On the hardness of approximate reasoning. Artif. Intell. 82(1 (2(1-2)2273 1996. doi: (94)
(16600004-3702/9940009222 URL 00092-1
Sergey Shirobokov, Vladislav Belavin, Michael Kagan, Andrei Ustyuzhanin, and Atilim Gunes Baydin. Black-box optimization with local generative surrogates. Advances in Neural Information Processing Systems, 111144650-6666 2020.
George Tucker, Andriy Mnih, Chris Maddison, Dieterich Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances Neural Information Processing Systems, 2017.
Laurens van der Maaten. Learning a parametric embedding by preserving local structure. In David van Dyk and Max Welling (eds.). Proceedings of the Twelth International Conference on Artificial Intelligence and Statistics, volume of Proceedings of Machine Learning Research, pp. 384-391, Hilton Clearwater Beach Resort, Clearwater Beach, Florida USA, 16-18 Apr 2009. PMLR.
Ronald J. Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn., 1-22222266 1992. doi: 10. 10077/B000999266.
Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. In Proceedings of the Twenty-Figtth International Joint Conference on Artificial Intelligence, IJCAI-19, pp. 3919-3925. International Joint Conferences Artificial Intelligence Organization, 7 2019. doi: 00.24995/jcaii2211/55444
Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. Advances in Neural Information Processing Systems, 33:20520-20531, 2020.
Zhe Zeng. Paolo Morettin, Fanqi Yan, Andrea Passerini, Guy Van den Broeck, et al. Is parameter learning via weighted model integration tractable? In The 4th Workshop on Tractable Probabilistic Modeling.
Zhe Zeng, Paolo Morettin, Fanqi Yan, Antonio Vergari, and Guy Van den Broeck. Probabilistic inference with algebraic constraints: Theoretical limits and practical approximations. Advances Neural Information Processing Systems, 33:11564- 11575, 2020.
12

---

Published conference paper at ICLR 2023
A PROOFS
k) be the probability of exactit--k of the unconstrained distribution parameterized by log probabilities O. For every Z,. the conditional marginal is
Proposition 2.
Let pe 2,
0
p0
E3 23
k)
00,
log pe (E k)).
(10)
Proof. We first rewrite the marginal pe Ei 2s
k) into summation as the probability for all
possible events by definition as follows.
po 2 2j
Σ. = 11 |=, exp(03) 1,,=,=001 exp(0;))
(11)
Here assume that the probability of 2, 0 constant term w.r.t parameter 0j i.e. 80, exp(0,)) = 0 Further, the derivative of po 12, 2j = k) w.r.t. 8, is as follows.
@ (1
00; P0 E, 23 k) = 00;
Σ >,E,
II f=y==
cxp(0,) II (1 - cxp((8))
j=,==0
i
exp(O.)
II exp(O,) jay=
Il (1 exp(O,)) =:
08,
x:>i, =
exp(0:)
II exp(0,) f=
II (1-cx((0,) ):2y
2;),; kA=
po
EE 2j
where the first equality holds since terms corresponding to 2; # 1 has their derivative to be zero w.r.t. 0g further holds that
a log po po(E) E3 23 = k)=
30,
8
E3 23 = k)
00;
ro( = k:)
po E5 2j k =
po(E, 23 k)
pe
2z) 1 23 =
which finishes our proof.
Given variables Zi, Zn and a k-subset distribution po(z E, k) parameterized by 0, Algorithm computes entropy of Po (z E 24 k).
Proposition 3. Let Entropy be defined as in Algorithm 4
Proof. In slight abuse of notation, let 2n denote 2, = and let 2, denote žn 0. Furthermore, we denote by oh o and x the events Σi=o = k. D1- h: and ET- k- respectively.
The entropy of the susbett distribution is given by
HZ)) E3 (o [log [(0](()] = Σ po(= ( log a
We start by simplifying the expression for po(z ( where, by the chain rule, the above is
Σ 20( o po(o o zn) p= | o po(o oo 2n)
"In practice this be easily implemented. For example, in framework Tensorflow it can be done setting stop_graad ents.

---

Published conference paper at ICLR 2023
Plugging the above the expression for the entropy, distributing the sum over the product, we get
Σ (ān ( ) pol on kn)
x:o
llog (Zn loa) po(o 1 on 2n) + po(zn -pooo o 2n.)
pa (Zn
(o)) po po 1 o 2n)
log (2n o po(on-11 on 2n) + pol ( po o 2n)]
where, since the two events &n and 2n are mutually exclusive, we can simplify the above to
- pa ( po(o2-1 o En) log [pa(E, o.)) po(o"-1 o..
+ po(=n [x po(o xx =n) log [zo ) po po(o zn)]
Expanding the logarithms, rearranging terms, and using that conditional probabilities get
po( (xai) log (2n o + po(=n I loa)) rogz |,,1) + po 1o/) po(oo oh 2n) log (a, o kn)
+ pal | 1.hh) po(o k-1 of 2n)) log (o (ot-1 oit zm)
-Ps,,-puss.lll [- pe( x + Ssa.soss lot))
(2 2n)
-Ho(Zn (.) + (o) H (Z =n)]
That is, simply stated, the entropy the s-subset distribution decomposes as the entropy of the construined distribution over Zn, and average entropy of the distribution on the remaining variables.
As the base case, the entropy of the susett distribution when k = n is O: there is only one way in which pick choose of n variables, and the k-subset distribution is therefore deterministic. I
B OPTIMIZED ALGORITHMS
Algorithm is the optimized version of Algorithm I both of which compute the marginal probability of the exactly-k: constraint. Algorithm 6]is the optimized version of Algorithm both of which sample faithfully from the k-subset distribution.
Algorithm 5
PrExactlyyk00,,.uu
Algorithm 6 Sample(0,1, u,)
Input:
The logits of the distribution, range of
Input: The logits of the distribution, range variable indices [1, u). and the subset size k
variable indices [2, u). and the subset size &
Output: The exact marginal probability of vari- ables summing up to k, P(EK) X; A))
Output: A sample = (21, 2.) from po(2
Σ 24 k))
if > u then return 0
define p(æ = m) Pm. m 0, ..... k
i then return Po(X: k)
with Pm defined in Algorithm
for m = to & do
sample m from p
Pm (r [u/2].m)=
#r|,//] = Sample 1, [u/2], m²)
PrExactlyk(A 0 [u/2] + +, m)
z [//2=+1 = Sample(O0 Iu/2 + u. & mm))
return > Em=oPm Pm
return Concat(*), w/2) [
C EXPERIMENTAL DETAILS
C.I SYNTHETIC EXPERIMENTS
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 the the gradients estimators. the variance of the gradient estimators, as well as the average deviation of each estimated gradient from the exact gradient. We used cosine distance, defined as I- cosine similarity measure of

---

Published conference paper at ICLR 2023
distance, as we typically care about the direction, not the magnitude of the gradient; the latter can be recovered using an appropriate learning rate. Following Niepert et al. (2021), we chose a tractable 5-subset distribution, where n 10, and were therefore limited to (s) 252 possible subsets. We set the loss ((0)- l8 z-pp (xIE, =c-R) k) [|| bll"), where b is the groundtruth logits sampled from (0, I). We used a sample size of 10000 to estimate each of our metrics.
C.2 DISCRETE VARIATIONAL ENCODER
We tested our SIMPLE gradient estimator in the discrete k-subset Variational
(VAE)
setting, where the latent variables model a probability distribution over k-subsets, and has a dimen- sionality of 20. The experimental setup is similar to those used in prior work on the Gumbel softmax tricks (Jang et al. 2017) and IMLE (Niepert et al. (022) The encoding and decoding functions the VA consist of three dense layers (encoding: 512-256-20x20; decoding: 256-512-784) 12-784). As is commonplace discrete VAΣs, the loss is the sum of the reconstruction loss (binary cross-entropy loss on output pixels) and KL divergence
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, use 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), 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 x 20) which outputs 0 and no non-linearity SIMPLE takes 0 as input and outputs a discrete latent code of size 20 x 20. The decoder network, which takes this discrete latent code input, consists of a dense layer with dimension 256 and ReLu activation, a dense layer with dimension 512 and ReLu activation, and finally a dense layer with dimension 784 returning the logits for the output pixels. Sigmoids are applied to these logits and the binary cross-entropy loss is computed. To obtain the best performing model of each of the compared methods, we performed grid search over the learning rate in the range 10-3,
[I x 10 3 5 x 10-4) À in the range [1 x 10-3 1 x 10 10-2, 1 x 10-1, I x 100, 1 x 10t x 10?, x 10%) and for SoG I-MLE, the temparature in the range [1 x 10 1, x 100, x 10t, x 10a]
We will present formal proof on how to compute the KL-divergence between the k-subset distribution and a uniform distribution tractably and exactly.
Proposition 4. Let po (z Σ; z = k) be a k-subsel distribution parameterized by O and (() be uniform distribution on the construined space C = {= 1E 2 = k). Then the KT-divergence
between distribution po (z Σ, 2: = k) and (z) can computed by
Σ = =) (12(()) = --[(2)) + log
where H denote the entropy of distribution pa (z Σ, = k).
Proof. By the definition of KL divergence, it holds that
D(po(z Σ 24 = k) (()))
Σno Co Σ k) log (2 (z)
Σ = k)
U(z)
#EC
= 2( pz(z
Σ = k) logpo(z Σ 21 = k))
Epli Σx = k) log (z))
sEC
-z( log
The last equality holds since (z) 1/)))
C.3 LEARNING TO EXPLAIN
The BEERADVOCATE dataset (McAuley et al. 2012) consists of free-text reviews and ratings for 4 different aspects of heer: appearance, aroma, palate, and taste. The training has 80k reviews for
15

---

Published conference paper at ICLR 2023
the aspect APPEARANCE and 70k reviews for all other aspects. The maximum review length 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. addition the ratings for all reviews, each sentence the test contains annotations of the words that best describe the review score with respect to the various aspects. Following the experimental setup of recent work (Paulus et al. 2020; Niepert al. (02T). we address the problem introduced by the L2X paper (Chen al. 2018) of learning 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 architecture from prior work work(Chen al. (2018); Paulus et al. (2020) with convolutional and one dense layer. This neural network outputs the parameters of the k-subset distribution over
k-hot binary latent masks with & É {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 the model with the best validation MSE among the 20 epochs.
hhtt::/pcople.coo
esail. mit. cduttaocc//bberr.


Reviews:
[{"review_id": "eAEr2ProSvU", "reviewer": "Reviewer_AhEr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper provides a gradient estimator when the computation graph includes a k-subset sampling module.", "summary_of_the_review": "The paper provides a gradient estimator when the computation graph includes a k-subset sampling module, which is new to me.", "strength_and_weaknesses": "* Strength \n\nThe paper provides a gradient estimator when the computation graph includes a k-subset sampling module.\n\n* Weakness\n\nComparisons with other methods can be made more comprehensive.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "bsVXCn6mTv", "reviewer": "Reviewer_XcRS", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents an approach for gradient estimation for discrete k-subset sampling, called SIMPLE.  For the forward pass, this approach involves discrete sampling.  For the backward pass, SIMPLE efficiently computes the gradient with respect to the exact marginals of the $k$-subset distribution.  Compared to competing approaches for gradient estimation in this setting, SIMPLE exhibits lower bias and variance.  In addition to an efficient algorithm for computing the gradient, the authors also present an algorithm for computing the exact ELBO for $k$-subset distribution, leading to state-of-the-art performance on sparse VAE learning.", "summary_of_the_review": "There are some important empirical and theoretical contributions in this paper. The experimental results are somewhat convincing, showing that the proposed SIMPLE method outperforms a number of competing methods for estimating gradients for the discrete k-subset setting, in terms of both reduced bias and variance, and improved predictive performance.  However, there are several issues with this paper that should be addressed by the authors, as mentioned in the list of weaknesses above.  Overall, in terms of acceptance, this is a borderline paper in its current state.", "strength_and_weaknesses": "Strengths:\n* Compared to competing approaches for computing gradients for k-subset sampling, SIMPLE provides significant reductions in bias and variance.  Compared to competing approaches, SIMPLE can also provide moderate to significant improvements in predictive performance when used for training models, as shown by the experimental results in this paper for the learning to explain task.\n* When used to compute the ELBO for a $k$-subset discrete VAE, compared to competing approaches SIMPLE can provide significantly lower ELBO loss on the test set.\n* The complexity analysis presented in this paper shows that the proposed SIMPLE algorithm appears to scale when computing the required marginal probabilities and sampling k-subsets, particularly when using vectorized computations for parallelization.\n\nWeaknesses:\n* The complexity analysis provided in the paper focuses on algorithms for computing marginals probabilities and for sampling $k$-subsets.  However, an explicit analysis of the complexity for computing the backward pass (from Algorithm 3) appears to be missing.  This should be provided in the paper.\n* This paper does not include a discussion of the weaknesses/drawbacks of the proposed approach.  For example, SIMPLE seems to assume that the $k$-subset distribution can be factorized as a product of the conditional marginals; that is, the subset is a multilinear function of its elements (as pointed out in the first paragraph of Sec. 3.1).  However, this assumption may only be valid in cases where the $k$-subset distribution does not involve higher-order interactions between elements within a subset, such as subsets that are modeled using a Transformer-based model.  An explicit discussion of this limitation should be provided in the paper.\n* An empirical evaluation of the runtime of SIMPLE, compared to the competing approaches, is missing and should be provided.  While the complexity analysis appears to show that SIMPLE scales, an empirical validation of this scalability is important.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RFHqq13fjp", "reviewer": "Reviewer_pM21", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes novel gradient estimator for k-subset sampling. Authors justify their method with rigorous proofs and conduct experiments to support their claims. Experiments have major flaws that should be addressed.", "summary_of_the_review": "The proposed sampler is novel while experimentation section of the paper is flawed and requires clarifications.", "strength_and_weaknesses": "Strength:\n-- efficient and novel gradient estimator\n-- experiments on synthetic and real data\n-- rigorous theoretical analysis\n\nWeaknesses:\n-- choice of baselines in experiments seems odd, authors claim that their estimator is better then Gumbel softmax estimator but do not present comparison of it's quality on real data\n-- improper comparison with Gumbel Softmax estimator, algorithm proposed by authors has asymptotic worse then of Gumbel SoftMax by a factor of k therefore I wonder whether their reduction in variance and bias will be significant over GS if we choose number of samples for it greater then 1 as authors did to match running time of their estimator.\n-- it's not clear why GS and proposed estimator have bias at all -- GS based on log-likelihood trick should be unbiased and proposed estimator is claimed to exactly compute probabilities therefore it should be unbiased too (or am I missing something here?). Backprop cannot increase bias if it's absent since chain rule is basically reduces to application of linear mapping on unbiased estimate.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "DcZcz1C2oj", "reviewer": "Reviewer_6jdF", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper suggest a k-subset sampling gradient estimator, namely SIMPLE. The proposed SIMPLE estimator can replace the reparameterization by relaxation with exact gradient estimation, and utilizes the gradient with respect to the exact conditional marginals in the backpropagtion. The authors conducted various experiments to demonstrate the performance of SIMPLE estimator: synthetic example, discrete variational autoencoder, learning to explain, and sparse linear regression.\n", "summary_of_the_review": "I agree with the methodology, but for some experiments, the explanation of the experimental method is insufficient, and the reliability of the results is insufficient. Also, the presentation of the paper difficult to understand. I therefore lean negatively on this paper.\n", "strength_and_weaknesses": "The paper studies the exact gradient estimator for the n-choose-k distribution. The idea of this paper seems quiet novel. Also, it seems that the proposed method achieved better result in bias/variance, ELBO of DVAE, etc. \n\nHowever, the paper is very hard to follow, not because of the difficulty of the paper, but due to the presentation of the paper. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster