Reparameterizable Subset Sampling via Continuous Relaxations

1 Introduction

Sampling a single item from a collection is common in machine learning problems such as generative modeling with latent categorical variables, attention mechanisms [Kingma et al., 2014; Xu et al., 2015]. These tasks involve optimizing an expectation objective over a latent categorical distribution parameterized by a neural network. Score-based methods such as REINFORCE [Williams, 1992] for estimating the gradient of such objectives typically have high variance. The reparameterization trick [Kingma and Welling, 2013] allows for low variance gradients for certain distributions, not typically including categorical distributions. The Gumbel-softmax trick [Jang et al., 2016] or Concrete distribution [Maddison et al., 2017] are continuous relaxations that allow for reparameterized gradients with respect to the parameters of the distribution. Among many others, this enabled generative modeling with latent categorical variables without costly marginalization and modeling sequences of discrete elements with GANs [Jang et al., 2016; Kusner and Hern´andez-Lobato, 2016].

In this paper, we consider the more general problem of sampling a subset of multiple items from a collection without replacement. As an example, choosing a subset is important in instance-wise feature selection [Chen et al., 2018], where the goal is to select a subset of features that best explain the model's output for each example. Sampling subsets of neighbors also enables implementing stochastic k-nearest neighbors end-to-end with deep features. Stochastic optimization involving subset sampling, however, does not typically have relaxations with low-variance reparameterization gradients as in Gumbel-softmax. To overcome this limitation, we develop a continuous relaxation for approximate reparameterized gradients with respect to the parameters of a subset distribution to enable learning with backpropagation. In our setting, the Gumbel-max trick (and thus Gumbel-softmax) is not directly applicable since it requires treating every possible subset as a category, requiring a combinatorial number of categories. We use an extension to the Gumbel-max trick which perturbs the log-probabilities of a categorical distribution with Gumbel noise and takes the top-k elements to produce samples without replacement. Ignoring ordering in these samples allows for sampling from a family of subset distributions using the same algorithm. We give a general algorithm that produces continuous relaxations with reparameterization gradients using top-k relaxations. We then show that a recent topk relaxation [Pl¨otz and Roth, 2018] can be used in our algorithm and study the consistency of this top-k relaxation.

Our main contributions are the following:

• We give an algorithm for a reparameterizable continuous relaxation to sampling subsets using top-k relaxations and a extension to the Gumbel-max trick.
• We show that the top-k relaxation of [Pl¨otz and Roth, 2018] is consistent in the sense that the ordering of inputs is preserved in the output in many practical settings.
• We test our algorithm as a drop-in replacement for subset selection routines in explaining deep models through feature subset selection, training stochastic neural knearest neighbors, and implementing parametric t-SNE without Student-t distributions by directly comparing neighbor samples. We improve performance on all tasks using the same architectures and metrics as the original 1 .

1 Code available at https://github.com/ermongroup/subsets.

Algorithm 1 Weighted Reservoir Sampling (non-streaming)

Output: Swrs = [ei1 , . . . , eik ] a sample from p(Swrs|w)

• 1: r ← [ ]
• 2: for i ← 1 to n do
• 3: ui ←Uniform(0, 1)
• 4: ri ← u1/wi i
• 5: r.append(ri)
• 6: end for
• 7: [ei1 , . . . , eik ] ←TopK(r, k)
• 8: return [ei1 , . . . , eik ]

# Sample random keys

2 Preliminaries

2.1 Weighted Reservoir Sampling

Reservoir sampling is a family of streaming algorithms that is used to sample k items from a collection of n items, x1, . . . , xn, where n may be infinite [Vitter, 1985]. We consider finite n and only produce samples after processing the entire stream. In weighted reservoir sampling, every xi is associated with a weight wi ≥ 0. Let w = [w1, . . . ,wn] and Z = ∑ni=1 wi be the normalizing constant. Let ej = [ej1 , . . . , ejn ] = [0, · · · , 0, 1, 0, · · · , 0] ∈ {0, 1}n be a 1-hot vector, i.e., a vector with only one nonzero element at index j, where ejj = 1. We define a weighted reservoir sample (WRS) as Swrs = [ei1 , . . . , eik ], a sequence of k 1-hot (standard basis) vectors where eij represents selecting element xij in the j-th sample. We wish to sample Swrs from

