Stochastic Subset Selection for Efficient Training and Inference of Neural Networks
Authors: Bruno Andreis
Tuan Nguyen
Juho Lee
Eunho Yang
Sung Ju Hwang
Venue:  (2021)
Keywords: efficient deep learning
meta learning
efficient training
data compression
instance selection

Abstract: Current machine learning algorithms are designed to work with huge volumes of high dimensional data such as images. However, these algorithms are being increasingly deployed to resource constrained systems such as mobile devices and embedded systems. Even in cases where large computing infrastructure is available, the size of each data instance, as well as datasets, can provide a huge bottleneck in data transfer across communication channels. Also, there is a huge incentive both in energy and monetary terms in reducing both the computational and memory requirements of these algorithms. For non-parametric models that require to leverage the stored training data at the inference time, the increased cost in memory and computation could be even more problematic. In this work, we aim to reduce the volume of data these algorithms must process through an end-to-end two-stage neural subset selection model, where the first stage selects a set of candidate points using a conditionally independent Bernoulli mask followed by an iterative coreset selection via a conditional Categorical distribution. The subset selection model is trained by meta-learning with a distribution of sets. We validate our method on set reconstruction and classification tasks with feature selection as well as the selection of representative samples from a given dataset, on which our method outperforms relevant baselines. We also show in our experiments that our method enhances scalability of non-parametric models such as Neural Processes.

Under review as a conference paper at ICLR 2021
STOCHASTIC SUBSET SELECTION FOR EFFICIENT TRAINING AND INFERENCE OF NEURAL NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
Current machine learning algorithms are designed work with huge volumes of high dimensional data such as images. However, these algorithms are being increasingly deployed resource constrained systems such as mobile devices and embedded systems. Even in cases where large computing infrastructure is available, the size of each data instance, as well as datasets, can provide a huge bottleneck in data transfer across communication channels. Also, there is a huge incentive both in energy and monetary terms in reducing both the computational and memory requirements of these algorithms. For con-parametric models that require to leverage the stored training data at the inference time, the increased cost in memory and computation could be even more problematic. In this work, we aim to reduce the volume of data these algorithms must process through an end- to-end two-stage neural subset selection model, where the first stage selects set of candidate points using a conditionally independent Bernoulli mask followed by an iterative coreset selection via conditional Categorical distribution. The subset selection model is trained by meta- meta-learning with a distribution of sets. We validate our method on set reconstruction and classification tasks with feature selection as well as the selection of representative samples from a given dataset, which our method outperforms relevant baselines. We also show in our experiments that our method enhances scalability of aan--araattic models such as Neural Processes.
INTRODUCTION
The recent success of deep learning algorithms partly owes to the availability of huge volume of data (Deng et al., 2009; Krizhevsky et al., 2009; Liu et al., 2015), which enables training of very large deep neural networks. However, the high dimensionality of each data instance and the large size of datasets makes difficult, especially for resource-limited devices (Chan et al., 2018; Li et al., 2019; Bhatia et al., 2019), to store and transfer the dataset, or perform on-device learning with the data. This problem becomes more problematic for non non-parametric models such as Neural Processes (Hensel, 1973; Kim et al., 2019a) which require the training dataset to be stored for inference. Therefore, it is appealing to reduce the size of the dataset, both at the instance (Dovrat et al., 2019; Li et al., 10;;b;) and the dataset level, such that selects only small number of samples from the dataset, each of which contains only few selected input features (e.g. pixels). Then, we could use the selected subset for the reconstruction of the entire (either each instance or the entire dataset) or for prediction task, such as classification.
The simplest way to obtain such a subset is random sampling, but it is highly sub-optimal in that it treats all elements in the set equally. However, the pixels from each image and examples from each dataset will have varying degree of importance (Katharopoulos & Fleuret, 2018) to target task, whether it is reconstruction or prediction, and thus random sampling will generally incur large loss of accuracy for the target task. There exist some work on coreset construction (Huggins et al., 2016; Campbell & Broderick, 2018; 2019) which proposed to construct a small subset with the most important samples for Bayesian posterior inference. However, these methods cannot be applied straightforwardly to deep learning with an arbitrary target task. How can we then sample elements from the given set to construct a subset, such that it suffers from minimal accuracy loss on any target task? To this end, propose to learn sampler that learns sample the most important samples for a given task, by training jointly with the target task and additionally meta-learn a sampler over a distribution of datasets for instance selection in the classification task.

---

Under review as a conference paper at ICLR 2021
Faster communication
Faster communication
Real-lime Facial Recogritiinn
Faster training or inference
Faster Training inference
(a) Features Selection
(b) Instance Selection
Figure 1: Concept: Our Stochastic Subset Selection method is generic and can be applied to many type of sets. (a) Selecting a subset of features (pixels) out of an image. This reduces communication cost in data transfer between devices and allows faster training inference resource-comee condtreined devices due the reduced computational cost. (b) Selecting subset of instances out a dataset. This helps resource- cooorce-onstrraiee system to train faster, or to make non parameeettic inference more scalable.
Specifically, we learn the sampling rate for individual samples in two stages. First we learn a Bernoulli sampling rate for individual sample to efficiently screen out less important elements. Then, to select the most important elements out of this candidate set considering relative importance, we use Categorical distribution to model the conditional distribution of sampling each element given a set of selected elements. After learning the sampling probability for each stage, we could perform stochastic selection of a given set, with linear time complexity. Our Stochastic Subset Selection (SSS) is a general framework to sample elements from set, and it can be applied to both feature sampling and instance sampling. SSS can reduce the memory and computation cost required to process data while retaining performance downstream tasks.
Our model can benefit from a wide range of practical applications. For example, when sending an image an edge device with low computing power, instead of sending the entire image, could send a subset of pixels with their coordinates, which will reduce both communication and inference cost. Similarly, edge devices may need perform inference on huge amount of data that could be represented as a set (e.g. video, point clouds) in real-time, and our feature selection could be used to speed up the inference. Moreover, our model could also help with on-device learning on personal data (e.g. photos), as it can select out examples to train the model at a reduced cost. Finally, it can help with the scalability of non-paranae non-parametri models which requires storage of training examples, such as Neural Processes, to scale up to large-scale problems.
We validate our SSS model on multiple datasets for 1D function regression and 2D image recon- struction and classification for both feature selection and instance selection. The results show that our method is able to select samples with minimal decrease on the target task accuracy, largely outperforming random or existing sampling method. Our contribution in this work is threefold:
We propose a novel two-stage stochastic subset selection method that learns to sample a subset from a larger set with linear time complexity, with minimal loss of accuracy at the downstream task.
We propose framework that trains the subset selection model via meta-learning. such that it can generalize to unseen tasks.
We validate the efficacy and generality of our model on various datasets for feature selection from instance and instance selection from dataset, which it significantly outperforms relevant baselines.
2 RELATED WORK
Set encoding Permutation invariant networks Recently, extensive research efforts have been made in the area of set representation learning with the goal of obtaining order-invarian (or equivari- ant) and size-invariant representations. Many propose simple methods to obtain set representations by applying non-linear transformations to each element before a pooling layer (e.g. average pooling or max pooling) (Ravanbakhsh et al., 2016; Qi et al., 2017b: Zaheer et al., 2017; Sannai et al., 2019). However, these models are known to have limited expressive power and sometimes not capable of capturing high moments of distributions. Yet approaches such Stochastic Deep Network (De Bie et al., 2018) 8) and Set Transformer (Lee et al., 2018) consider the pairwise (or higher order) interactions among set elements and hence can capture more complex statistics of the distributions. These methods often result in higher performance in classification/eeression tasks; however, they have run time complexities of O(n2) or higher.
2

