Task Conditioned Stochastic Subsampling
Authors: Bruno Andreis
Seanie Lee
A. Tuan Nguyen
Juho Lee
Eunho Yang
Sung Ju Hwang
Venue: ICLR 2022 Submitted (2022)

Abstract: Deep Learning algorithms are designed to operate on huge volumes of high dimensional data such as images.  In order to reduce the volume of data these algorithms must process, we propose a set-based two-stage end-to-end neural subsampling model that is jointly optimized with an \textit{arbitrary} downstream task network such as a classifier. In the first stage, we efficiently subsample \textit{candidate elements} using conditionally independent Bernoulli random variables, followed by conditionally dependent autoregressive subsampling of the candidate elements using Categorical random variables in the second stage. We apply our method to feature and instance selection and show that our method outperforms the relevant baselines under very low subsampling rates on many tasks including image classification, image reconstruction, function reconstruction and few-shot classification.  Additionally, for nonparametric models such as Neural Processes that require to leverage whole training data at inference time, we show that our method enhances the scalability of these models. To ensure easy reproducibility, we provide source code in the \textbf{Supplementary Material}.


Under review as a conference paper at ICLR 2022
TASK CONDITIONED STOCHASTIC SUBSAMPLING
Anonymous authors
Paper under double-blind review
ABSTRACT
Deep Learning algorithms are designed to operate on huge volumes of high dimen- sional data such as images. In order to reduce the volume of data these algorithms must process, we propose set-based two-stage end-to-end neural subsampling model that is jointly optimized with an arbitrary downstream task network such as a classifier. In the first stage, we efficiently subsample candidate elements using conditionally independent Bernoulli random variables, followed by conditionally dependent autoregressive subsampling of the candidate elements using Categorical random variables in the second stage. We apply our method to feature and instance selection and show that our method outperforms the relevant baselines under very low subsampling rates on many tasks including image classification, image recon- struction, function reconstruction and few-shot classification. Additionally, for nonparametric models such as Neural Processes that require leverage whole training data at inference time, we show that our method enhances the scalability of these models. To ensure easy reproducibility, we provide source code in the Supplementary Material.
INTRODUCTION
Deep learning algorithms operate large volume of high-dimensional dense inputs such the pixels of an image (Deng et al., 2009; Krizhevsky et al., 2009; Liu et al., 2015). Training or evaluating a model with such data is computationally expensive and several works (Balın et al., 2019; Huijben et al., 2019; Yoon et al., 2018) have proposed subsampling techniques to subsample such dense inputs. Subsampling methods have the potential to drastically reduce the data acquisition effort and reduce the inference time of algorithms that operate on dense inputs. Additionally, subsampling techniques have found applications in medical research for the purpose interpretation (Ribeiro et al., 2016).
However, these methods have major drawback in that they are only applicable when each feature (such as a pixel) of the imput is 1-dimensional. This restriction is imposed by the way the subsampling methods are designed: the input (e.g. an image), is flattened to single vector and a model predicts a binary mask for each feature (e.g. pixels). This becomes problematic when we consider a channel image. Flattening out the image results in ambiguities as to which pixels to select since channels are treated independently. For instance, in order to perform subsampling on CIFARIO 10 images, INVASE (Yoon al., 2018) and DPS (Huijben al., 2019) convert the images into single channel images (e.g- grey-scaled images) before subsampling pixels. In more extreme cases such as subsampling (different from instance-wise feature selection) of training instances, each feature is itself an image (each possibly multi-channel) and hence those subsampling techniques cannot be used. Finally, as we show in our experiments, most of these methods fail under extremely low subsampling rates and their performance is similar to random sampling in this setting.
In order to tackle these limitation, we propose to consider each feature or instance as an element of a set. We formulate the subsampling problem as selecting a subset of features. or instances that minimizes the performance degradation of an arbitrary model on an arbitrary task such as image classification, regression or instance subsampling for tasks such as few-shot classification. As result, there are several advantages compared to the previous works. First, we can handle multi-dimensional feature, such as RGB pixel value, if we use set functions parameterized with expressive neural networks (Zaheer et al., 2017; 7; Lee et al., 2018). Second, a subsampling model with set function can process arbitrary number of elements. Thus the model is robust to wide range of subsampling rates at test time even when trained with fixed sampling rate, Lastly, the set-based formulation unifies the feature and instance subsampling tasks under single framework.

---

Under review as a conference paper at ICLR 2022
Interaction Network
Interaction Network
Interaction
Network
V
(a] Candidate Subsampling
(b) Autoregressive Subsampling
Figure 1: Concept: Two-stage stochastic subsampling process. (a) In the first stage, we screen out less important data points to construct the candidate subset. (b) In the second stage, we autoregressively select instances from the candidate subset.
However, it is prohibitively expensive to process all the set elements (e.g. all the pixels in image) with expressive set function such as Set Transformer (Lee et al., 2018) due to the quadratic computational cost with respect to the number of elements. Hence selection method. In the first stage, as shown in Fig. 1-(a). we learn the Bernoulli sampling rate for individual samples and efficiently screen out less important ones resulting subset which we call the candidate set. The second stage is more fine-gra fine-grained and designed to select a smaller subset from the candidate set by considering the relative importance of the samples the candidate set using conditionally dependent Categorical distribution through an iterative procedure as shown in Fig. 1-(b). Once optimized, the resulting subsampling model can perform stochastic subsampling of given input with linear time complexity. We call the resulting model Stochastic SubSampling (SSS) which a general subsampling framework and can be applied to both feature and instance selection.
propose efficient two-stage
We validate SSS multiple datasets and tasks such ID function regression, 2D image reconstruc- tion and classification for both feature and instance selection. The experimental results show that SSS is able to subsample with minimal degradation on the target task performance, largely outperforming the relevant baselines. We summarize our contribution as follows.
We reformulate the feature/instance subsampling problem by treating all the fea- tures/instances as members of a ser. This allows us apply set-based functions to the subsampling problem and extend its range of applicability.
We propose a set-based two-stage stochastic subsampling method that learns to efficiently subsample set with minimal performance degradation
target task.
We validate the efficacy and generality of our method on various datasets for feature selection in the input space (e.g. pixels) and instance selection from dataset and show that it significantly outperforms the relevant baselines.
2 RELATED WORK
Set Functions Recently, extensive research efforts have been made in the area set representation learning with the goal of obtaining order-invariant (or equivariant) and size-invariant representations. Many propose simple methods to obtain set representations by applying non-linear transformations to each element before a pooling layer (Ravanbakhsh al., 2016; Qi al., 2017 7b; Zaheer al., 2017; Sannai et al., 2019). However, these models have limited expressive power. Yet, approaches such as Set Transformer (Lee et al., 2018) consider the pairwise interactions among elements and hence can capture more complex statistics of set distributions. They often show better performance but require least (n?2) time complexity.
Feature Selection Recent interest in deep learning based subsampling methods has produced many works mostly applied feature selection. In Balın al. (2019), continuous approximation of the Concrete Distribution (Maddison et al., 2016) is used for global feature selection where a fixed set of features are sampled across an entire dataset. In Chen et al. (2018), instance-wise feature selection is used for interpretation of deep learning models applied medical data. Dovrat et al. (2019) propose learning to subsample by generating virtual points, then matching them back to the original input. Several works (Qi et al., 201 7a;e; Li et al., 2018b; Eldar et al., 1997; Moenning & Dodgson, 2003) also propose farthest point sampling, which selects & points from an input by ensuring that the selected samples are far from each other metric space. However work is most similar with
2

---