which corresponds to sampling without replacement with probabilities proportional to item weights. Modeling samples without replacement allows for sampling a sequence of distinct items. For k = 1, p(Swrs|w) is the standard softmax distribution with logits given by log(wi).

[Efraimidis and Spirakis, 2006] give an algorithm for weighted reservoir sampling (Algorithm 1). Each item xi is given a random key ri = u1/wi i where ui is drawn from a uniform distribution between [0, 1] and wi is the weight of item xi. Let the top k keys over the n items be ri1 , . . . , rik . We define the function TopK(r, k) which takes keys r = [r1, . . . , rn] and returns [ei1 , . . . , eik ] associated with the topk keys. The algorithm uses TopK to return the items Swrs = [ei1 , . . . , eik ] as the WRS. Efraimidis and Spirakis proved (Proposition 5 in [Efraimidis and Spirakis, 2006]) that the output of Algorithm 1 is distributed according to p(Swrs|w).

2.2 Gumbel-max Trick

Given w as in (1), log(wi) are logits for a softmax distribution p(xi|w) = wi/Z. The Gumbel-max trick [Yellott, 1977] generates random keys ˆ ri = log(wi) + gi by perturbing logits with Gumbel noise gi ∼ Gumbel(0, 1), then taking xi∗ such that i∗ = argmaxi ˆ ri as a sample. These samples are distributed according to p(xi|w) = wi/Z.

The idea is to reparameterize the sample as a deterministic transformation of the parameters w and some independent noise gi. Then by relaxing the deterministic transformation (from max to softmax), the Gumbel-softmax trick allows for training with backpropagation [Maddison et al., 2017; Jang et al., 2016]. Similarly, we use an extension of the Gumbel-max trick to decouple the deterministic transformation of the parameters (in our case, a top-k selection function) and the randomness (Gumbel noise gi), and we relax the topk function to allow for backpropagation.

3 Reparameterizable Continuous Relaxation for Subset Sampling

3.1 Setup

We represent a subset S ∈ {0, 1}n as a k-hot vector, which is a vector with exactly k nonzero elements that are all equal to 1. We define the probability of a subset S as the sum of the probabilities of all WRS with the same elements

where Π(S) = {Swrs : S = ∑kj=1 Swrs[j]} is the set of all permutations of elements in S represented by sequences of 1-hot vectors. Here, Swrs[j] is the j-th 1-hot vector in the sequence. By simply ignoring the order of elements in a WRS, we can also sample from p(S|w) using the same algorithm. Note that this is a restricted family of subset distributions. Since each distribution is over (nk) subsets of size k, the full space of distributions requires (nk) - 1 free parameters. Here, we reduce the number of free parameters to n-1. While this is a restriction, we gain tractability in our algorithm.

3.2 Gumbel-max Extension

We extend the Gumbel-max trick to sample from p(S|w). The main intuition is that the outputs of the reservoir sampling Algorithm 1 only depend on the ordering of the random keys ri and not their values. We show that random keys ˆ ri generated from the Gumbel-max trick are monotonic transformations of the random keys ri from Algorithm 1. Because a monotonic transformation preserves the ordering, the elements that achieve the top-k Gumbel-max keys ˆ ri have the same distribution as the elements that achieve the top-k weighted reservoir sampling keys ri. Therefore we can sample from p(S|w) by taking the top-k elements of ˆ ri instead.

To make the procedure differentiable with respect to w, we replace top-k selection with a differentiable approximation. We define a relaxed k-hot vector a = [a1, . . . , an] to have ∑ni=1 ai = k and 0 ≤ ai ≤ 1. We relax Swrs by replacing all eij with relaxed 1-hot vectors and relax S by a relaxed k-hot vector. Our continuous relaxation will approximate sampling from p(S|w) by returning relaxed k-hot vectors. We use a top-k relaxation RelaxedTopK, which is a differentiable function that takes ˆ ri, k, and a temperature parameter t > 0 and returns a relaxed k-hot vector a such that as t → 0, RelaxedTopK(ˆr, k, t) → ∑kj=1 TopK(ˆr, k)[j], where TopK(ˆr, k)[j] = eij is the 1-hot vector associated with