---

Under review as a conference paper at ICLR 2021
dg
dg
d2
dg
2
d2
dg
dg
d.
d
P1
03
Cat
d
d2
d3
dg
d
4
d2
Z
23
(3)
Cat
d2
dg
d4
d2
Set Encode
74
Cat
d
7
Zig
(a) Candidate Selection with Berouuli distribution
(b) Iterative Coreset Selection with stack of Categorial distributions
Figure 2: Overview. The subset D, is sampled by a stochastic 2-stage subset selection process. Dark shaded boxes correspond to selected elements while lightly shaded boxes correspond to non-selected elements in the given set. Also, Ber and Cat correspond to the Bernoulli and Categorical distributions respectively. (a) Candidate(in blue) Selection. (b) Core Subset(in green) Selection using continuous relaxations of discrete distributions.
Subset sampling There exist some works which have been proposed handle large sets. Dovrat et al. (2019) proposed to learn to sample a subset from a set by generating virtual points, then matching them back subset of the original set. However, such element generation and matching process is highly inefficient. Our method on the other hand only learns to select from the original elements and does not suffer from such overhead. Wang al. (2018) proposed distill the knowledge of large dataset to small number of artificial data instances. However, these artificial data instances are only for faster training and doesn't capture the statistics of the original set. Moreover, the instances are generated artificially and can differ from the original set making the method less applicable to other tasks. Also several works (Qi et al., 7a;c; Li al., 2018b; Eldar et al., 1997; Moenning farthes farthest point sampling, which selects l points from set by ensuring that the selected samples are far from each other on given metric space.
Dodgson, 2003) propose
Image Compression Due to the huge demand for image and video transfer over the internet, number of works have attempted to compress images with minimal distortion. These models (Toderici et al., 2017; Rippel & Bourdev, 2017; Mentzer et al., 2018; Li et al., 2018a) typically consist of pair of encoder and decoder, where the encoder will transfer the image into compact matrix to reduce the memory footprint and communication cost, while the decoder used to reconstruct the image back. These methods, while achieving huge successes in the image compression problem, are less flexible than ours. Firstly, our model can be applied to any type of sets (and instances represented sets), while the aforementioned models mainly work for images represented tensor form. Furthermore, our method can be applied both at the instance and dataset level.
Representation learning Our instance-sampling model is also related to the Variational Auto Encoder (VAE) (Kingma & Welling, 2013). However, while VAE learns a compact representation of a data point, our model learns a compact representation of set. Balın et al. (2019) learns a global feature selection model for reconstruction of the input data from selected features via unsupervised learning. Chen et al. (2018) learns instancewise feature selection with the goal of model interpretation by extracting subset of features most informative for given sample. Our method also falls this category.
Active Learning Active learning methods are aimed at selection of data points for labeling given small labelled set. This domain different from our method since active learning does not consider the label information but our method does utilize label information. Also, our motivation is quite different. We focus on efficiency in inference and training of non- pormmaraettric models by reducing the sizes of the inputs, be it pixels or instances and this greatly differs from the goal of active learning. Methods such as (Sener Savarese, 2017; Coleman et al., 2019; Wei et al., 2015) all tackle the data selection problem in the active learning setting.
3

---