Under review as a conference paper at ICLR 2022
Dataset
Candidate Samples
Mulni(f, (+)
Muliil (x)
Add to Subset
Add to Subset
Set Ropresentation
D
Masking
Interaction Net g=ff.)
Ber((((d))
f(o)
Score Function
p()
Candidate Samples
D Subset
Candidate Samples
D
D
Subset
Mask Probability
[=
[a] Candidate selection with Bernoulli distribution
[b) Autoregresiive subsamples selection
Figure 2: Overview. Two-stages of Stochastic Subsampling. (a) Construction of candidate subset D. with sampling mask for each instance from Bernoulli distribution. (b) Iteratively sampling elements from the candidate set with multinomial distribution to construct the subset D..
the recent works of Yoon al. (2018) and Huijben al. (2019) which learn subsampling model conditioned on an given task. However these models have limitations both in terms of their range of applicability and poor performance under extremely low subsampling rates. Our method on the other hand is flexible, more general and applicable to a wide range of subsampling problems.
Image Compression Due to the huge demand for data transfer over the internet, some works attempt to compress images with minimal distortion. These models (Toderici et al., 2017; Rippel & Bourdev, 2017: Mentzer al., 2018; Li et al., 2018a) typically consist of an encoder and decoder, where the encoder transforms the image with compact matrix and the decoder reconstructs the image. These methods, while highly successful for the image compression problem, are less flexible than ours. Our model can be applied to any type of set structured data while the aforementioned models mainly work for images represented in tensor form. Furthermore, our method can be applied both at the instance and dataset level.
Active Learning Active learning aims to select data points for labeling given a small labeled set. This domain is different from since active learning does not consider the label information (for the selected data points) but our method does utilize label information. Also, our motivation is quite different. We focus on optimal subsampling conditioned on an arbitrary task and this greatly differs from the goal of active learning. Methods such as (Sener & Savarese, 2017; Coleman et al., 2019; Wei al., 2015) all tackle the data selection problem in the active learning setting.
APPROACH
3.1 PRELIMINARIES
We consider set D {di,............. (di...., ,......,d.]. an input with D p(-) for some unknown data distribution, where individual d, 's either represent pair of input 2; and target y or a feaeaurre such as the pixel value of an image. We assume there exists subset D. {si..... $ C D such that E(. D) E(. D.) for an arbitrarily defined loss function {(. D) that we are interested in optimizing over the full set with < n. what follows, we present method that learns the conditional distribution Pe( (D.(D) of the subset D, via two-stage subsampling method consisting of candidate selection and autoregressive subset selection as illustrated in Fig. 2. In general, we minimize the loss function (p.D)) D.)I] with respect to the parameters of the subsampling model.
3.2 STOCHASTIC SUBSAMPLING
To select D., we need to model the pairwise interactions among the elements of and then choose a few representative elements in D based on the relative sample importance computed using the interaction scores. However, when the cardinality of D is large or its elements, d,'s d''s, are high dimensional, modeling pairwise interactions becomes computationally infeasible since we need to compare every element in with all the other elements. This computational bottleneck motivates the first stage of SSS of which the goal is to construct smaller subset D.. which refer to as the candidate set, at a coarse level without considering pairwise interaction. We call the first stage candidate selection and the second stage, which is more fine-grained, autoregressive subset selection.
3

---

Under review as a conference paper at ICLR 2022
3.3 CANDIDATE SELECTION
We formulate the candidate selection problem as a random Bernoulli process where the parameters of the Beronulli distribution are conditioned on the set representation D and the individual elements d3 E D. Specifically, we first encode the dataset D to single representation De as follows:
De D 1 Σ g(d.),
(1)
where 9 is a neural network (Set Encoder in Fig. 2-(a)), which projects each element in indepen- dently to a lower dimension and = ID]. This encoding scheme is similar the one proposed in DeepSets (Zaheer et al., 2017) except that we do not perform message-passingg between the set elements. We then concatenate every d, with De. denoted as d,. This ensures that each element of D has global view of all the other elements in the set coarse level. For every d; E D, we can sample a mask 2 as follows:
σ (h(di))
Pe(zi D) = Ber(z;)p(d.)),
p(d.)
Σド= a (h(d))
where h is a neural network that outputs the logits for the probability that d; is in the candidate set De and (.) is the Sigmoid function, and Ber denotes the Bernoulli distribution. 24 is a binary random variable where 2; indicates that d; is element in De. We concatenate all 2/'s 's to obtain single vector 7 = (21,-.., ). Note that we feed the output of h into the Sigmoid function and normalize its output over all samples in D obtain valid probability distribution and induce sparsity in the selected candidate set De. Since sampling from the Bernoulli distribution is not differentiable, during training, we use the continuous relaxations of the Bernoulli distribution (Maddison et al., 2016; Jang al., 2016; Gal al., 2017) to sample z; Ber(p(d,) for de as shown in the black boxes in Fig. 2-(a). Although pairwise interactions are not considered in this stage, the ablation studies show that learning Pş (2 d,, D) leads to selecting highly informative samples compared to random selection.
Constraining the size of D. For computational efficiency, we want to restrict the size of D. to save computational cost when constructing D,, Hence introduce sparse Bernoulli prior r(Z) and minimize the KL. divergence along with the target downstream task loss E(, D.) w.r.t as follows:
EP(D) (ED((..D))((( D.)] + SKL (p
where (e( = II
ps (= D)
(3)
1=1
where 3 > is hyperparmeter used to control the sparsity level in 2.
3.4 AUTOREGRESSIVE SUBSET SELECTION
At this stage we have a set D. with m = DD < [D], which is small enough to perform fine-grained subset selection through pairwise modeling. To select a subset with elements from De, we requires k iterative steps. As shown in Fig. 2-(b), at time step I, we have the subset - constructed from
the previous iteration with ( 0 and D {((²) ..... s00 De (t-1) Assuming have a function ip o f (Interaction Net in Fig 2-(b)) for modeling pairwise interactions between the elements of an input set, autoregressively compute the interactions time step follows:
# (0)
(+
#(0)) = ( S(D(") D(--))))
where (-)) denotes the Sigmoid function and &p f is composition of two neural networks, one that computes interaction scores between elements in D and the other which outputs element-wise logits using the interaction scores. Further, #(0) is the vector of interaction scores for all elements in !!! at the current time step f. Given # > 0 for all =1,.. 1,...,m, my we can compute the probability
of element 8; (t) being selected from ( as:
(1)
p(s (0) (( D( = π
n
(t)
#
(t) (
m.
(5)
j=1
4

---

Under review as a conference paper at ICLR 2022
pixels
Subset pixels
Classification
Full-set of pixels
Subset pixels
Reconstruction
(c)
(d)
VS
VS
Dog
Full Dataset
Subset of Dataset
Few few-shot dataset
Subset selection Few-shot Learning
Figure 3: Target Tasks: (a) Feature selection for reconstruction.
(b) Feature selection for prediction.
(e)Instance selection representative data points. (d) Instance selection few-shot classification.
That is, we normalize #(t) over all the elements in D at time step to obtain a valid probability distribution. The key to avoiding redundant elements in D, lies in the fact that for each element added Da, its selection is conditioned on both the candidate set e and all the elements in the subset D as described in Eq. 4 & 5. For the choice the function f. we use MultiHead Attention Block (MAB) (Lee et al., 2018) which describe in detail in Appendix H.
With Eq. 5, we can sample an element s()) Cat(T] ..... Ti. from the candidate subset D and construct ! = - U {s(*)), where Cat is Categorical distribution. During training, it can be expensive to sample times from the Categorical distribution since involves computing Eq. k times. We remedy this by selecting 1 elements from _ at once, which reduces the number of iterations to k// for selecting elements. We may sample / elements from the multinomial distribution with probability =10) without replacement. However, this sampling procedure is differentiable, and hence it cannot be trained with backpropugation. Instead, we independently sample I elements from the continuous relaxation of Categorical distributions (Maddison al., 2016; Jang et al., 2016) using the same probabilities in Eq. 5 to approximate sampling from the multinomial distribution as shown Fig. 2-(b). Since we want to simulate sampling without replacement, we discard all elements sampled more than once. This sampling procedure guarantees that get at most / elements at each iteration. A similar sampling procedure is adopted in previous works (Balın et al., 2019; Chen et al., 2018). We present this training algorithm in Appendix C.
Inference Complexity of SSS. The inference complexity of SSS depends heavily on the choice of the function f. Using MAB the inference complexity of SSS is Onn) O(k2m/l) where n, m, & correspond to D) D and [D,| respectively.
3.5
TASKS
Set Classffication/Peeeittioo As shown in Fig. 3-(b), we train a neural network to predict a single target value yn for the subset D, of the given dataset D, where the set D is collection of the features from single instance such as pixels of an image. For instance, the target is either the class of an image for classification the statistics of the set for regression. Here, our goal is learning to select the most representative subset D, C D such that we can maximize the log likelihood log Pe (yD D.) with computational efficiency. In order to achieve this goal, we jointly train the SSS model and the neural network which predicts the target value yD for D, to minimize the negative log-likelihood, the loss function f(., D.) described in Eq. and enforce sparsity on Z, the selection masks for the follows:
candidate set, by minimizing the KL-divergence
Ep(D) [Epe(D. [o) log pe (yD D.)] + 3KL ()]I((Z))
(D))
(6)
We provide experimental results in Section 4.4 and corresponding graphical model in Appendix E.
Set Reconstruction Given a dataset D = {X, Y} consisting of 2d coordinates X = { and corresponding pixel values Y {y/ a 1 we want to reconstruct all the RGB pixel values y R for each coordinate x; e R² from the subset D. = (X,, Y.) with X, C X and Y, C Y as shown in Fig. 3-(a). We jointly train the SSS model and neural network predicting pixel values to minimize
5

---

Under review as a conference paper at ICLR 2022
the loss function w.r.t I and follows:
By(D)
(og, P.)] (2)))
We enforce sparsity on the subset D, by minimizing the KL-divergence between the mask probability Pc Pe(2| (Z)D) and sparse prior r r()). Moreover, minimizing the negative log likelihood, which corresponds E(., D&) in Eq. 3, ensures that the constructed D is the most representative for the downstream tasks. We implement P xX D.) as an Attentive Neural Process (ANP) (Kim et al., 2019). The ANP takes Ds as input and predicts a distribution the elements the original set D. It mimics the behaviour of Gaussian Process but with reduced inference complexity. We present experimental results for this task in Section 4.3 and a corresponding graphical model depiction in Appendix E.
given collections of datasets D fid [DO),... [(m)) with D n D()) = 0 for * j and D(i) ~ p(D), and the goal is to select the most representative subset D, with
Dataset Distillation: Instance Selection In this task,
D& < [D] for each dataset D = [di...., dn D, where d, is a data point uniformly sampled from the entire datasets D. Using CelebA dataset as an illustrative example, shown in Fig. 3, consists of n randomly sampled faces from the entire dataset and the task is to construct subset, D., most representative of D.
In order to learn to select the subset D, from each D E with unsupervised learning, we jointly train the SSS model and generative model such that the SSS model chooses the most representative subset so that the generative model can reconstruct all the images d; E 1) from the subset. Naively, can minimize the sum of negative bog-likelihood Ed
log po (o D2) for the loss function ((., D.) and KL divergence in Eq. 3. However, we find that the generative model outputs mean images for all d, To capture variations of different images, we introduce three latent variables as, Ca. and W, which both depend on d,. We provide graphical model illustration of this task in the Appendix E. Since is intractable to compute the log likelihood log pa D.) by marginalizing over all the latent variables, we derive the upper bound of the marginal likelihood using variational inference and plug the upper bound into the loss function E(,, D.) in Eq. 3 as follows:
Bp(D) Bp((D.10)
[ (1o,, log (o((:) c)] (Llgg(m:ldd)llv((u)) d.ED
Σ
(8)
+ [K]gg(a//d))lpp(a))) + KLlgo(ciDD,,ai)lloooor + (KL(P((11D))I112)
where (.(()) are priors on their respective latent variables, r(2 r(Z) is sparse prior over the mask for candidate set selection in SSS. po(-) is the decoder to reconstruct d, and all variational posteriors (-) are parameterized with neural networks. All priors are chosen as the standard normal distribution.
In summary, we jointly train both the SSS and generative model minimize the objective in Eq. w.r.t 0, ob and & for all D E 'D and leverage the optimized SSS to select few representative instances of the dataset, which results in distilled dataset. We report all the experimental results in Section 4.5
Dataset Distillation: Classification Finally for the dataset distillation task, we consider the problem of selecting prototypes for few-shot classification as shown in Fig. 3-(d). We adopt Prototypical Networks (Snell et al., 2017) and deploy the SSS model for selecting representative prototypes from the support set for each class. We minimize the objective in Eq. 3, where we use the distance loss induced by the metric space from Prototypical Networks for the target task loss É(- Dz). to jointly train the Prototypical Networks and SSS. Note that we use D., the subset of the support set, for computing loss and prediction. By learning select the prototypes, we can remove outliers that would otherwise change the class prediction boundaries in the classification task where we need to predict the label y, for an unseen instance I.. Experimental results for this task are in Section 4.5 and an accompanying graphical model description is provided in Appendix E.
EXPERIMENTS
Since SSS is two-stage method, we perform extensive Ablation in Appendix I.

---

Under review as a conference paper at ICLR 2022
Random
SSS
DPS
INVASE
Full Image
LTS
SSS(REC)
0.90
Ac0
3B (a)
100
16
10
Sobort Size
105
305
205
600
(b)
(d)
Figure 4: (a) MNIST Classification. (b) ID Function Reconstruction. (c) Image Reconstruction on CelebA. (d)
Attribute classification with fcature selection on CelebA
4.1 BASELINES
Feature Selection for Classification We compare SSS with the following models. 1) Random Selection: it randomly subsamples features. 2) DPS (Huijben et al., 2019): this model jointly optimizes the sampling parameters along with the parameters of task model. 3) INVASE (Yoon et al., 2018): this model uses actor-critic (Peters & Schaal, 2008) to optimize the parameters of the sampling network, as well as a model that makes use of the full set of pixels in each image on the MNIST classification task (Section 4.2).
In Section 4.4, we perform attribute classification on the CelebA dataset using the selected features from given image. We cannot apply DPS and INVASE to this experiment since each feature is tuple of 3-d RGB pixels since flattening these features and applying DPS and INVASE results in ambiguities to which features select and which to ignore. For instance, applying these methods to a 3-channel image can result in some channels being selected while others are zeroed out. However this leads the entire pixel being preserved and hence violate the subsampling objective.
Feature Selection for Reconstruction In Section 4.3 we compare SSS against the following: 1) Random Selection. 2) LTS (Dovrat et al., 2019): which is a model that learns to generate &s virtual elements which can be matched to elements in D and optimized for the downstream task. We use LTS for both the function reconstruction and the image reconstruction tasks. DPS and INVASE are not applicable for these experiments since each feature is multi-ddmenssonal. We also note that L.TS is not applicable for the classification tasks since the virtual points generated by LTS cannot be converted back into image form to serve as input to an image classifier.
Instance Selection In Section 4.5, we compare SSS with 1) k-Center-Greedy: this algorithm iteratively selects elements in D closest to set of centroids and 2) FPS: this algorithm iteratively selects the most distant elements to a randomly initialized D, and Random Selection on the instance selection tasks. Here also, DPS, INVASE, and LTS are all inapplicable.
Multiple Subsampling Rates For all the neural network based baselines (DPS, INVASE, LTS), we train a separate model for every subsampling rate. For instance, select 15, 20, 25, 30. 50 and 100 pixels from MNIST images in Section 4.2, we need to train 6 different models for DPS, INVASE and LTS each with the corresponding target subsampling rate. However for SSS, train single model and vary the sampling rate on each iteration. During evaluation, we use this single model for all the different sampling rates. We find that applying a similar training technique to the baselines result in drastic performance degradation. Thus the set formulation of SSS makes it generalize to varying subsampling rates at test time and offers train time efficiency.
4.2 FEATURE SELECTION FOR MNIST CLASSIFICATION
Given an MNIST image with 784 pixels. the task is to subsamlpe 15, 20, 25, 30, 50 and 100 pixels to be used as input to train and evaluate a classification model (CNN). Since images in the MNIST dataset have single channel, each feature is of dimension and thus we can train classifier with DPS or INVASE. We keep the pixels values for the selected pixels and set all the none selected pixels
(a)
(b)
(c)
(d)
(e)
Figure 5: (a) Full-image. (b) Random, (c)
DPS. (d) INVASE, (e) SSS. All models select 15 pixels of the original image.
7

---

Under review as a conference paper at ICLR 2022
Predicted
Ground Truth
Selected Subset
RANDOM
LTS
SSS
-2-2
2 -2
2 -2
2
Figure 6: Visualization of ID function reconstruction with Random, LTS, and SSS.
to zero. Note that we set subsampling rates to be much lower than the experimental setup from the papers (Yoon et al., 2018; Huijben et al., 2019).
As shown in Fig. 4a, SSS significantly outperforms all the baselines with large margin. It reaches 89% accuracy using only 15 pixels and shows better performance than the baselines with 100 pixels. Moreover under these low subsampling rates, the performance on the baselines are on par with random selection.
Lastly, we provide qualitative results in Fig. 5 where we visualize the selected 15 pixels by all the baselines and SSS. Again find that SSS selects representative pixels (Fig. 5-(e)) that the classifier can predict the correct label of the input image. However, all the baselines tend to select background black pixels, which are not informative for digit classification.
4.3 FEATURE SELECTION FOR REGRESSION
Function Reconstruction Suppose that we have function f [a, 6] R. We first construct set of data points with D {m11, y1 f(x=)), (In: Yn f(x,))), where (æ1,-....
3n are
uniformly sampled from the interval (a, b] and f is Gaussian process. We sample (v)² (O,, Kxx + - for = ......... where Kxx is a qquared-exponentiin kernel with the set of inputs = = {=r,,-.,, In. and of is variance for small likelihood noise. This leads to a collection of sets (D(1) ..... D(N)). We train our model which consists of the subset selection model Pe D) and a task network po (Y,,, Dz), which is Attentive Neural Process (ANP) (Kim et al., 2019), on this dataset and report the negative log- eigllikllihood (NLL).
(6) üid
yo
Fig. 4b shows the performance (NLL) of SSS compared to the baselines - Random Selection and LTS. As shown in Fig. 4b, SSS outperforms the baselines, verifying that the subset selection model pe [2 D) learns a meaningful distribution over subsets. We visualize a sample reconstructed function and the selected points by each models in Fig. 6. As shown in the rightmost figure (Fig, 6), SSS tends to pick out more elements (presented as red dots) in the drifting parts of the curve, which reasonable since those are harder to reconstruct than the others. However, the other baselines sometimes fails to do that, which leads to inaccurate reconstructions.
(D.
Image Reconstruction Given an image, we learn to select representative subset of pixels that best reconstructs the original image. Here, æ; is the 2d pixel coordinates and E IR3 is the RGB pixel value. We use an ANP to reconstruct the remaining pixels from the subset D. = {m Ui, ) constructed by each subsampling model. We conduct the experiment on the CelebA dataset (Liu et al., 2018). Fig, 4c shows that our model significantly outperforms Random Selection and LTS in terms of NLL. We provide qualitative examples for the competing methods in Appendix F.2.
4.4 FEATURE SELECTION FOR CLASSIFICATION
In this subsection, we validate our model on the image classification task illustrated in Fig. 3-(b). The goal is to select a subset of pixels of an image and predict the label of the chosen subset. We jointly train the subsampling models and classifiers on the CelebA dataset, where each classifier performs binary classification for 40 attributes of face. Note that only the selected pixels of image by the subsampling models are used for prediction and the other pixel values are set to zeros.
We report the mean AUC score all 40 attributes for varying sizes of D.. Fig. 4d shows that using only 500 pixels (~x1.3% of total pixels in an image), SSS achieves a mean AUC of 0.9093 (99.3% of the accuracy obtained with the full image). SSS achieves higher AUC score than Random Selection, showing the effectiveness our subset selection method. We also include another baseline, namely SSS-rec. This is the SSS model trained for image reconstruction, but then later used for classification
8

---

Under review as a conference paper at ICLR 2022
Table 1: FID Score with varying the number instances
#lostances
20
30
K-Greedy
8,8800 5,5857
4,4306 + 1,3313
4,2199 + 1,4214
3.7160 + 1.1314
3.2431 + 1,3881
2.7554 + 0.8554
FPS
6,5014
4.3502
4.5098 2.3809
3,0746 1.0979
2.7458
0.6201
277118
1.0410
2.2943
0.8010
Random
3.7309 1.1690
1.1575 0.6532
0,8970 0,4677
0.3843
0.3877
0.196
0.1980 0.1080
SSS
2.5307 + 1,3583
1.0166 +: 0.1982
0.5922 + 0.3181
0,3331 = 0.1169
0.2381 = 0.1153
0.1679 = 0.0807
without any finetuning. Our model also outperforms this variant, showing the effectiveness of training with the target task. Note that we cannot apply LTS to this experiment because during training, the generated virtual points with LTTS cannot be converted back to an image in matrix form due to the virtual coordinate, thus we cannot train the LTS model with CNN-based classification for this task.
Inference Efficiency in Nonparametric models In all the experiments where we used an ANP, we greatly improve the inference time complexity. By design, these models need to leverage the full training data at inference time. However by subsampling highly informative samples at low sampling rates (like in our experiments), these models must process very few elements at inference time. Similar gains can be obtained for models in the Neural Process family of models (Garnelo et al., 201 20;8;;
4.5 DATASET DISTILLATION
Instance Selection The goal is to select only a few representative images from the given dataset as described in Section 3.5 and Fig. 3. We split the CelebA dataset into m disjoint sets D {D(1) D(m) and jointly train SSS and the generative model minimize the objective in Eq. with D. After training, we discard the generative model and leverage the subsampling model to choose a few representative images from the full CelebA dataset.
We evaluate the selected subset with the Fréchet Inception Distance (FID) Heusel et al. (2017), which measures similarity and diversity between two datasets and compare SSS and Random Selection. We report the experimental results in Table where SSS achieves the lowest FID score for all selection sizes. Specifically, SSS outperforms all the baselines for selecting very few instances since SSS is able to model the interactions within the dataset and hence selects the most representative subset. Additionally, given that the dataset highly imbalanced, -.Cnter-Greedd and FPS perform worst since by selecting extreme or similar elements in the given set cannot capture the true representation of the whole dataset. We provide selected images by SSS from the full dataset in Appendix G.
Greedd, FPS
Classification In this task, we perform few-shot clas- sification with the minilmageNet dataset (Vinyals et al., 2016) where the models select 1, 2, or in- stances from the support set with size 20. As shown
Table 2: Accuracy on minilmageNet
Alnstances
2
FPS
0.402
0.501_0.002
Random
0.44410.003
0.025500005
0.618 0003
in Table 2, we compare SSS against Random Selec- tion, FPS, and k-Center-Grreedy SSS learns select more representative prototypes than the others especially for small D, where the choice of prototypes matters more. Notably, the K-Greedy method performs poorly for small subset sizes given that the model overfits to a few samples and does not generalize to unseen examples. We show samples of selected prototypes in Appendix G.I.
K-Greddy
0. 0.29020.006
0.413-0.005
0.570=00002
SSS
0.475=00006
0.545=0.011
0.625=0.006
CONCLUSION
In this paper, we have proposed stochastic subsampling method to reduce the size of an arbitrary dense input while preserving performance on a target task. Our subsampling method utilizes a Bernoulli mask to perform candidate selection, and stack of Categorical distributions iteratively select a representative subset from the candidate set. As a result, the selection process does take the dependencies of the input's 's elements into account. Hence, it select compact set that avoids samples with redundant information. By using the compact subset in place of the original set for a target task, we can save memory, communication and computational cost.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Muhammed Fatih Balm, Abubakar Abid, and James Zou. Concrete autoencoders: Differentiable feature selection and reconstruction. In International Conference on Machine Learning, pp. 444 453. PMLR, 2019.
Jianbo Chen, Le Song, Martin Wainwright, and Michael Jordan. Learning to explain: An information- theoretic perspective on model interpretation. In International Conference on Machine Learning, pp. 883-892. PMLR, 2018.
Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matci Zaharia. Selection via proxy: Efficient data selection for deep learning, arXiv preprint arXiv: 11299, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248- 255. Icve, 2009.
Oren Dovrat, Itai Lang, and Shai Avidan. Learning to sample. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2760-2769, 2019.
Yuval Eldar, Michael Lindenbaum, Moshe Porat, and Yehoshua Y Zeevi. The farthest point strategy for progressive image sampling. IEEE Transactions on Image Processing, 1997.
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. In Advances in neural information processing systems, pp. 3581 3581-3590, 2017.
Marta Garnll,, Dan Rosenbaum, Chris J Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Eslami. Conditional neural processes. arXiv preprint orXiv:1807.0166 2018a.
Marta Gamelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv: 2018b.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. In Advances in neural information processing systems, pp. 6636-6637, 2017.
Iris AM Huijben, Bastiaan S Veeling, and Ruud JG van Sloun. Deep probabilistic subsampling for task aas-tidetive compressed sensing. In International Conference on Learning Representations, 2019.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizatiion with gumbel-softmax. arXiv preprint arXiv:1611.0 111.611111111011 144, 2016.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv:1901.05766 2019.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar- 10 and cifar-100 datasets. URI: https://www: cs. toronto. edulkriz/cifar. html, 2009.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer. arXiv preprint I810 00825, 2018.
Mu Li, Wangmeng Zuo, Shuhang Gu, Debin Zhao, and David Zhang. Learning convolutional networks for content- weighttd image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3214-3223, 2018a.
Yangyan Li, Rui Bu, Mingchao Sun, Wei Wu, Xinhan Di, and Baoquan Chen. Pointenn: Convolution on x-transformed points. In Advances in neural information processing systems, pp. 820-830, 2018b.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
10

---

Under review as a conference paper at ICLR 2022
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba)
dataset. Retrieved August, 2018.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous
relaxation of discrete random variables. arXiv preprint arXiv:/61 2016.
Fabian Mentzer, Eirikur Agustsson, Michael Tschannen, Radu Timofte, and Luc Van Gool. Condi- tional probability models for deep image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4394-4402, 2018.
Carsten Moenning and Neil A Dodgson. Fast marching farthest point sampling. Technical report, University of Cambridge, Computer Laboratory, 2003.
Jan Peters and Stefan Schaal. Natural actor-critic, Nearocomputing, 7107-9):1180-1111 2008.
Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 652-660, 2017a.
Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 652-660, 2017b.
Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas J Guibas. Pointnet++:: Deep hierarchical feature learning on point sets in metric space. In Advances neural information processing systems, pp- 5099-5108, 2017c
Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Deep learning with sets and point clouds. arXiv preprint arXiv: 2016.
Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. why should i trust you?" explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1135 1144, 2016.
Oren Rippel and Lubomir Bourdev. Real-time adaptive image compression. In Proceedings of the 34th International Conference Machine Learning- Volume 70, pp. 2922-2933. JMLR. org, 2017.
Akiyoshi Sannai, Yuuki Takai, and Matthieu Cordonnier. Universal approximations of permutation invarian//equivariant functions by deep neural networks. arXiv preprint arXiv:1903.0// 1903.0199, 2019.
Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach. arXiv preprint arXiv: 170.00770,, 2017.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Advances in neural information processing systems, pp. 4077-4087, 2017.
George Toderici, Damien Vincent, Nick Johnston,
Sung Jin Hwang, David Minnen, Joel Shor, and Michele Covell. Full resolution image compression with recurrent neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5306-5314, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NIPS, 2017.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. In Advances in neural information processing systems, pp. 3630-3638, 2016.
Kai Wei, Rishabh Iyer, and Jeff Bilmes. Submodularity in data subset selection and active learning. In International Conference on Machine Learning, pp. 1954-1963, 2015.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. Invase: Instance-wise variable selection
using neural networks. In International Conference Learning Representations, 2018.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in neural information processing systems, pp. 3391-3401, 2017.

---

Under review as a conference paper at ICLR 2022
A ORGANIZATION
Organization This supplementary file is organized as follows. We provide the full pseudo-code for the Greedy Training Algorithm and Fixed-size Subset Selection Algorithm and illustrate the generative process of each task using graphical models. We then provide qualitative results for the 1D function reconstruction task, qualitative results the ablation studies and sample visualizations for the CelebA reconstruction task, Instance Selection and on Instance Selection as applied to few-shot classification on the minlmaaeeeet dataset. Finally we provide more details on specifications of models used in the experiments.
Code. We provide source code in the Supplementary Material.
Algorithm 1 Greedy Training Algorithm
Input
& (max subset size)
(mini-batch size)
I (# elements selected at each iteration)
pl D) (distribution of sets)
(learning rate)
((, D.) (loss function)
Output trained model with converged 0 and ø
1: Randomly initialize 0, E.
2: while not converged do
3: Sample n sets (11) ((m) from p(D)
4: Sample Z [2(0) ..... (1) (e(Z) for j = 1...m
5: Construct D (d()) E D()) 200 = 1} for = 1...m
6: Sample integer Unif(0,
16) + random i-element subset of D for j = 1...m m
7:
QU) + select l-element subset from (i) ( (with the auto-regressive model)
8:
9:
m O + 0 a VV.
a j=1 ((,10)) U QQ1)
10:
Ø + Φ a v 1=1
f(,, 16) u UQU))
11:
{ { V
j=1 ((,, 700 UQU)
12: end while
B SOCIETAL IMPACT AND LIMITATION
Since stochastic subsampling models learns choose the most task relevant elements of set, does not explicitly handle any biases induced by the set where certain attributes are highly correlated with genders or races, and may select highly biased samples.
C GREEDY TRAINING ALGORITHM
In order to reduce the computational cost at training time, stochastic gradient descent as described in Algorithm 1. It selects only the top elements from the candidate set to train the auto-regressive model by minimizing the target loss the selected samples and randomly selects the remaining do have to run the auto-regressive model k/l time during training, which significantly reduces the computational cost and shows reasonable performance.
greedy training algorithm with
elements from the candidate set. As a result, we
D FIXED-SIZE SUBSET SELECTION
At test time, we run the fixed size subset selection algorithm to choose the most task relevant elements from the set D. described in Algorithm 2. We do not the greedy training algorithm. Instead, autoregressively select elements from the candidate set D, as described in line 12 from Algorithm construct the representative subset D...
12

---

Under review as a conference paper at ICLR 2022
Algorithm 2 Fixed Size Subsampling. l is the required subset size. / is the number of elements to the full input set and D, selected subset after running SSS.
select at each iteration.
Input: k. 1, D
...
2: Output: D. = [s).....8k] 8k
3: procedure SSS(k,
4: De f +>: g(44). Concat(d,, D.)
5:
Ber(= pld)) for 1, .,...,7
6: D. - {d: D| for n
D. 4
8:
for 1= to k/l do
9: D. < D. U AUTOSELECTII, D., Dc)
10: end for
end procedure
12: procedure AUTOSELECT((, D D.))
13:
D [wi,.... Woyi De
D
14:
7(1)) o f(w, D 1)
15:
(mT (6)
..... ain (7t
# 1= (())
16:
Q +- Select elements
_ with (t)
17: return Q
18:
end procedure
E GRAPHICAL MODEL
In Figure 7, we illustrate the generative process using graphical models for each tasks (a) feature selection for set reconstruction, (b) feature selection for prediction (c) Instance selection for repre- sentative data points and (d) instance selection for few-shot classification. Only the shaded circles denote observable variable and the others are latent variables.
D
D2
OK
D
a
UD
W
(b)
(c)
Figure 7: Graphical Models: (a) Feature selection for reconstruction. (b) Feature Selection for prediction
task. (c) Instance selection for representative data points. (d) Instance selection for few-shot classification.
F INSTANCE SELECTION SAMPLES
In this section, show more examples of our 1D and CelebA experiments on how the models select the set elements for the target task.
F.1 ID FUNCTION RECONSTRUCTION
Figure 8 shows the reconstruction samples of our model on the ID function dataset, which is objectively better than that of Learning to Sample (LTS) or Random Subset (RS). Since RS selects the set elements randomly, can leave out important part of the ID curve leading wrong reconstructions. LTS also selects insufficient amount of set elements in some parts of the curves, resulting in suboptimal reconstructions.
F.2 CELEBA
Figure 9 shows samples of reconstructed images for varying selection number of pixels. Additionally in Figure 10, we show the selected pixels our model for both the classification and reconstruction task. For the attribute classification task, the model tends select pixels mainly from the face, since
13

---

Under review as a conference paper at ICLR 2022
Predicted
Ground Truth
Selected Subset
RANDOM
LTS
555
-2
-2
2
-2
-
-2
2
RANDOM
2
LTS
SS$
0
-2
2
-2
-2
2
2
RANDOM
LTS
SSS
-2
-2
RANDOM
LTS
SSS
-2
-2
2
2
-2
0
-2
0
2
RANDOM
2
LTS
SSS
-2
2
RANDOM
LTS
SSS
-2 -2
2
-2
-2 -2
2
-2 -
2
-2
-2 -
0
RANDOM
2
LTS
555
0
-2
-2_2
RANDOM
LTS
-2
2
-2
Figure 8: Visualization of 1D function reconstruction with three different subset selection models. Fach method selects 15 out 400 elements. As can he seen, SSS selects elements that result in better reconstructed functions.
14

---

Under review as a conference paper at ICLR 2022
555
LTS
Random
Original
Random
Original
555
LTS
Random
Original
Original
555
LTS
Random
Original
Random
Original
555
LTS
Random
Original
Figure 9: Visualization of reconstructed images for the CelebA dataset. Each model selects 40, 60, 80, and 100 pixels from a 218 x 178 image and reconstruct the full image using only the selected pixels.
Original Image
Pixels Selected Classification (SSS) Reconstruction (SSS)
Pixels Selected for
Pixels Selected for Reconstruction (LTS)
Pixels Selected (RS)
Figure 10: Zoom-In for best view. Selected pixels for different tasks on CelebA. As can be seen from selected pixels, SSS adaptively selects different pixels for both reconstruction and classification. Pixels for reconstruction more spread out to include the background since this contributes to the reconstruction loss. For classification. almost all the pixels focused on the face since most of the attributes can be found there.
15

---

Under review as a conference paper at ICLR 2022
the task is classify characteristics of the person. For reconstruction, the selected pixels are more evenly distributed, since the background also contributes significantly to the reconstruction loss.
G DATASET DISTILLATION: INSTANCE SELECTION
For the instance selection experiments, we construct a set by randomly sampling 200 face images from the full dataset. To evaluate the model, we create multiple such datasets and run the baselines baselines(Randon (Random Sampling. K-Center Greedy and FPS) and SSS on the same datasets. The FID metric is then computed on the instances and averaged on all the randomly constructed datasets. For FPS, we use the open- source implementation
ht pps: github. com/ rust mastyls// byytorch_ooo cluster. Further, we provide qualitative results a single dataset in Figure 11 where show how our model picks 5 instances from the full set of 200 images face images.
G.1 DATASET DISTILLATION: CLASSIFICATION
In Figure 12 provide visualizations for the instance selection problem applied the few-shot classification task. Here, we from 20-shot to a 1-shot classification problem where the prototype selected from the support using SSS. The selected subset then used in place of the support and used to classify new query instances.
H MODEL SPECIFICATIONS
In this section, we describe the main components our Stochastic Subset Selection models g(d), p(d) and 4 f(D(", ( D!
For all the experiments, we use g feedforward neural network with ReLU to project each instance d to lower dimension and average it to encode set representation De. following DeepSets (Zaheer et al., 2017).
We parameterize p(-) with a 3 layered feedforward neural network h followed by sigmoid function as follows as described in Equation (2) from Section 3.3. For p o f. we use conv-net to extract feature map for each instances in ( D! and feed it to set transformer (Lee et al., 2018) for set classification as follows:
S(D(0) 1(L-)))= = MA D!-
MAB(D! D( LayerNorm( H + +FF(HH)
(9)
H LayerNormmDD" Multihead(D ((t (D(-"),((-1)) D((-
where rFF is row-wise feedfoward layer which processes each instance independently and Multihead denotes Multihead Attention (Vaswani et al., 2017) with each slot of Multihead(.;- -. representing query, key, and value, respectively. For the other experiments, we use Deepsets for f. We use linear layer for to output logits for each element in D (t-1)
H.1 ATTENTION
We details on Attention and Multihead Attention here for completeness. For more thorough exposition, we refer the reader to Vaswani et al. (2017) and Lee et al. (2018).
attention module computes the following interactions using the dot product:
At(Q, K, V;w) w(QKT) )
(10)
where Q Raxdd the query vectors each with of dimension d,. K Ra.xd, and V 1R..dd, the keys and values respectively. Interactions modelled through QKT and w is activation function such as Softmax or Sigmoid.
Multihead attention projects Q, K, V to h different vectors each with a' d aM dimensions and computes different attention modules according to the following:
Multihead QQ, K, V; w)) concat(O),..... On)
(11)
16

---

Under review as a conference paper at ICLR 2022
Full Set
Candidate Set
Coreset
Figure II: Visualization of a set with 200 images for instance selection. The two stage selection method in SSS is visualized as Candidate Set and coreset. A subset of size 5 is visualized.
17

---

Under review as a conference paper at ICLR 2022
Full Set
Candidate Set
Coreset
Figure 12: Sample visualization of prototype selection for the mindlmagenet dataset on the few-shot classifica- tion task. Each row represents a ser that corresponds to the support from which prototype is selected for the few-shot classification task.
where
O, A KWjK, vwy: :wy)
(12)
The Multihead Attention module has learnable parameters À where À = (W WKK wy and j WKK E x4, wy E RRd and D E RR xd In all our experiments, we use Sigmoid as the activation function.
I ABLATION
We perform extensive ablation studies on the two stages of SSS using the function reconstruction task presented in Section 4.3. First we explore the contribution of the candidate selection and autoregressive subset selection stages. We then replace the candidate selection stage with random selection in the SSS model and compare it with the full SSS model.
Candidate Selection Only In Fig. 13, the model with only the candidate selection stage shows poor performance. This is because it is not always desirable to select only highly activating samples in the set without considering any dependencies among the others, which may leads to constructing candidate set with redundant elements.
Random Selection with Autoregressive Subset Selection To show how much the candidate selec- tion stage contributes to the performance of SSS, we replace it with random selection. As shown in Fig. 13, we find that while this model performs better than the model with only candidate selection, performs worse than SSS and the autoregressive subset selection stage used alone. This shows that filtering elements with candidate selection helps the model to select more informative instances from the input set in the autoregressive subset selection stage than random selection.
Regressee Subset Selection Only As shown in Fig. 13 we observe that the model with only autoreressive subset se- lection performs significantly better than the model with the candidate selection and the model with random selection fol- lowed by autoregressive subset selection. While this model performs well, it is not very practical due the high computa- tional cost when the size of the set becomes large. Additionally, processing lots of elements with the interaction model can lead
Stage 1 Only Random
Stage Only
SSS[Stage & 2)
Subset Size Figure 13: Ablation SSS.
18

---

Under review as a conference paper at ICLR 2022
Model
# Pixels
Storage
mAUC
Full Image
All 38804
114KB
0.9157
RS
500
5KB
0.8471
SSS(rec)
500
5KB
0.8921
SSS(MC)
500
5*5KB
0,9132
SSS(ours)
500
5KB
0.9093
Table 3: CelchA Attributes Classification.
to flattened attention scores over the elements, which makes it difficult to distinguish important elements from the others. The candidate selection stage plays important role of remedying this problem.
Generally, we find that SSS performs better than the variants considered here and has a better tradeoff between model performance and computational requirements.
Stochasticity of SSS Since our method is stochastic, the pre- dictive distribution is Epe(D, (D.(D.,D) [po(yD) D.)] for classification and we approximate with Monte Carlo sampling. However in all experiments, we only report the result with one sampled sub- set, since it gives the best tradeoff in memory and computation. We compare it against another variant: SSS-MC which use S samples of subset for MC sampling and obtains a mean AUC of 91.32%. Note SSM-MC increases the computational cost (inference) and memory requirement up to 5 times. Our model SSS with single subset sample achieves 90.93% (Fig. 4d) accuracy, which is comparable to SSS-MC. The result justi- fies that model SSS achieves good performance for target tasks with memory and computational efficiency. Storage com- parisons for the selected pixels with these variants of SSS are presented in Table 3. Additionally, in Figure 14 we show how uncertainty of subset sampling method decreases as we increases subset size. For each subset size, we draw 5 different subsets from the subsampling model trained on MNIST classification task described in Section 4.2 and average prediction of the sampled subsets, which requires forward pass of the model. After certain subset size, the model with single subset shows similar performance with the one with multiple draws of subsets.
1.0
Acuracy
0.9
55S
SSS(MC)
15
30
50
100
Subset Size
Figure 14: Accuracy with varying size of subset and the number of particles for MCMC.
Cost of Running SSS We report the cost of running SSS for pixel subsampling, which has the largest set size (38804 pixels). In this experiment, we select 500 pixels in total with / = 20, i.e., we select 20 pixels at once in the second stage described in Section 3.4. We measure the FLOPS and memory requirements for forward pass of SSS. We find that the computational cost of running SSS is 8.38 GMac (40% of the FLOPS for the full model which is 20 GMac) with 217.09k (compared to 958.85k for the full model) which shows that SSS is computationally cheap to run.
In the ablation studies, we investigated the contribution of the Candidate Selection stage (Stage 1) and the AugoRegres subset selection stage (Stage 2). Additionally, we replaced Stage I with random selection in SSS and compare the performance of these three variants with SSS. In Figure 15, we provide qualitative results on the function reconstruction tasks for these 3 variants as well as SSS. As can be seen from the figure, using the Candidate Selection stage only select elements generally focuses on the difficult sections of the curves ignoring the other parts which are relatively casier to reconstruct and this results in poorer overall reconstructions. The Autoregressive Subset Selection stage on the other hand selects elements that reconstruct the functions well. The downside to this model is that it requires a large amount of computation since we have to execute the interactions model for all elements in the input set and this renders this model impractical. On the other hand, the model that replaces the Candidate Selection stage in SSS with Random Selection performs better than using only the Candidate Selection stage only but results in function reconstruction that -ar compared to the Autoregressive Stage only and SSS. This is because using random selection in place of the first stage can ignore elements from certain parts of the function when randomly sample the candidate
19

---

Under review as a conference paper at ICLR 2022
Predicted
Ground Truth
Selected Subset
CAND
RANAUTO
AUTO
555
-2
-2
2
CAyD
RANAUCO
AU10
-2
-2
-2
-2
-2
-2-22
-2
CAND
RANAUTO
AVTO
S$5
-2
-2
-2
-2
CAND
RANAUTO
AUTO
555
-2 -2
-2 -2
-2-2
-2
2
2 22
2
-2-2
22
CAND
RANAUTO
AUTO
SSS
-2
-2
22222
-2
-2-2
-2-22
-2-2
CAND
RANAUTO
AUTO
555
Figure 15: Ablation: Visualization of ID function reconstruction. In the ablation studies, we compare the first stage (CAND) and the second stage (AUTO) with SSS. Additionally, we replace the Candidate Selection stage (Stage 1) in SSS with random selection (RANAUTO) and compare the performance these models. As can be seen from the visualized reconstructed outputs, the combination of the Candidate Seletion stage with the Autoregressive stage results in the best subset selection for the reconstruction task.
set and hence these parts cannot be reconstructed since the second stage cannot select elements from those regions. This behaviour can be seen in some of the functions in the second column of Figure 15, In conclusion, SSS offers a good trade-off between computation and performance and results in better reconstructions since both the Candidate Selection stage and AutoRegressive Subset Selection stages work together to select the most relevant elements for the task.
20

---

Under review as a conference paper at ICLR 2022
LTS
Invate
1..0
(a)
Figure 16: (a) CIFARIO Reconstruction. (b) CIFARIO Classification.
DISCUSSION OF LTS PERFORMANCE IN FIGURE 4C
For the image reconstruction task in Figure 4c, the LTS model generates virtual points for each pixel (RGB values and two coordinate values). In our experiments, we find that the generated coordinates values by LTS are imprecise compared to sampling directly from the original pixel space (as done in SSS and Random subsampling). This makes matching the virtual points with the original pixel values extremely difficult. Such inaccurate coordinate values result in poor performance as the subsampling rate increases even compared to random subsampling as depicted in Figure 4c. We observe similar pattern in the CIFARIO reconstruction task presented in Section K and Figure 16a.
The need to predict the two coordinate values for the virtual points is necessity only because the ANP model used for the reconstruction task requires the context points which consist of RGB values and their corresponding coordinates in the image space. Interestingly in the function reconstruction task in Section 4.3, the point matching stage in LTS is fairly easy and hence the LTS model shows better performance than random subsampling.
K EXPERIMENTS ON CIFAR10 DATASET
We provide further experimental results the CIFAR10 dataset in which we subsample pixels for both image reconstruction and image classification. In order to compare with both Huijben et al. (2019) and Yoon et al. (2018), we convert all the images to grayscale, following Huijben al. (2019) (see Section 4.3).
The experimental result of the CIFARIO reconstruction task is presented in Figure 16a, where for the same subsampling rates, SSS outperforms the competing baselines (DPS, INVASE, LTS and Random Sampling) terms of the negative log-likelilne likelihood. For this task, we use ANP model as we have done in CelebA and function reconstruction tasks in Section 4.3. For this task, we do not compare with INVASE since INVASE requires two copies of the reconstruction model and requires more GPU memory.
In Figure 16b, we present the results for the CIFAR10 classification task. Again we observe that for the same subsampling rate, SSS performs better than DPS, INVASE and Random Sampling and the performance of the same classification model trained on the full input image is 0.70 0.02. Note that we cannot compare with LTS for the same reasons given in Section 4.1.
21


Reviews:
[{"review_id": "gTVO4sQSGVd", "reviewer": "Reviewer_Lpik", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "The main contribution of this paper is to propose a 2-stage sampling method. The core method of each stage does not have much novelty. The main weakness of this paper is that the core details of the proposed method are missing in the main text, which makes it very difficult for readers to assess its true value.\n\nDetails:\n1) p. 3, last paragraph of Sec. 3: It's not accurate that \"active learning does not consider the label information.\" Action learning selects examples to be labeled, but which examples to select is based on the error/loss of the existing model, which in turn depends on the label information.\n2) p. 3, Sec. 3.1: The authors state \"we minimize the loss function...\" Minimizing the loss function is to identify the subset $D_S$ or train the model for the downstream task? If it's the former, the previous description in Sec. 1 indicates that the authors would like to minimize the \"performance degradation\", which is the difference in losses from the original dataset $D$ and the subset $D_S$. If it's the latter, then why there're expectations here? Do you minimize over multiple draws or multiple selected subsets?\n3) p. 4, Sec. 3.3: How do you train the neural network $h$?  It's not clear at all.  Based on Eq. (2), we have $E(\\sum_{i=1}^{n} z_i) = 1$, which mean the effective sample size is 1? Also, you mentioned \"sparse Bernoulli prior $r(Z)$\". What does a sparse Bernoulli prior look like? Could you give an example?\n4) p. 4, Sec. 3.4: For the neural networks $\\varphi$ and $f$, how are they trained?\n5) p. 6, Sec. 3.5: \"Representativeness\" is defined as allowing a generative model to reconstruct the other examples in the dataset. How do you use a generative model to reconstruct the remaining examples? There's no such guarantee by simply sampling from a VAE.\n", "summary_of_the_paper": "This paper proposes a 2-stage sampling method that can find its use in either feature selection or instance (training examples) subsampling. At a high level, the authors frame the problem as a subset selection problem with the goal to minimize the performance difference for a downstream task between using the original set and the downsampled set. The proposed method consists of two stages: 1) select a subset (called candidate set) using Bernoulli variables as the binary mask, and the Bernoulli distribution's parameter depends on the corresponding example and the original set, 2) select elements from the candidate set iteratively to form the final selected subset, and the selection here is based on a softmax whose parameters depend on the interactions between the remaining candidate set and the selected subset so far. To support the merits of the proposed method, the authors provide experiment results on feature selection and dataset distillation tasks.", "summary_of_the_review": "The authors try to pack lots of information into this paper, but the key details to support the usefulness of the proposed method are not fully explained. Without such details, it's hard for readers to appreciate its intellectual merits, and it's equally hard for readers to implement it in real applications. Therefore, this paper in its current form, doesn't seem to meet the bar of ICLR, and requires major improvements.", "strength_and_weaknesses": "", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "M1NdtxxG-rg", "reviewer": "Reviewer_7mpS", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "Strengths:\n1. The authors validate their approach on many different classification tasks.\n\n1. The idea is simple and feasible.\n\n1. The method offers performance improvements, although I'm skeptical about the empirical significance of the experiments.\n\nWeaknesses:\n1. The experiments are only on celebA faces and MNIST digits, both unchallenging datasets.\n\n1. I do not think the numerical comparisons in Fig 4 are fair. The authors technically get access to the entire dataset and then choose what to show to the classifier, based on evaluations on different subsets. The performance gains over competing methods tell me that the baselines are most likely unfairly weaker.\n\n1. I have concerns about the technical novelty of this work. It essentially boils down to sampling from a softmax distribution over features, then putting a sparsity prior on the distribution to make sure few features and data points are preserved.", "summary_of_the_paper": "This paper proposes a data summarization technique that can be used for arbitrary downstream classification tasks. The authors show that their methods can be used to reduce the number of features required by a model while preserving accuracy. The authors additionally show that similar ideas can be used to extend this to selecting a subset of images in dataset that serves as a training dataset that provides good accuracy.\n\nThe algorithm proceeds by sampling features according to some softmax distribution defined by the classifier on each feature. The algorithm then learns to sample such that the accuracy of the classifier does not fall significantly when trained on a small subset of features. (The small subset of features are selected by adding a regularizer that rewards sparse selection of features)", "summary_of_the_review": "Overall I am unconvinced by the technical novelty of the work and its implications. I do not necessarily buy the argument that existing feature selection techniques fail on RGB images. Additionally, I do not see the point of the experiments where a subset of the data is selected to train the classifier without sacrificing the accuracy -- each training technically requires the full dataset, even though the classifier gets to see only a subset, and training a new classifier would require the entire dataset to select a different sample.", "strength_and_weaknesses": "", "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": "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": "ISihlWV65jf", "reviewer": "Reviewer_oSsn", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "The major weakness of the paper is that it lacks theoretical justification for the proposed method. For example, it is unclear if the learned decision function (e.g., regression function and classification boundary) is consistent or not. Uncertainty of the learned decision function is unclear, either. ", "summary_of_the_paper": "This paper presents an interesting method for selecting subsamples or subfeatures used for deep learning. More precisely, the paper proposes to regularize the task loss by a Kullback-Leibler divergence between the selected samples/features and a sparse Bernoulli prior, and employs a neural network to learn the weights used for sample/feature selection.  ", "summary_of_the_review": "This paper present an interesting method and promising results for subsample selection used in deep learning, while lacking theoretical justification for the proposed method.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "D02N372MLnL", "reviewer": "Reviewer_zeQG", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Strength:\n- The paper writing is good, and the presentation is very clear.\n- The idea of SSS is interesting, and its set-based nature allows it generalizes to various tasks.\n- Extensive experiments have been conducted to support the claim that SSS is advantageous over other baselines.\n\nWeaknesses:\n-  In Figure 4 (c): for the task of image reconstruction on CelebA, I observe that LTS performs even worse than random subsampling, which is counter intuitive and is inconsistent with other tests, e.g., Fig 4 (b). Discussions need to be added here.\n- In Section 3.5, I would suggest re-order the paragraphs so that they match the order of the illustrations in Figure 3 to increase the readability. Now it starts with Fig 3 (b), then (a), (c), (d). Why not (a), (b), (c) and (d)? Similar things could be improved in Section 4 when discussing each tasks and the experiments.\n\n-Typo:\n- In the third to last line of the paragraph after Eq. (2), there is a right parenthesis missed for \"Ber(\\rho(\\bar d_i)\"\n\n", "summary_of_the_paper": "The paper proposes a Stochastic SubSampling method (SSS) that is helpful in reducing the size of input dataset while preserving reasonable performance on a target task. SSS is two-stage and set-based, which can be jointly optimized with arbitrary downstream tasks. Extensive experiments verify that SSS do perform well in the sense of efficiency and generalization a variety of tasks and datasets.", "summary_of_the_review": "The paper writing is clear, and the proposed method is clearly advantageous over other baselines in terms of generalization and efficiency. Extensive experiments have been conducted to support the claim. I would vote for an acceptance of the paper.", "strength_and_weaknesses": "", "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": "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": {}}]

Author rebuttal:

Meta-review: The paper proposes a set-based neural subsampling model that selects both features and instances using a two-stage model. The motivation is to allow for scaling to large datasets by first subsampling using an initial stage that does not model second-order interactions (which would require work quadratic in the dataset size to model), and then following up with a second more sophisticated pass that includes second-order interactions. Its results show an empirical improvement over previous methods in the case of very small subsample sizes.

The responses from the reviewers in discussion were varied—and often off-base for all reviewers—and as a result I took an even more deep look at this paper than usual. I think the variance in reviewer response is a symptom of the fact that the paper is somewhat confusingly written, and sometimes has parts that give the opposite impression to what the authors intend. For example, the author response says "we emphasize again that our method is a subsampling method. This is very different from core-set selection methods" but then Figures 11 and 12 explicitly label the subset produced by the algorithm a "coreset." There is general confusion as to what exactly is being subsampled (features or instances) and even what datasets were being evaluated on—it's not that the information is not there, but rather that it's easy to miss while reading the paper. We can see this happening where most of the reviewers were misunderstanding or making factual errors about the paper, and I can see how this happened by reading the paper.

The reviewers also shared some concerns about the baselines, and indeed some things about the baseline comparisons are confusing: for example, Figure 4 seems to report DPS having below 70% accuracy on MNIST while subsampling to a size of 25, but the original DPS paper (Huijben et al, ICLR 2020) in their Figure 2 reports a percent error of 6.6% (at Pixels removed: 96.8%, which I believe corresponds to keeping 25 pixels as 28*28*(1-0.968) = 25). This does seem to back up the reviewer's speculation that "the baselines are most likely unfairly weaker." The presentation of the results should make this sort of thing more clear (if the setup isn't the same as DPS's MNIST setup, how does it differ? if the setup is the same, as seems to be the case the way the paper is presently written, why are the result accuracies so different from what is published in the DPS paper?).

The reviewer comment about theory is not one I count against the submission. Although it is certainly true that this paper would be greatly strengthened by some theoretical backing, it is also part of a line of work that eschews theory—so we cannot reasonably disqualify it for doing so.

To sum: although the technical contributions of this paper do seem to be significant, I expect that if the paper is published as presently written, it will confuse ICLR readers just as it has confused our reviewers. This leads me to lean against accepting this paper at this time.
Decision: Reject