Algorithm 2 Relaxed Subset Sampling

Input: Items x1, . . . , xn, weights w = [w1, . . . ,wn], subset

size k, temperature t > 0

Output: Relaxed k-hot vector a = [a1, . . . , an], where

1: ˆr ← [ ]

∑ni=1 ai = k, 0 ≤ ai ≤ 1

2: for i ← 1 to n do

3: ui ←Uniform(0, 1)

# Random Gumbel keys

4: ˆ

ri ← - log(- log(ui)) + log(wi)

5: ˆr.append(ˆ

ri)

6: end for

7: a ←RelaxedTopK(ˆr, k, t)

8: return a

the j-th top key in ˆr. Thus Algorithm 2 produces relaxed khot samples that, as t → 0 converge to exact samples from p(S|w). Note that we can also produce approximate samples from p(Swrs|w) if an intermediate output of RelaxedTopK is a sequence of k relaxed 1-hot vectors [ai1 , . . . , aik ] such that as t → 0, [ai1 , . . . , aik ] → TopK(ˆr, k). This means that the intermediate output converges to a WRS.

Proposition 1. Let RelaxedTopK be defined as above. Given n items x1, . . . , xn, a subset size k, and a distribution over subsets described by weights w1, . . . , wn, Algorithm 2 gives exact samples from p(S|w) as in (2) as t → 0.

Proof. Let the random keys in Algorithm 2 be ˆr and the random keys in weighted reservoir sampling be r. For any i,

Fixing ui, since -log(-log(a)) is monotonic in a, ˆ ri is a monotonic transformation of ri and TopK(ˆr, k) = TopK(r, k). Let TopK(r, k) be samples from Algorithm 1. By construction of (2), ∑kj=1 TopK(r, k)[j] is distributed as p(S|w). As t → 0, Algorithm 2 produces samples from p(S|w) since RelaxedTopK(ˆr, k, t) → ∑kj=1 TopK(ˆr, k)[j] = ∑kj=1 TopK(r, k)[j].

This fact has been shown previously in [Vieira, 2014] for k = 1 and in [Kim et al., 2016] for sampling from p(Swrs|w) without the connection to reservoir sampling. Kool et al. concurrently developed a similar method for ordered sampling without replacement for stochastic beam search. Note that Algorithm 2 is general to any choice of top-k relaxation.

3.3 Differentiable Top-k Procedures

A vital component of Algorithm 2 is a top-k relaxation that is differentiable with respect to the input keys ˆ ri (a random function of w). This allows for parameterizing w, which governs p(S|w), using neural networks and training using backpropagation. We propose to use a recent top-k relaxation based on successive applications of the softmax function [Pl¨otz and Roth, 2018]. For some temperature t > 0,

where aji is a sample at step j from the distribution

for j = 1, . . . , k steps. In the relaxation, the aji is replaced with its expectation, p(aji = 1), such that the new update is

and the output is ai = ∑kj=1 p(aji = 1) for each i. Let aj = [p(aj1 = 1), . . . , p(ajn = 1)] be relaxed 1-hot outputs at step j and a = ∑kj=1 aj be the relaxed k-hot output. Pl¨otz and Roth show that as t → 0, [a1 , . . . , ak ] → TopK(ˆr, k) and thus a → ∑kj=1 TopK(ˆr, k)[j] so that this is a valid RelaxedTopK. Thus the relaxation can be used for approximately sampling from both p(S|w) and p(Swrs|w). Next we show that the magnitude of values in the output relaxed k-hot vector a preserves order of input keys ˆ ri for t ≥ 1.

Theorem 1. Given keys ˆr, the top-k relaxation of [Pl¨otz and Roth, 2018] produces a relaxed k-hot vector a where if ˆ ri ≤ ˆ rj , then ai ≤ aj for any temperature t ≥ 1 and k ≤ n.