Under review as a conference paper at ICLR 2021
7
Zi
Do
D.
D
D
CH
D.
Ul
UD
W
(b)
(c)
(a)
Figure 3: Graphical Models: (a) Feature Selection for Reconstruction. (b) Feature Selection for Prediction
Task. (c) Instance Selection model. (d) Instance Selection model for Classification
APPROACH
3.1 PRELIMINARIES
[di,...,dn) dn where individual d''s possibly represented as impul I; and target yi. D is the complete set and we assume that within D, there C D such that < n and that for an arbitrarily defined loss function E(., D) that we are interested in optimizing over the full set D, D2 can be used as proxy for such that E(., D) R E(.,, D.). what follows, we present a method that learns the conditional distribution p( D. DD) of the subset D2 via two stage selection procedure dubbed candidate selection and autoregressive subset selection. The overall objective is then to minimize the loss function with respect to the subset D.. Ep(D, (D) [e(,, D.)J. When the set D itself follows distribution of sets as in the meta-leat letanlea framework, then the objective becomes Ep [Pp [e(., Dz)]]. In essence, we seek to construct subset D2 that is optimally representative of the full set D w.n. E())
In this work, consider data of the type D =
exists subset D, {st, .5.....,8. .kk
3.2 STOCHASTIC SUBSET SELECTION
In order to select D.. we need to model the interactions among the elements of D and construct D, based on said interactions. However, when the cardinality [D] of the set D is large it's elements d, are high dimensional, modeling such pairwise interactions becomes computationally infeasible. As such, we first present the candidate selection procedure used construct a smaller set, De. without considering inter-sample dependencies. This is then followed by the autoregressive subset selection procedure used to construct D, from D. by modeling inter-sample dependencies. The complete model is depicted in Figure 2.
3.3
CANDIDATE SELECTION
We model the task of candidate selection random Bernoulli process where the logits of the Bernoulli function are conditioned on the set representation of the full set D and the individual elements d, D. For set D with cardinality we define Z {2 such that 2 E {0, 1) and z = 1 implies that d2 E D. and for each di, 2; is computed according to:
plzildi, D) Be((s;;
(1)
where r(D) is permutation-invarial function that compresses D into a single vector set representa- tion and p(d,, r(D)) computes the logits used to calculate the probability of d, belonging to D.. We implement both r r(D) and p(d,, r(D)) neural networks and specifically for r r(D), Sets (Zaheer et al., 2017). Since Ber is non-differentiable, we use the continuous relaxations of the Bernoulli distribution introduced in (Maddison et al., 2016; Jang al., 2016; Gal et al., 2017).
Deep
Specifically, to sample 2;, we execute the following computational routine:
o ( log
+ log
क. p(d;,r(D)), -~ Unif(i 1),
where a is the Sigmoid function, J is the temperature for the continuous relaxation and u is sampled from the uniform distribution. is set to 0.05 in all experiments. Given that pair-wise interactions
4

---

Under review as a conference paper at ICLR 2021
Algorithm Fixed Size Subset Selection
Input k(subset size), qff elements selected each iteration), {di, ....... da) (Full Set) (selected subset)
Output D: = {s1, ................................ S&
I: procedure STOCHASTIC SUBSET SELECTION(K, q, D)
2: (mi, m2,..., TT) <- (p(da, (D)). po.....,,ddd,,, p(d,,r(D)))
3:
2i ~ Ber(m.) for ,,...,
4:
De + {d, for 1: if
Candidate Selection
D. +
6: for iori= 1....,k/44 k/q do do
AutoRegressive Subset Selection
7:
D. c D. AUTOSELECT(Q, D.. D.)
Select elements
8:
return D.
9:
procedure AUTOSELECT(Q, Da, D.)
10:
C'= [Wi, W2, Wms +- De D.
11:
(PI, P2.. pm) (f(w1, De, Ds), fww2, De, Ds), f(um, De, Ds))
12:
(21. Pm & (P..P...... .P.. 23
13:
Q Select elements from C' with probability (p1,p2.... Pm))
14:
return Q
between elements are not considered in this stage, learning p(2z D) ensures that highly activating samples selected instead of random subset of the original set.
3.4 AUTOREGRESSIVE SUBSET SELECTION
The candidate selection stage introduce samples with redundant information in De since effort was made to compare the informativity of the elements. To alleviate this issue, we must first model the interactions between the elements of D, and construct D, based the relative importance of individual elements. To construct representative subset D, with |D.| = k, iterative steps are required and step the probability of an element in De \ (i-1) belonging to D, is computed according to:
f(d, De, Ed EXEDADE-N f(d', De, D.i-1) Vd D. AD!-"
p(s: dd D(:
(3)
where (i=11) the constructed subset at iteration 1 and is a positive function. The key to avoiding samples with redundant information in (( lies in the fact that for each element added to De. it's selection is conditioned on both D. and all elements in ( We further propose method that samples elements from Cat(pI.... Pm) in a single pass for efficient training. Specifically, instead of sampling 9 times from the categorical distribution, we sample the selection mask the element j being selected is
for element j from Ber(q ). In this routine, the probability q Pj which very close the original distribution. Algorithm details the entire procedure. The inference complexity depends heavily on the choice of the function f. If considers the pairwise interactions between all candidate elements and the selected elements, the inference complexity is q)) where n. d, & correspond to D). Del and D,, respectively. In our experiments, for the choice the function f. we utilize either a Set Transformer(Lee al., 2018) DeepSets(Zaheer et al., 2017) to model the pairwise interactions between the elements of given set.
3.5 CONSTRAINING THE SIZE OF De
For computational efficiency, we may desire to restrict the size of De to save computational cost when constructing D.. We adopt the idea of Information Bottleneck and constrain the distribution of 2 for Dr. Specifically,
(pcD, P.)] (KL[(
5

---

Under review as a conference paper at ICLR 2021
-08-
60-
Mean AUC
0.85
Image
-02-
0.50
505
16
30
100
150
200
300
400
800
S58
Numbor of Selected Points
Numbor of Salected Poxels
Number of Selected Pixals
(a)
(b)
(c)
Figure 4: Models' performance reconstruction and classification task (lighter-color areas present the
standard deviations). (a) ID Function reconstruction. (b) CelebA (reconstruction), (c) CelchA (classification)
2
2
LTS
RS
SSS
0
0
-2
-2
-2
-2
0
2
-2
0
2
-2
0
2
Figure 1D Function Reconstruction.
where r(Z) is a sparse prior. In our experiments, we set the parameter of the Bernoulli spares prior r(Z) to either 0.1 0.01 for different levels of sparsity and is set 1.0 or 0.001.
3.6 TASKS
We now present four tasks for which the described subset selection method is applied.
Set Reconstruction Given D3 and network pe p X, D.) parameterized by 0, the task to reconstruct D = (X,Y). The objective function for this task is given as:
Eytn) [pp log (Y D.)] (KLip(ZID))|<Z2]]
Minimizing this objective ensures that we learn a compact subset(D,) most representative of D and D, can then be used for other tasks. We implement (Y D,) as an Attentive Neural Process (ANP) (Kim et al., 2019b). An ANP takes as input a context( D, in this case) and predicts distribution of the elements in the original set D). It mimics the behaviour of Gaussian Process but with reduced inference complexity. The complete model is depicted in Figure 3a.Exxerimental results for this task can found in Section 4.
We can also opt to train the network to predict single target yn for the set D. For instance, the target could be the class of an (mage(classffiattio) or the statistics the set(regression problem). Here, pe (yD D&) is neural network that predicts the target yp- A set in this task, may be the features from a single example like an image and experimental results can be found in Section 4.2. The model for this task is depicted in Figure 3b. The objective function for this task given as:
Set Classification/Prediicii
(pxD)_ppD.I)) log pogpo(un)D.) + (KLL(Z(DD)
(6)
Dataset Distillation: Instance Selection For this task, we are given a dataset D (D....., D,) where each D; set of data points sampled from the entire dataset. Using CelebA as illustrative example, some D, may consist of [D. randomly sampled faces from the whole dataset. The goal is to construct D, for each D, E D. We describe model capable of taking as input D; E 'D to perform a task such the reconstruction of all elements in the given dataset.
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 1: CelebA Attributes Classification.
For single dataset D; E D. apply the subset construction method already described to D, that can be used to reconstruct all the elements in D, In essence, D, is distilled into a new dataset D. with k < [Da] elements. The task then is reconstruct the entire D, back conditioned only

---

Under review as a conference paper at ICLR 2021
#Instances
10
15
20
30
#Instances
2
FPS
6.50
4.51
3.07
2.75
2.71
2.29
FPS
0.432
0.501
0.598
0.636
Random
3.73
1.16
0.90
0.38
0.39
0.20
Random
0.444
0.525
0.618
0.663
SSS(ours)
2.53
1.02
0.59
0.33
0.24
0.17
SSS(ours)
0.475
0.545
0.625
0.664
Table 2: FID Score for varying Instance Selection
Table 3: Accuracy on minilmagenet
on D.. As first step, we represent D, as a unique representative vector e for each element in the dataset akin to the statistics network used in the Neural Statistician (Edwards & Storkey, 2016) model. Specifically, reconstruct an element d; D; given D,, C computed by applying a stochastic cross-attention
mechanism on D& where the stochasticity is supplied by query a which is computed using di- To obtain varying styles in the generated images, we additionally learn a latent variable w used to perturb e and both are combined to obtain a new element æ. The graphical model for this process depicted in Figure 3c. Additionally, to ensure that is properly learnt, we add an informativity loss by reconstructing from the generated samples from the given dataset. The objective for the model depicted in Figure 3c for single dataset D is
C(0,0,4)
ne (wn, (po ct] - (Llg Σ d,ED
where (-) are priors on their respective latent variables and qs (.)ss are implemented with neural networks. All priors are chosen to be Gaussian with zero mean and unit variance. This objective is combined with the informativity loss on all samples in D;. It is important to note that is computed using only D, for every element in D. In addition to Equation and the informativity loss, the model is optimized together with the subset selection model already described. When the model is fully optimized, it is applied to the instance selection task on the given dataset. In summary, the purpose of the generative model introduced is train the subset selection module for the instance selection task. Experimetal results for this task can be found in Section 4.3.
Dataset Distillation: Classification Finally in the dataset distillation task, we consider the problem of selecting prototypes to be used for few-shot classification. Here, we adopt Prototypical Net- works (Snell al., 2017) and apply the subset selection model the task of selecting representative prototypes from each class to be used for classifying new instances. By learning to select the proto- types, we can remove outliers that would otherwise change the class prediction boundaries in the classification task. The complete graphical model for this task is given in Figure 3d where again D, corresponds to the selected prototypes and r. and y. correspond to query and class label respectively. Experimental results for this task can be found in Section 4.3.
4 EXPERIMENTS
In this section, we present our experimental results. Model architectures and training hyper parameters are specified in the Appendix C.
4.1 FEATURE SELECTION EXPERIMENTS
Function Reconstruction Approximation Our first experiment is on 1D function reconstruction. Suppose that we have a function f [a, bj + R. We first construct set of data points of that function:
D = {(01, 1/1 = f(x1)), (I2, 3/2 = f(x2))..... (In, Vn. = f(x,), (xx)))} where (II, T2, In) are
uniformly distributed along the x-axis within the interval [a, bj. Now if have family of functions (f(), f(²).... ... ......(N)) (()), this will lead family of sets (D(1), ( ..... (NN).. We train model which consists of the subset selection model p DD) and a task network p(Y|X D.) (e.g. ANP), on this data set and report the reconstruction loss, which is the negative likelihoodd.
((()......DN)))
7

---

Under review as a conference paper at ICLR 2021
20
40
60
80
100Poens
20
80
80
100Po
20
40
60
80
100Pleels
Figure 6: CelebA reconstruction samples with varying number of pixels.
Bandon
Figure 7: Instance Selection Samples from dataset of size 200.
We compare SSS with Random Select(RS) (randomly selects subset of and uses an ANP to reconstruct the set), and Learning to Sample (LTS) (Dovrat et al., 2019) to sample k elements and uses an ANP reconstruct the set). Figure 4a shows the performance (reconstruction loss) of our models(SSS) and the baselines.
SSS out tutt-perfoms Random Select (RS), verifying that the subset selection model p( (DD,)D) can learn meaningful distribution for the selected elements. Our model also out-performs the Learning to Sample (LTS) baseline.
Through the visualization of the selected points in Figure 5, we can see that out model tends to pick out more points (presented as red dots) in the drifting parts of the curve, which is reasonable since these parts are harder to reconstruct. The other two baselines sometimes fails to do that, which leads to inaccurate reconstructions.
Image Reconstruction Given an image learn to select the original image. Here, x is 2-dimensional and y is 3-dimensional for RGB images. An ANP is then used to reconstruct the remaining pixels from set of context elements (selected subset in our case). We conduct this experiment on the CelebA dataset (Liu et al., 2018). Figure 4b shows that model significantly outperforms ANP with RS (as in the original ANP paper) and the LTT baseline. Figure 6 shows the reconstruction samples of our model which are visually better than the reconstruction of the baselines for the same number of pixels.
subset of pixels that best reconstructs
4.2 CLASSIFCATTIO//RRRESSST
In this subsection, we validate our model on the prediction task. The goal is to learn to select a subset for a target task such classification or regression. We again use the CelebA dataset, but this time the selected pixels are used to give predictions for 40 attributes of a celebrity's 's face (in multi-task learning setting). For our proposed model, only the selected pixels are used for prediction (other pixels' values are set to zeros). Table shows that using only 500 pixels (~1.3% of total pixels in an image), we can achieve a mean AUC of 0.9093 (99.3% of the accuracy obtained with the full image). Figure 4c shows the classification performance (in terms of mean AUC) versus the number of pixels selected. The AUC with selected pixels learned from our SSS is significantly higher than that of the random pixels baseline, showing the effectiveness of subset selection method. We also include another baseline, namely SSS(rec). This our stochastic subset selection model trained for reconstruction, but then later used for classification. Our model outperforms this variant, showing the effectiveness of training with the target task. Note that LTS cannot be applied to this experimental setup because during training, the generated virtual points cannot be converted back an image in matrix form (due to the virtual coordinate), thus we cannot train the LTS model with CNN-based classification on the target task.
Ablation Study Since our method is stochastic, the predictive distribution can be written as EPCDD D) (o((D)DD)) and we can use Monte Carlo sampling get the prediction in practice. However, throughout the experiment section, we only reported the result with one sampled subset, since gives the best reduction in memory and computational cost. This can be seen as MC sampling with one sample. We compare it against another variant: SSS(MC) with MC sampling (5 samples). It should be noted that by doing MC sampling with samples, the computational cost (inference) is increased by 5 times, and the memory requirement be increased by up to 5 times too. Table 1 shows that our model achieves comparable performance with that variant, thus justifying that it can achieve good performance for target tasks, while reducing memory and computation requirement.
8

---

Under review as a conference paper at ICLR 2021
4.3 DATASET DISTILLATION
Instance Selection We present results on the instance selection task applied to a whole dataset. In this task, we use the CelebA dataset since it has an imbalance both in terms of gender and race. A dataset is constructed by sampling 200 random images from the full dataset. In this experiment, we seek select only few(5-30) representative images from these generated datasets. On this task, our subset selection module is trained via the procedure detailed in Section 3.6 on instance selection. To evaluate the effectiveness of the SSS model, we evaluate the model in terms of the diversity in the selected subset using the Fréchet Inception Distance(FID Score) (Heusel et al., 2017) which measures the similarity and diversity between two datasets. We compare our model with the model that randomly samples instances from the full dataset. Additionally, we compare our method with the Farthest Point Sampling(FPS) algorithm which selects &o points from a given set by computing distances metric space between all elements and selecting those elements that are furthest from each other. FPS in general seeks to obtain a wide coverage over a given set and hence is a suitable baseline. The results of this experiment is presented in Table 4 where a lower FID score compared to FPS and Random Sampling. Additionally, given that the dataset highly imbalanced, FPS performs worst since by selecting the furthest elements in the given set it cannot capture the true representation of the whole dataset even when compared with Random Sampling. Also for small sample selection, our method outperforms FPS and Random Sampling significantly since our method is able to model the interactions within the full dataset and hence can select the most representative subset.
selection method achieves
Classification We use the minelmageNet dataset (Vinyals et al., 2016) and go from 20 shot classification task to one of 1,2,5 or 10 shot classification task. We again compare with Random Sampling and FPS and apply them together with SSS for the reduction in shot. The results for this experiment is shown in Table 3, where it can be observed that SSS can learn to select more representative prototypes compared to the other methods especially in the few-shot problems where the choice of prototypes matters more. All models were trained for 300 epochs and the best model was picked using a validation set.
CONCLUSION
In this paper, we have proposed a stochastic subset selection method to reduce the size of an arbitrary set while preserving performance on a target task. Our selection method utilizes a Bernoulli mask to perform candidate selection, and stack of Categorical distributions to iteratively select subset from the candidate set. As a result, the selection process does take the dependencies of the set's members into account. Hence, it can 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. We hope that this can facilitate the use of machine learning algorithm in resource-limite limited systems such as mobile and embedded devices.
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, 2019.
Abhinav Bhatia, Pradeep Varakantham, and Akshat Kumar. Resource constrained deep reinforcement learning. In Proceedings of the International Conference on Automated Planning and Scheduling. volume 29, pp. 610-620, 2019.
Trevor Campbell and Tamara Broderick. Bayesian coreset construction via greedy iterative geodesic ascent. arXiv preprint arXiv: 1802.01777 2018.
Trevor Campbell and Tamara Broderick. Automated scalable bayesian inference via hilbert coresets. The Journal of Machine Learning Research, 20(1):551-588 2019.
Michael Chan, Daniel Scarafoni, Ronald Duarte, Jason Thornton, and Luke Skelly. Learning network architectures of deep enns under resource constraints. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 1703-1710, 2018.

---

Under review as a conference paper at ICLR 2021
Jianbo Chen, Le Song, Martin J Wainwright, and Michael I Jordan. Learning explain: An cnoomatiion-thoreee perspective on model interpretation. arXiv preprint arXiv:1802.07844 2018.
Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis,, Percy Liang Jure Leskovec, and Matei Zaharia. Selection via proxy: Efficient data selection for deep learning. arXiv preprint arXi:: 18829, 2019.
Gwendoline De Bie, Gabriel Peyré, and Marco Cuturi. Stochastic deep networks. arXiv preprint arXiv:1811.0422) 2018.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scane scale hierarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255, leee, 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.
Harrison Edwards and Amos Storkey. Towards a neural statistician. arXiv preprint arXiv:1606.02185 2016.
Yuval Eldar, Michael Lindenbaum, Moshe Porat, and Yehoshua Y Zeevi. The farthest point strategy for progressive image sampling. IEEE Transactions on Image Processing, 609):1305-1315, 1997.
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in neural information processing systems, pp. 3181-358199 2017.
HERBERT Hensel. Neural processes thermoregulation. Physiological Reviews, 53(4):948-1017, 1973.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Berhaard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Jonathan Huggins, Trevor Camphell, and Tamara Broderick. Coresets for scalable bayesian logistic regression. In Advances in Neural Information Processing Systems, pp. 4080-4088, 2016.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparamcterizzation with gumbel-softmax arXiv preprint arXiv:1611 01 14444,10000 2016
Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. arXiv preprint arXiv:1803.00 00942, 2018
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint 19X1v:1901.05761, 2019a.
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.05761 2019b.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.611 14, 2013.
arXiv preprint
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar- 10 and eifar- 100 datasets. URI: http:://www. cs. toronto. edu/kriz/cifar. himl, 6, 2009.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam R Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer. arXiv preprint arXiv:1810.0087 2018.
Mengtian Li, Ersin Yumer, and Deva Ramanan. Budgeted training: Rethinking deep neural network training under resource constraints. arXiv preprint 190.v::905.04: 2019.
Mu Li, Wangmeng Zuo. Shuhang Gu, Debin Zhao, and David Zhang. Learning convolutional networks for content- content-weighted image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3214-3223, 2018a.
10

---

Under review as a conference paper at ICLR 2021
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 the wild. In Proceedingss of International Conference Computer Vision (ICCV), December 2015.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15:2018, 2018.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:I611.000 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. 4399-4402, 2018.
Carsten Moenning and Neil A Dodgson. Fast marching farthest point sampling. Technical report, University of Cambridge, Computer Laboratory, 2003.
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 point sets for 3d classification and segmentation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 652-660, 2017b.
Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas Guibas. Pointnet+++: Deep hierarchical feature learning point sets in metric space. In Advances in neural information processing systems, pp. 5099-5108, 2017c
Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Deep learning with sets and point clouds. arXiv preprint 161iv:1611... 1..500000,,,,.1 2016.
Oren Rippel and Lubomir Bourdev. Real-time adaptive image compression, In Proceedings of the 34th International Conference Machine Learning- Volume 70, pp. 2922-2930. JMLR. org, 2017.
Akiyoshi Sannai, Yuuki Takai, and Matthieu Cordonnier. Universal approximations of permutation invariant/equivariant functions by deep neural networks. arXiv preprint arXiv:
2019.
Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach. arXiv preprini arXiv: (oXXv:/708..04440 170890048, 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.
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.
Tongzhou Wang, Jun-Yan Zhu, Antonio Torralba, and Alexei A Efros. Dataset distillation. arXiv preprint arXiv: I. 10959, 2018.
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.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in neural information processing syslems, pp. 3391-3401, 2017.
11

---

Under review as a conference paper at ICLR 2021
A APPENDIX
This supplementary file organized as follows. We provide the full pseudo-code for the Greedy Training Algorithm. We then show some visualization our method for feature selection (in both ID function and CelebA dataset) and report the results with multiple runs of the instance selection experiment, as well its visualization. Qualitative results for Instance Selection as applied to the few-shot classification task are provided together with model specifications.
Organization
A.1 GREEDY TRAINING ALGORITHM
Algorithm shows our greedy training algorithm with stochastic gradient descent. The idea of the greedy training algorithm is to train the regresseee model to select the best next q elements from the candidate set to minimize the target loss on the selected samples. By doing this, we do not have run the auto-regressive model k/g time during training, thus reducing the computational cost.
Algorithm 2 Greedy Training Algorithm
Input
k(max subset size)
q( elements selected at each iteration)
pl D) (distribution of sets)
a (learning rate)
a target task with loss function ((.;))
Output
trained model with converged 0 and
I: 0, o + initialization
2: while not converged do
Sample minibatch with m sets ((2) ..... ((m) from p(D)
D! p(D! (D() for ....m .. D(2)
i random sample from (O...., k -0)
6: 16) ~ random i-element subset of _ for j 1...m
Q select a elementt subset from D 10) (with the auto-regressive model)
8:
-oo Em C(., (0) ) & aV E l(- 160 Qu1)
B INSTANCE SELECTION SAMPLES
In this section, we show more examples of our ID and CelebA experiments on how the models select the set elements for the target task.
B.0.1 ID FUNCTION RECONSTRUCTION
Figure 8 shows the reconstruction samples of our model on the 1D function dataset, which is objectively better than that of Learning to Sample (LTS) or Random Subset (RS). Since RS selects set elements randomly, it can leave out important part of the ID curve leading wrong reconstructions. LTS also selects insufficient amount of set elements some parts of the curves, resulting suboptimal reconstructions.
B. CELEBA
Figure 9 shows the selected pixels of our model for both the classification and reconstruction task. For the attribute classification task, the model tends select pixels mainly from the face, since the task is to classify characteristics of the person. For reconstruction, the selected pixels are more evenly distributed, since the background also contributes significantly to the reconstruction loss.
B.2 DATASET DISTILLATION: INSTANCE SELECTION
In Table 4, represent the full results for the Instance Selection model the CelebA dataset. For these experiments, we construct a set by randomly sampling 200 face images from the full dataset. To evaluate the model, create multiple such datasets and the baselines(Random Sampling
12

---

Under review as a conference paper at ICLR 2021
2
0
2
0
(a) SSS
(b) LTS
(c) RS
Figure 8: Reconstruction samples of ID functions with different selection methods
Original Image
Pixels Selected for Classification (SSS)
Pixals Selected for Reconstruction (SSS)
Pixels Selected for Reconstruction (LTS)
Pixels Selected (RS)
Figure 9: Selected pixels for different tasks CelebA.
and FPS) and SSS on the same datasets. The FID metric then computed on the instances and averaged on all the randomly constructed datasets. For FPS, we use the open-source implementation cluster. Further, we provide qualitative results on a single dataset in Figure 10 where show how model picks instances from the full set of 200 images face images.
in wowptuu...o/ruuttlllast.
Table 4: FID Score for varying Instance Selection
#lastances
15
20
30
FPS
6.5014
4,3502
4.5098 + 2.3809
3,0746 + 1.0979
2.7458
0.6201
2.7118
1.0110
2.2943 + 0.8010
Random
3,7309
1.1690
1.1575
0.6532
0.8970 0.4867
0.3843
0.2171
0.3877
0.1906
0.1980
0.1080
SSS
2.5307
1.3583
1.186 0.1982
0.5922 0.3181
0.3331
0.1169
0.2381
0.1153
0.1679
0.0807
B.3 DATASET DISTILLATION: CLASSIFICATION
In Figure 11 we provide visualizations for the instance selection problem as applied to the few-shot classification task. Here, we go from 20-shot a I-shot classification problem where the prototype is selected from the support using SSS.
C MODEL. SPECIFICATIONS
SSS consists of r(D), p(d;;r(D)) and [(d, Day D (1-1) We describe the models in this section.
13

---

Under review as a conference paper at ICLR 2021
Full Set
Candidate Set
Co reset
Figure 10: Visualization of set with 200 images for instance selection. The two stage selection method in SSS visualized as Candidate Set and SSS. A coreset of size is visualized.
14

---

Under review as a conference paper at ICLR 2021
Full Set
Cand Set
Coreset
Figure 11: Sample visualization of prototype selection for the min/Imagenet dataset the few-shot classifica- tion task. Each row represents a sef that corresponds to support from which prototype selected few-shot classification task.
For all experiments, r(D) is implemented DeepSets. This means that take the mean of all the samples in set to obtain the set representation.
p(d,,r( r()) implemented as a neural network with the following specifications: there are Linear layers each followed by ReLU activation. Also, all inputs are projected into feature space using Linear layers, each followed by ReLU activation.
In the set classification task, fd, De, D(--1) is implemented as Set Transformer network. All other experiments use DeepSets.
15


Reviews:
[{"review_id": "dHi4RWijfJZ", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Final recommendation**\nI support accepting this paper. While I am a bit disappointed the authors did not add in the paper the results they discussed during the rebuttal, I think the paper is interesting and clearer than at submission.\n\n**Summary**\nstochastic subset selection (SSS) is a method to learn to compress a set $D$ by selecting a subset $D_s$ such that the loss of a task performed on $D_s$ is as close as possible to the loss if the task had been performed on the original $D$. The paper show how this general method can be applied to several tasks. For example, $D$ can contain the pixels of an image and the task can be to reconstruct the original image from the subset. This corresponds to learning to compress. As another example, with the same $D$ but a label prediction task, SSS learns to dynamically select sparse features for a classification task. The other tasks considered are dataset distillation (compressing a data set) and selecting prototypes for few-shot classification. $D_s$ is constructed in two steps. A candidate set is first constructed by considering elements of $D$ independently and in a second step elements of $D_c$ are included in $D_s$ iteratively by considering other elements in $D_c$ and in $D_s$. Experiments are performed for one example of each task previously discussed and SSS is shown to outperform baselines. SSS is motivated by the need to reduce bandwidth, computation and memory footprint of deep learning in edge devices. \n\n---\n**Strong points**\n\n1) Code is provided (I did not try to run it).\n\n2) The proposed approach is applied to four different tasks. This convinced me of its applicability.\n\n3) Empirical experiments seem well conducted (with some caveat, see weak point 3) and results are good. \n\n4) The problems tackled are interesting.\n\n**Weak points**\n\n1) The paper is very dense, which makes it hard to follow. Related to that, I think that some parts of the paper are too concise to understand them well. The supplementary material partly compensates that.\n\n2) While the experiments show that the proposed approach works well, I think that some experiments should compare full data pipelines. As an example, one experiment shows that using SSS to dynamically select pixels of an image and using this sparse representation as input leads to good classification accuracy. The paper implicitely suggests that this reduces inference cost if classification is performed on the device and bandwidth if inference is performed on another device. However, the cost of the pixel selection by SSS is never analyzed. I think comparing the memory footprint, the number of operations, bandwidth used and accuracy for all steps of the following pipelines would give a much clearer picture of the interest of the approach for deep learning on mobile devices:\n    - classification is done directly on the device, using an optimized network.\n    - pixels are randomly selected and classification is performed on the device.\n    - pixels are randomly selected, transmitted and classification is performed on another device.\n    - pixels are selected by SSS, which involves computing multiple embeddings, and classification is performed on the device.\n    - pixels are selected by SSS, transmitted and classification is performed on another device.\n\n---\n**Recommendation**\nOverall, I vote for accepting. I am on the fence with this paper. The current results are good and I think the proposed approach has some potential and will be interesting for the community. On the other hand the paper is hard to read and the experiments, while good to show the strength of the approach, do not cover usefullness for the mobile use case. I think the former can be adressed in time and that the latter, while important, does not prevent the paper to be published.\n\n---\n**Details**\n\nI like that each task is first formally described and then evaluated in a specific case with the experiments. However I only understood most tasks when reading the experiment sections. I would suggest to either provide a short illustration of an actual instance of the task when describing the task formally, to directly provide the experimental result after describing the task or to direct the reader to the relevant experimental section(s) for illustration.\n\nThe paper states that model descriptions are available in appendix A, but I could not find them.\n\nFor the *Set Classification/Prediction* task, I would suggest clarifying that for this task a set typically contains the features of a single example, so $D=x$. I only understood this when reading the experimental section.\n\nThe *Dataset Distillation: Instance Selection* task was not clear to me. Here are a few points that puzzled me.\n- The text suggest the proposed approach is applied to each $D_i$, but figure 3(c) has no $D_i$ and $i$ seems to be used as an index of $x$ rather than $D$. Furthermore, in figure 3(c) a single D_s is constructed whereas I was expecting one per $D_i$. \n- The paper states that *It is important to note that $c$ is computed using only $D_i$ for every element in $D$*. However $c$ also depends on *a query $\\alpha$ which is computed using $d_i$*, which, if I understood correctly, is the element of $D$ and can be outside $D_i$. So it seems that $c$ is not computed using only $D_i$. \n\nFor the *Dataset Distillation: Classification* task, in figure 3d, the labels are not observed. How can the model learn? \n\n1) For this task, if the label is observed, how is this different from the first task?\n\n2) For the Image Reconstruction experiment, why is $x$ 2 dimensional? Are the inputs black and white images?\n\nI think that the image reconstruction experiment should include a comparison to compression algorithms. I agree with the paper that the proposed approach is applicable to many other types of data than images, but it would be very interesting to see how the approach compares to expert defined methods. Furthermore, implementation in relation to mobile devices is the motivation of this work and as far as I know classical compression algorithms are still used on mobile devices.\n\n---\n**Questions**\n\nQuestions 1 and 2 in *details* above.\n\nI would also be interested to hear authors' opinions on weak point 2 and a comparison to compression. \n\nFor classification, pixels not selected are set to 0. Wouldn't the inference cost be the same when using convolutional networks?\n\n---\n**Minor details**\n\nFor the *set reconstruction* task, would it also be interesting to learn a model on $D_s$ and use it to predict $Y$ from $X$?\n\n\n*typos*\nforo", "extra_scores": {}, "extra_text": {}}, {"review_id": "WIhRln6n2G", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe authors present a stochastic algorithm for selecting a subset of a large dataset, while trying to preserve statistics from the original dataset. The algorithm is a two-step process, first selecting a set of \"candidates\" based on individual features, then filtering to a final subset which may account for interactions between candidate items.\n\n### Strengths\n\nThe paper provides empirical experimental results which suggest the superiority of the proposed method against two baselines: Learning to Sample (LTS) and random selection (RS).\n\nVery few grammatical mistakes.\n\n### Concerns and Issues that Need Clarification\n\n1. A significant amount of active learning literature deals with subset selection. However, this paper fails to discuss or even cite any results from the active learning literature. It would be helpful for the authors to compare/contrast their proposed method with existing active learning results. For example, consider\n   - [*Active Learning for Convolutional Neural Networks: A Core-Set Approach* by Sener and Savarese (2018)](https://openreview.net/forum?id=H1aIuk-RW)\n   - [*Selection via Proxy: Efficient Data Selection for Deep Learning* by Coleman et al. (2019)](https://openreview.net/forum?id=HJg2b0VYDr)\n   - [*Submodularity in Data Subset Selection and Active Learning* by Wei et al. (2015)](http://proceedings.mlr.press/v37/wei15.html)\n   - [*On Statistical Bias In Active Learning: How and When to Fix It* (2020)](https://openreview.net/forum?id=JiYq3eqTKY) - this paper is a submission to ICLR 2020, so obviously there is no expectation that the authors would be familiar with this work. However, I bring this paper up because it discusses a similar topic of trying to select data points from a larger dataset that provide a reasonable estimate of a loss function on the \"population\" vs. a \"sample.\"\n\n2. The algorithm, as described, states that individual elements of a dataset $D$ are \"possibly represented\" as a pair $(x_i, y_i)$. What does \"possibly represented\" mean? Assuming that the dataset elements are given as a $(x_i, y_i)$ pair, the stochastic subset selection algorithm seems to describe instance selection. However, the paper also claims that the algorithm works for feature selection. What are $D$, $x_i$, and $y_i$ in the feature selection tasks? The paper makes an attempt at clarifying this by stating, \"Here, x is 2-dimensional and y is 3-dimensional for RGB images.\" But this still leaves the reader clueless - why is $x$ only 2-dimensional? What are those 2 dimensions?\n\n3. The end of the introduction claims that the proposed subset selection method \"learns to sample a subset from a larger set with linear time complexity,\" yet Section 3.4 claims that the inference complexity is $O(n) + O(k^2 d/q)$, which is nonlinear if either $k$ or $d$ is chosen to be dependent on $n$. Could the authors please clarify?\n\n4. The authors claim that their method uses meta-learning, but no mention of meta learning is made beyond the \"Preliminaries\" subsection of the \"Approach\" section. How is this method a example of meta-learning?\n\n5. The paper provides experimental evidence that the proposed method outperforms other subset selection baselines, such as Learning to Select (LTS) and random selection (RS). However, the paper provides little (if any) intuition or formal justification for why the proposed method performs better. The authors need to explain the motivation behind the proposed algorithm.\n\n6. The writing and figures lack clarity and are often difficult to follow. Some examples:\n   - The context for the mathematical expression in (4) is very unclear. Is (4) an objective function? If so, is this expression supposed to be minimized or maximized? And what is the random variable $Z$ supposed to represent? And why is there a period in the middle of the expression?\n   - Figure 2 is poorly explained and difficult to understand. I'm assuming that the blue vertical vector is supposed to represent $D_c$. Is that correct? And are the dark shaded boxes supposed to represent the \"candidates\" and the lightly shaded boxes supposed to represent the items that were not selected? As a reader, I should not have to guess what the colors mean. They should be explained clearly in the caption, or at the very least, in the paper text. Furthermore, it seems to me (again it's never made explicit) that Figure 2b is depicting the Autoregressive Subset Selection (\"AutoSelect\") routine. In this case, why is there a white row vector of $D = d_1, \\dotsc, d_n$? In Algorithm 1, it seems as though AutoSelect only takes as input $D_s$ and $D_c$, which are subsets of $D$ and not $D$ itself.\n\n### Original Rating and Confidence\n\n**Rating** - 4: Ok but not good enough - rejection\n\n**Confidence** - 4: The reviewer is confident but not absolutely certain that the evaluation is correct\n\n### Updated Review after Author Response\n\nThank you for the thoughtful responses to my questions, and apologies for the delay in updating my review. In general the responses have addressed my questions, and I have updated my rating accordingly. However, there are still several issues in my mind:\n\n1. **Connection to active learning** - The authors correctly point out that active learning does not assume access to labels, whereas the proposed method does. However, I believe the comparison against selection algorithms such as k-center greedy should be included in the paper for completeness sake, and not just left as a comment in this forum.\n\n2. **Dataset Distillation: Instance Selection** - Upon re-reading the paper carefully, I found this subsection in Section 3.6 to be confusing. Why are there multiple $D_i$? The section seems to only talk about one $D_i$.\n\n3. **Set vs. Tensor**: In the new \"representation learning\" subsection of Section 2, the paper stresses that the model learns a representation of a set, instead of an individual data point. However, many tasks (e.g., image reconstruction) do impose some sort of ordering. The paper lacks clarity about where permutation-invariance is used vs. not used.\n\n4. **Other places that need clarification**:\n  - Equation (2): presumably you threshold the output of the sigmoid function, right? Because in Algorithm 1, it seems like $z_i$ are supposed to be boolean-valued.\n  - Appendix: there are several missing closing parentheses", "extra_scores": {}, "extra_text": {}}, {"review_id": "J7iMmLQRbni", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work introduces a method to select instances from any set (stochastic subset selection, or SSS). The experiments demonstrate a diverse set of use-cases, including feature selection and core-set selection. The proposed approach is a two-stage method involving candidate selection (learning a function $\\rho$ to determine a Bernoulli probability for each input) and AutoRegressive subset selection (learning a function $f$ to generate probabilities for sampling elements from a reduced set); both stages use the Concrete distribution to ensure differentiability.\n\nThe paper includes a broad set of experiments, but I have some questions and concerns regarding the baselines and the proposed SSS method.\n\n1. Using the Concrete distribution enables differentiability, but Algorithm 1 seems to implement the candidate selection and AutoRegressive Subset selection steps in ways that eliminate differentiability. Line 4 assigns elements to $D_c$ based on $z_i$, but when is $z_i$ differentiated? Similarly, $Q$ is selected based on sampling using the $p_i$ parameters (line 13), but at what point are the samples differentiated after assigning elements to $Q$? The method apparently works so I'm sure there's a good answer, but it is not well explained.\n\n2. Some aspects of SSS seem arbitrary, and the authors could make a more compelling case for this particular approach with a more serious ablation study. A couple ideas: (i) Running SSS just using the first stage (candidate selection). (ii) Running SSS using the second stage directly (if it's not too slow). (iii) Removing aspects of the models, such as removing $r(D)$ as an input to $\\rho$, or $D_C$ as an input to $f$.\n\n3. As a further baseline for SSS, could the authors simply have implemented an algorithm that learns separate Bernoulli probabilities for each example (without the model $\\rho$)? This seems far simpler than SSS and it would be interesting to see how it compares.\n\n4. When attempting to constrain the size of $D_C$, why use a KL divergence penalty and not simply penalize each $z_i$'s probability, or the sampled $z_i$'s themselves? The latter could be viewed as a version of an $\\ell_0$ penalty on the subset size $|D_C|$, and both would remove the need to choose a hyperparameter for $r(Z)$.\n\n5. Obtaining $q$ examples by independently sampling from Bernoulli distributions with probabilities given by $q \\times p_j$ is odd. You can't guarantee that you get $q$ samples, and you can't even guarantee that $q \\times p_j$ is a valid probability. There must be a better solution to this problem; [1, 2] both use clever sampling tricks with the Concrete distribution, perhaps they suggest a better way.\n\n6. The paper is missing a couple citations and comparisons when it comes to feature selection using the Concrete distribution. Concrete Autoencoders (CAEs) [1] are a method for performing differentiable feature selection using a similar trick with the Concrete distribution; while CAEs are used for global feature selection (the same features for every instance), Learning to Explain (L2X) [2] is a method for instance-wise feature selection that also uses the Concrete distribution. Both should be mentioned, and L2X could be used a comparison in certain experiments.\n\nThe motivation around reducing communication overhead and allowing for faster training is interesting, but these claims could be more effectively grounded in the experiments performed in the paper. When accounting for the overhead involved in running SSS, is this method faster than conventional training, and does it require less communication overhead? The clearest application from these experiments is feature selection, but the comparisons with baselines are not extensive.  \n\nOverall, the method seems promising but I hope the authors can clarify several points and offer more competitive baselines.\n\n[1] Abid et al. \"Concrete Autoencoders for Differentiable Feature Selection and Reconstruction\" (2019)\n\n[2] Chen et al. \"Learning to Explain: An Information-Theoretic Perspective on Model Interpretation\" (2018)", "extra_scores": {}, "extra_text": {}}, {"review_id": "0DSP5-RGg21", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a stochastic subset selection method for reducing the storage / transmission cost of datasets. The proposes method minimizes the expected loss over selected datasets. The data selection algorithm consists a candidate selection stage and an autoregressive selection stage, parameterized with neural networks, and are trainable by gradient methods. The authors formulate and tested their approach on four tasks. The problem formulation and methodology are technically sound. The proposed method also seems to be more general than competing methods, such as coreset.\n\nHowever, I think the experimental results do not support the paper well.\n1. In Table 1 and Table 2, SSS doesn't seem to have big advantage over RS. From the selected images (e.g., Fig. 10 and 11), I also don't gain an intuition on what does SSS actually select, and how do the selected instances differ from random subsampling. \n2. The paper claims to address the dataset redundancy. From to my understanding, that is, removing similar instances. However, this claim is not explicitly supported by the experiments.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposed a two-stage method to select instances from a set, involving candidate selection (learning a function  to determine a Bernoulli probability for each input) and AutoRegressive subset selection (learning a function  to generate probabilities for sampling elements from a reduced set); both stages use the Concrete distribution to ensure differentiability. The experiments show the performance of the proposed method on several use-cases, including reconstruction of an image from a subset of its pixels, selecting sparse features for a classification task, and dataset distillation for few-shot classification. I read the paper and I agree with the reviewers that in its current format the paper is hard to follow. I strongly encourage the authors to add more discussion and intuition on the proposed method and extend the experiments with more baseline comparison and ablation studies in the revised version. 
Decision: Reject