Proof. By induction on k. Fix any ˆ ri ≤ ˆ rj. For step k = 1, we have a1i ≤ a1j since the softmax function preserves ordering. Assuming the statement holds for 0, . . . , k, we want to show that αk+1 i ≤ αk+1 j , which suffices to imply ak+1 i ≤ ak+1 j by the order-preserving property of softmax. Define

Comparing αk+1 j and αk+1 i through the ratio,

where we can view c = ∑nm=1 exp(αkm/t) - exp(αk

i /t) - exp(αk j /t) ≥ 0 as a non-negative constant in this analysis. Note that exp(αkj ) exp(αki ) = exp(αk j -αk i ) ≥ exp((αk j -αk i )/t) = exp(αkj /t) exp(αki /t) for t ≥ 1. Therefore

for any c ≥ 0. Thus αk+1 i ≤ αk+1 j , implying ak+1 i ≤ ak+1 j .

Therefore, this top-k relaxation is consistent in the sense that the indices with the top k weights in w are also the top k values in the relaxed k-hot vector for many reasonable t. In the limit as t → 0 the consistency property holds for the exact k-hot vector, but the relaxation is not consistent in general for 0 < t < 1. To see a concrete example of the loss of consistency at lower t, take ˆr = [e, e2 ], k = 2, and t = 0.4. The output on this input is [1.05, 0.95], which has inconsistent ordering. Consistency at higher temperatures is important when considering the bias-variance tradeoff in the gradients produced by the approximation. While higher temperatures allow for lower variance gradients, the bias introduced may cause the model to optimize something far from the true objective. Consistency in ordering suggests that the top-k information is not lost at higher temperatures.

Note that any top-k relaxation can be used in Algorithm 2. For example, [Grover et al., 2019] give a relaxation of the sorting procedure through a relaxed permutation matrix P. To produce a relaxed k-hot vector in this framework, we can take a = ∑kj=1 Pj where Pj is the j-th row of P. Similarly, when the associated temperature parameter approaches 0, P approaches the exact permutation matrix so that the relaxed k-hot vector a also approaches the exact top-k k-hot vector.

4 Experiments

4.1 Synthetic Experiments

We check that generating samples from Algorithm 2 results in samples approximately from p(S|w). We define a subset distribution using weights w = [0.1, 0.2, 0.3, 0.4] and take subset size k = 2. Using Algorithm 2 and the top-k relaxation from [Pl¨otz and Roth, 2018], we sample 10000 relaxed k-hot samples for each temperature in {0.1, 1, 10} and take the the top-k values in the relaxed k-hot vector as the "chosen" subset. We plot the empirical histogram of subset occurrences and compare with the true probabilities from the subset distribution, with 95% confidence intervals. The relaxation produces subset samples with empirical distribution within 0.016 in total variation distance of p(S|w) for all t. This agrees with Theorem 1, which states that even for higher temperatures, taking the top-k values in the relaxed k-hot vector should produce true samples from (2).

i just saw this at the - film festival it was the most wonderful movie the best i've seen in quite a while the - character of is an open - and as - in the film love filled person portrayed by the beautiful - the - music is - and - a - to the soul each character seems to be - in some way and their - w conflicts make for a great story the tale told by a k - as the old man was most - done i wanted it to go on and on that - seemed to remember his place throughout added get power to the story a refreshing change to the - headed plots of many modern writers all and all an excellent film go see

Figure 2: An example of an IMDB review where our explainer model predicts the k = 10 words and the sentiment model outputs a similar prediction (positive) when only taking the 10 words as input. Words not in the vocabulary are replaced with dashes.

read the book forget the movie

Figure 3: A review where the L2X explainer model [Chen et al., 2018] can select 10 words (whole review) but subsamples, reducing the sentiment model's confidence from 0.93 to 0.52 when taking the subset of words as input. Our explainer returns the whole review.

4.2 Model Explanations

We follow the L2X model [Chen et al., 2018] and set up the problem of explaining instance-wise model outputs by training an auxiliary explainer model to output the k features with the highest mutual information with the model output. For example, given a movie review and a sentiment model, we want to select up to k words from the review that best explains the sentiment model's output on the particular review (see Figure 2). Since optimizing the mutual information is intractable, L2X optimizes a variational lower bound instead:

where E : Rd → Pk is an explainer model, parameterized as a neural network, mapping from an input to the space of all k-hot vectors S. The approximating distribution is q(XS) where q is a neural network and XS = S ⊙X ∈ Rd is X with the elements not corresponding to S zeroed out. L2X approximates the subset sampling procedure by sampling k independent times from a Concrete distribution [Maddison et al., 2017] and taking the elementwise maximum. Since each independent Concrete sample is a relaxed 1-hot vector, the L2X explainer model may suffer by independently sampling the same feature many times, resulting in selecting less than k features. Our model differs by replacing this sampling procedure with Algorithm 2. Figure 3 shows an instance where the L2X explainer model chooses an ambiguous subset of the review even when the review is shorter than the maximum number of words to be selected (k), which reduces the sentiment model's confidence significantly (0.93 to 0.52). Our model selects the entire review in this case, so the sentiment model's output is not affected.

We test our results on the Large Movie Review Dataset (IMDB) for sentiment classification [Maas et al., 2011], where we select the most important words or sentences that contribute to the sentiment prediction for the review.

The original model for word-based sentiment classification (IMDB-word) is a convolutional neural network [Kim, 2014], while the original model for sentence-based sentiment prediction is a hierarchical LSTM (IMDB-sent) [Li et al., 2015]. The explainer and variational distribution models are CNNs with the same architectures as in L2X [Chen et al., 2018]. Following L2X, we use k = 10 for IMDB-word and k = 1 sentences for IMDB-sent. At test time, all explainer models deterministically choose subsets based on the highest weights instead of sampling. We evaluate using post-hoc accuracy, which is the proportion of examples where the original model evaluated on masked features XS matches the model on unmasked X. We use cross validation to choose temperatures t ∈ {0.1, 0.5, 1, 2, 5} according to the validation loss. Our model (RelaxSubSample) improves upon L2X by up to 1% by only changing the sampling procedure (Table 1).

4.3 Stochastic K-Nearest Neighbors

We give a stochastic k-NN algorithm where we use deep features tuned end-to-end for computing neighbors. We follow the setup of Grover et al., using m = 100 randomly sampled neighbor candidates and the same loss. We defer details, including architecture, to Grover et al..

We compare to NeuralSort, which implements kNN using a relaxation of the sorting operator [Grover et al., 2019]. We fix k = 9 nearest neighbors to choose from m candidates and search over temperatures t = {0.1, 1, 5, 16, 64} using the validation set, whereas NeuralSort searches over both k and t. RelaxSubSample approaches the accuracy of a CNN trained using the standard cross entropy loss on MNIST (99.3% vs. 99.4%) and increases accuracy by 0.6% over the NeuralSort implementation on CIFAR-10 (Table 2).

Note that NeuralSort implements a relaxation to the sorting procedure, while in kNN we only require the top-k elements. We use the top-k relaxation from [Pl¨otz and Roth, 2018], which computes k softmaxes for a runtime and storage of O(km). NeuralSort requires O(m2 ) time and storage as it produces a m × m permutation matrix for each input. Table 3 shows forward pass runtimes of a CNN using our method and NeuralSort for different values of m and k = 5 neighbors. While runtimes for small m are comparable, our method scales much better for larger m (Table 3).

4.4 Stochastic Neighbor Embeddings

We consider the problem of learning a parametric embedding which maps a high-dimensional space into a lowerdimensional space while preserving local neighbor structure. This problem is addressed by parametric t-SNE [van der Maaten, 2009], which represents pairwise densities of datapoints in the high-dimensional space as symmetrized Gaussians with variances tuned so that the perplexity of each point is equal. The pairwise densities in the low-dimensional space are modeled by Student-t distributions to address the crowding problem, a result of the volume shrinkage from the high to low-dimensional spaces. Student-t distributions have heavier tails, allowing for distant points in the high-dimensional space to be modeled as far apart in the low-dimensional space. The objective is to minimize the KL divergence between the pairwise distributions in the two spaces.

We propose to learn such a mapping without the use of Student-t distributions. Our insight is that the goal of preserving local structure can be achieved by preserving the neighbor rankings, which is insensitive to scaling of distances. In our RelaxSubSample-based stochastic neighbor embedding (RSS-SNE), we aim to preserve the distribution of k neighbors around each point. We define the neighbor distributions as follows. Let X = {x1, . . . , xn} be the training data. Let w(i, j) = exp(-‖xi - xj‖22 ) be exponentiated negative pairwise squared distances. For datapoint xi, let w(i) ∈ Rn-1 be the pairwise distances from xi to other points. Then we model the neighbor distribution for xi as p(Swrs|w(i)) as in (1). Note that we sample sub-sequences because we want to preserve neighbor rankings. Let h be our parametric embedding function. Similarly, letting ˆw(i, j) = exp(-‖h(xi) - h(xj)‖22 ) and the pairwise distances involving h(xi) be ˆw(i), the neighbor distribution in the low dimensional space is p(Swrs| ˆw(i)). We aim to match these distributions by comparing the neighbor samples to avoid the crowding problem. For each xi, let a sample from p(Swrs|w(i)) be [ei1 , . . . , eik ] where eij are 1-hot vectors corresponding to a selected neighbor xij . Let a relaxed sample from p(Swrs|ˆw(i)) be [ai1 , . . . , aik ], k intermediate relaxed 1-hot outputs of the top-k relaxation from [Pl¨otz and Roth, 2018]. We minimize the objective

where we aim to match samples from the neighbor distributions in the two spaces, putting more weight on matching closer neighbors. While we can directly match the neighbor distributions p(Swrs|w(i)) and p(Swrs|ˆw(i)) without sampling, they are defined in terms of pairwise distances that cannot be matched due to the crowding problem. We find that sampling from both distributions is necessary to keep the loss agnostic to scaling of distances in the different spaces.

We compare with parametric t-SNE [van der Maaten, 2009] on the MNIST [LeCun and Cortes, 2010] and a small version of the 20 Newsgroups dataset [Roweis, 2009] 2 . Following van der Maaten, we compare the trustworthiness and performance of 1-NN classifiers of the low dimensional embeddings on the test set. Trustworthiness [Venna and Kaski, 2006] measures preservation of local structure in the low dimensional embedding and is defined as

where r(i, j) is the rank of datapoint j according to distances between datapoint i and other datapoints in the highdimensional space, and N (k) i is the set of k nearest neighbors in the low dimensional space. Trustworthiness decreases when a datapoint is in the k nearest neighbors in the low dimensional space but not the original space. As in van der Maaten, we use T (12), comparing 12 nearest neighbors.

We use the same feedforward networks as in van der Maaten as embedding functions. For all experiments, we set t = 0.1 and train for 200 epochs with a batch size of 1000, choosing the model with the best training loss. We sample neighbors only within each training batch. We find that k = 1 is sufficient to learn the local structure. For k > 1, we observe a trade-off where 1-NN accuracy decreases but trustworthiness either stays the same or increases. It was important to add a small bias (1e - 8) to the relaxed k-hot vectors for better optimization. We compare two versions of RSS-SNE, one trained from scratch and another using layerwise pretraining. We pretrain layer l by treating layers 1, . . . , l as an encoder and adding a 1 layer decoder to the original space; then we optimize the MSE autoencoder objective for 10 epochs. Note that the original parametric t-SNE used a similar layerwise pretraining scheme using RBMs. RSS-SNE models consistently have higher trustworthiness and competititive 1-NN test errors when compared to parametric t-SNE (Tables 4, 5). Since trustworthiness compares 12 nearest neighbors, this suggests that our embedding has better overall structure as opposed to focusing on the immediate neighbor.

5 Conclusion

We present an algorithm for relaxing samples from a distribution over subsets such that the procedure can be included in deep models trained with backpropagation. We use the algorithm as a drop-in replacement in tasks requiring subset sampling to boost performance. Our algorithm has the potential to improve any task requiring subset sampling by tuning the model end-to-end with the subset procedure in mind.

Acknowledgements

This research was supported by NSF (#1651565, #1733686), ONR, AFOSR (FA9550- 19-1-0024), and Bloomberg. We are thankful to Aditya Grover for helpful comments.

2 Data at https://cs.nyu.edu/∼roweis/data/20news w100.mat