{'fraction_non_alphanumeric': 0.029943350418127866, 'fraction_numerical': 0.007903965470731049, 'mean_word_length': 4.532985074626866, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 2, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Stochastic Subset Selection for Efficient Training and Inference of Neural Networks

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 be a 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 nonparametric models that require to leverage the stored training data at 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. We first efficiently obtain a subset of candidate elements by sampling a mask from a conditionally independent Bernoulli distribution, and then autoregressivley construct a subset consisting of the most task relevant elements via sampling the elements from a conditional Categorical distribution. 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 nonparametric models such as Neural Processes.


Introduction
The recent success of deep learning algorithms partly owes to the availability of huge volumes of data such as ImageNet, CIFAR10/100, CelebA [10,23,28] to name a few, which enables training of very large deep neural networks. However, the high dimensionality of each data instance and the large size of datasets in general makes it difficult, especially for resource constrained/limited devices [5,26,2], to store and transfer the dataset, or even to perform on-device learning. This problem becomes more problematic for nonparametric models such as Neural Processes [15,20] which require the training dataset to be stored for inference. It is therefore appealing to reduce the size of the dataset, both at the instance [11,27,27] and the dataset level, such that we select only a small number of samples from the dataset, called the subset, each of which contains only a few selected input features (e.g. pixels) or data points. The selected subset could then be used for tasks such as reconstruction of the original dataset or for prediction tasks such as classification of individual instances.
A simple and naïve way to obtain such a subset is random selection, but it is highly sub-optimal given that it treats all the elements in the set equally. However, the pixels from a given image and examples from a dataset will have varying degree of importance [19] to a target task, whether it is We validate SSS on multiple datasets and tasks such as 1D function regression, 2D image reconstruction and classification for both feature and instance selection. The results show that SSS is able to select samples with minimal degradation on the target task performance, largely outperforming random selection and other selection methods. Our contribution in this work is threefold: • We propose a novel two-stage stochastic subset selection method that learns to select a subset from a larger set with linear time complexity and with minimal degradation in performance on the downstream task.
• We propose a framework that trains the subset selection model over a distribution of sets so that it can generalize to unseen tasks such as in few-shot classification.
• We validate the efficacy and generality of our model on various datasets for feature selection from an instance and instance selection from a dataset and show that it significantly outperforms the relevant baselines.

Related Work
Set encoding Recently, extensive research efforts have been made in the area of 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 [36,34,45,38]. However, these models have limited expressive power. Yet, approaches such as Stochastic Deep Network [9] and Set Transformer [24] consider the pairwise interactions among set elements and hence can capture more complex statistics of set distributions. They often show better performance but require at least O(n 2 ) time complexity.
Subset Sampling There exist some works to handle large sets. Dovrat et al. [11] propose learning to sample a subset by generating virtual points, then matching them back to a subset of the original set. Such generation and matching process is highly inefficient, however, our method learns to select from the original elements and does not suffer from such overhead. Wang et al. [43] propose to distill the knowledge of a large dataset to the small number of artificial data points. However, they are only for faster training and do not capture the statistics of the original set. Moreover, they are artificial and may differ from the original set, making it less applicable to other tasks. Several works [33,35,27,12,32] propose farthest point sampling, which selects k points from a set by ensuring that the selected ones are far from each other on a metric space.
Image Compression Due to the huge demand for data transfer over the internet, some works attempt to compress images with minimal distortion. These models [41,37,31,25] typically consist of encoder and decoder, where the encoder transfers the image with a compact matrix and the decoder reconstructs the image back. These methods, while highly successive in the image compression problem, are less flexible than ours. Our model can be applicable to any type of sets 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.
Representation learning Our instance-sampling model is also related to the Variational Auto Encoder (VAE) [22]. However, while VAE learns a compact representation of a data point, our model learns a compact representation of a set. Balın et al. [1] propose a global feature selection model for reconstruction of the input data from selected features via unsupervised learning. Chen et al. [7] learns instancewise feature selection with the goal of model interpretation by extracting subset of features most informative for a given sample.
Active Learning Active learning methods aim at selection of data points for labeling given a small labeled set. This domain is different from ours 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 nonparametric models by reducing the sizes of the inputs -pixels or instances, and this greatly differs from the goal of active learning. Methods such as [39,8,44] all tackle the data selection problem in the active learning setting.

Preliminaries
We consider a set D = {d 1 , . . . , d n } as an input with D ∼ p(·) for some unknown data distribution, where individual d i 's either represent a pair of input x i and target y i or a feature such as the pixel value of an image. We assume there exists a subset D s = {s 1 , . . . , s k } ⊂ D such that (·, D) ≈ (·, D s ) for an arbitrarily defined loss function (·, D) that we are interested in optimizing over the full set D with k n. In what follows, we present a method that learns the conditional distribution p ξ (D s |D) of the subset D s via a two stage subset selection method consisting of candidate selection and autoregressive subset selection as illustrated in Fig. 2. In general, we minimize the loss function with respect to ξ the parameters of the subset selection model.

Stochastic Subset Selection
To select D s , we need to model the pairwise interactions among the elements of D 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 i 's, are high dimensional, modeling pairwise interactions becomes computationally infeasible since we need to  compare every element in D with all the other elements. This computational bottleneck motivates the first stage of SSS of which the goal is to construct a smaller subset D c , which we 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.

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 of D and the individual elements d i ∈ D. Specifically, we first encode the dataset D to a single representation D e as follows: where g is a neural network (Set Encoder in Fig. 2-(a)), which projects each element in D independently to a lower dimension and n = |D|. This encoding scheme is similar to the one proposed in DeepSets [45] except that we do not perform message-passing between the set elements. Then we concatenate every d i with D e , denoted as d i . It ensures each element of D has a global view of all the other elements in the set at a course level. For every d i ∈ D, we can sample a mask z i as follows: where h is a neural network that outputs the logits for the probability that d i is in the candidate set D c and σ(·) is Sigmoid function, and Ber denotes the Bernoulli distribution. z i is a binary random mask where z i = 1 implies that d i is an element in D c . We concatenate all z i 's to obtain a single vector Z = {z i } n i=1 . Note that we feed the output of h into the Sigmoid function and normalize its output over all samples in D to obtain a valid probability distribution and induce sparsity in the selected candidate set D c . Since sampling from the Bernoulli distribution is not differentiable, during training, we use the continuous relaxations of the Bernoulli distribution [30,18,13] to sample the mask z i ∼ Ber(ρ(d i ) for d i as illustrated in the black boxes in Fig. 2-(a). Although pairwise interactions are not considered in this stage, the ablatation studies show that learning p ξ (z i |d i , D) leads to selecting highly representative samples compared to random selection.
Constraining the size of D c For computational efficiency, we want to restrict the size of D c to save computational cost when constructing D s . Hence we introduce a sparse Bernoulli prior r(Z) and minimize the KL divergence along with the target downstream task loss (·, D s ) w.r.t ξ as follows: where β > 0 is a hyperparmeter used to control the sparsity level in Z.

Autoregressive Subset Selection
At this stage we have a set D c with m = |D c | |D|, which is small enough to perform fine-grained subset selection through pairwise modeling. To select a subset with k elements from D c , we requires k iterative steps. As shown in Fig. 2-(b), at the time step t, we have the subset D  Fig 2-(b)) for modeling pairwise interactions between the elements of an input set, we autoregressively compute the interactions at time step t as follows: where σ(·) denotes the Sigmoid function and ϕ • f is a composition of two neural networks, one that computes interaction scores between elements in D (t) c and the other which outputs element-wise logits using the interaction scores. Further,π (t) is the vector of interaction scores for all elements in D (t) c at the current time step t. Givenπ (t) > 0, we can compute the probability of an element s That is, we normalize the vector of interaction scoresπ (t) over all the elements in D (t) c at time step t to obtain a valid probability distribution. The key to avoiding redundant elements in D s lies in the fact that for each element added to D s , its selection is conditioned on both the candidate set D Using Eq. 5 we can sample an element where Cat is Categorical distribution. During training, it can be expensive to sample k times from the Categorical distribution since it involves computing Eq. 4 k times. We remedy this by selecting l elements from D (t) c , which reduces the number of iterations to k/l for selecting k elements. We may sample l elements from multinomial distribution with the probability π (t) without replacement. However, this sampling procedure is not differentiable, and hence it cannot be trained with backpropagation. Instead, we independently sample l elements from the continuous relaxation of Categorical distributions [30,18] using the same probabilities in Eq. 5 to approximate sampling from the multinomial distribution as shown in Fig. 2-(b). Since we want to simulate sampling without replacement, we discard all elements sampled more than once. This sampling procedure guarantees that we get at most l elements at each iteration. A similar sampling procedure is adopted in previous works [1,6]. We present this training algorithm in the Appendix.
Inference Complexity of SSS. The inference complexity of SSS depends heavily on the choice of the function f . Using MAB as f , the inference complexity of SSS is O(n) + O(k 2 m/l) where n, m, k correspond to |D|, |D c | and |D s | respectively. Fig. 3-(a). We jointly train the SSS model and a neural network predicting pixel values to minimize the loss function w.r.t θ and ξ as follows:

Tasks
We enforce sparsity on the subset D s by minimizing the KL-divergence between the mask probability p ξ (Z|D) and sparse prior r(Z). Moreover, minimizing the negative log likelihood, which corresponds to (·, D s ) in Eq. 3, ensures that the constructed D s is the most representative for the downstream tasks. We implement p θ (Y |X, D s ) as an Attentive Neural Process (ANP) [21]. The ANP takes D s as input and predicts a distribution of the elements in the original set D. It mimics the behaviour of a Gaussian Process but with reduced inference complexity. We present experimental results for this task in Section 4.1 and a corresponding graphical model depiction in Appendix D.
Set Classification/Prediction As shown in Fig. 3-(b), we train a neural network to predict a single target value y D for the subset D s of the given dataset D, where the set D is a collection of the features  from a single instance such as pixels of an image. For instance, the target is either the class of an image for classification or the statistics of the set for regression. Here, our goal is learning to select the most representative subset D s ⊂ D such that we can maximize the log likelihood log p θ (y D |D s ) with computational efficiency. In order to achieve this goal, we jointly train the SSS model and the neural network which predicts the target value y D for D s to minimize the negative log-likelihood -loss function (·, D s ) described in Eq. 3 and enforce sparsity on Z, the selection masks for the candidate set, by minimizing the KL-divergence as follows: We provide experimental results in Section 4.2 and a corresponding graphical model Appendix D.
Dataset Distillation: Instance Selection In this task, we are given a collections of datasets D = , and the goal is to select the most representative subset D s with k = |D s | |D| for each dataset D = {d 1 , . . . , d n } ∈ D, where d i is a data point uniformly sampled from the entire datasets D. Using CelebA dataset as an illustrative example, shown in Fig. 3, D consists of n randomly sampled faces from the entire dataset and the task is to construct a subset, D s , most representative of D.
In order to learn to select the subset D s from each D ∈ D with unsupervised learning, we jointly train the SSS model and a generative model such that the SSS model chooses the most representative subset so that the generative model can reconstruct all the images d i ∈ D from the subset. Naïvely, we can minimize the sum of negative log-likelihood di∈D − log p θ (d i |D s ) for the loss function (·, D s ) and KL divergence in Eq. 3. However, we find that the generative model outputs a mean images for all d i . To capture variations of different images, we introduce three latent variables α i , c i , and w i which both depend on d i . We provide graphical model illustration of this task in the Appendix D. Since it is intractable to compute the log likelihood log p θ (d i |D s ) 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 (·, D s ) in Eq. 3 as follows: where p ψ (·) are priors on their respective latent variables, r(Z) is sparse prior over the mask for candidate set selection in SSS, p θ (·) is the decoder to reconstruct d i , and all variational posteriors q φ (·) 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 to minimize the objective in Eq. 8 w.r.t θ, φ, and ξ for all D ∈ D and leverage the optimized SSS to select a few representative instances of the dataset, which results in distilled dataset. We report all the experimental results in Section 4.3 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 [40] 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 (·, D s ), to jointly train the Prototypical Networks and SSS. Note that we use D s , the subset of the support set, for computing loss and prediction. By learning to 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 x * . Experimental results for this task are in Section 4.3 and an accompanying graphical model description is provided in Section D of the Appendix.

Experiments
Detailed model architectures and hyperparamters are given in Section of the Appendix.
Baseline for feature selection 1) Random Selection: this model randomly selects k elements from the set D for the subset D s . 2) Learning to Sample (LTS) [11]: this model learns to generate k virtual elements, which can be matched to elements in D, that optimized for the downstream task.
Baseline for instance selection 1) Random Selection and two instance selection algorithms 2) k-Center-Greedy: this algorithm iteratively selects elements in D closest to a set of centroids. 3) FPS: this algorithm iteratively selects the most distant elements to a randomly initialized D s . iid ∼ N (0, K XX + σ 2 y I n ) for i = 1, . . . , N where K XX is a squared-exponential kernel with the set of inputs X = {x 1 , . . . , x n } and σ 2 y is variance for small likelihood noise. This will lead to a collection of the sets (D (1) , . . . , D (N ) ). We train our model which consists of the subset selection model p ξ (D s |D) and a task network p θ (Y |X, D s ), which is Attentive Neural Process (ANP) [21], on this dataset and report the negative log-likelihood (NLL). Fig. 4a shows the performance (NLL) of SSS compared to the baselines -Random Selection and LTS. As shown in Fig. 4a, SSS outperforms the baselines, verifying that the subset selection model p ξ (D s |D) learns a meaningful distribution over subsets. We visualize a sample reconstructed function and the selected points by each models in Fig. 5. As shown in the rightmost figure (Fig. 5), SSS tends to pick out more elements (presented as red dots) in the drifting parts of the curve, which is reasonable since those are harder to reconstruct than the others. However, the other baselines sometimes fails to do that, which leads to inaccurate reconstructions.

Image Reconstruction
Given an image, we learn to select a representative subset of pixels that best reconstructs the original image. Here, x i is the 2d pixel coordinates and y i ∈ R 3 is the RGB pixel value. We use ANP to reconstruct the remaining pixels from the subset D s = {x i l , y i l } k l=1 constructed by each subset selection models. We conduct the experiment on the CelebA dataset [29]. Fig. 4b shows that our model significantly outperforms Random Selection and LTS in terms of NLL. We provide qualitative examples for the competing methods in the Appendix.

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 subset selection models and classifiers with CelebA dataset, where each classifier performs binary classification for 40 attributes of a face. Note that only the selected pixels of an image by the subset selection models are used for prediction and the other pixel values are set to zeros.
We report the mean AUC score on all 40 attributes for varying sizes of D s . Fig. 4c shows that using only 500 pixels (∼1.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 of 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 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 LTS 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.

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 to minimize the objective in Eq. 8 with D. After training, we discard the generative model and leverage the subset selection model to choose a few representative images from the full CelebA dataset.
We evaluate the selected subset with the Fréchet Inception Distance (FID) [16], which measures similarity and diversity between two datasets and compare SSS to k-Center-Greedy, FPS and Random Selection. We report the experimental results in Table 1 where our selection method achieves the lowest FID score for all selection sizes. Specifically our model outperforms all the baselines for selecting very few instances since our method is able to model the interactions within the dataset and hence selects the most representative subset. Additionally, given that the dataset is highly imbalanced, k-Center-Greedy 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 the Appendix. Classification In this task, we perform few-shot classification with the miniImageNet dataset where the models select 1, 2, or 5 instances from the support set with size 20. As shown in 2, we compare our model against Random Selection, FPS, and k-Center-Greedy. SSS learns to select more representative prototypes than the others especially for small D s 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 the Appendix.

Ablation
We perform extensive ablation studies on the two stages of SSS using the function reconstruction task presented in Section 4.1. 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. In Fig. 6, 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 a candidate set with redundant elements.

Candidate Selection Only
Random Selection with Autoregressive Subset Selection To show how much the candidate selection stage contributes to the performance of SSS, we replace it with random selection. As shown in Fig. 6, we find that while this model performs better than the model with only candidate selection, it 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. AutoRegressive Subset Selection Only As shown in Fig. 6 we observe that the model with only autoreressive subset selection performs significantly better than the model with the candidate selection and the model with random selection followed by autoregressive subset selection. While this model performs well, it is not very practical due to the high computational cost when the size of the set becomes large. Additionally, processing lots of elements with the interaction model can lead to flattened attention scores over the elements, which makes it difficult to distinguish important elements from the others. The candidate selection stage plays an 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 predictive distribution is E p ξ (Ds|D) [p θ (y D |D s )] for classification and we approximate it with Monte Carlo sampling. However in all experiments, we only report the result with one sampled subset, since it gives the best tradeoff in memory and computation. We compare it against another variant: SSS-MC which use 5 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 a single subset sample achieves 90.93% (Fig. 4c) accuracy, which is comparable to SSS-MC. The result justifies that our model SSS achieves good performance for target tasks with memory and computational efficiency. Storage comparisons for the selected pixels with these variants of SSS are in the Appendix.
Cost of Running SSS We report the cost of running SSS in selection of subset of pixels, which has the largest set size (38804 pixels). In this experiment, we select 500 pixels in total with l = 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.

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 a stack of Categorical distributions to iteratively select a representative 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 a 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 our work can facilitate the use of machine learning algorithm in resource-limited systems such as mobile and embedded devices.

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 on 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 miniImageNet dataset. Finally we provide more details on specifications of models used in the experiments.

Algorithm 1 Greedy Training Algorithm
Input k (max subset size) m (mini-batch size) l (# elements selected at each iteration) p(D) (distribution of sets) α (learning rate) (·, D s ) (loss function) Output trained model with converged θ and φ 1: Randomly initialize θ, φ, ξ.  In order to reduce the computational cost at training time, we use a greedy training algorithm with stochastic gradient descent as described in Algorithm 1. It selects only the top l elements from the candidate set to train the auto-regressive model by minimizing the target loss on the selected samples and randomly selects the remaining k − l elements from the candidate set. As a result, we do not have to run the auto-regressive model k/l time during training, which significantly reduces the computational cost and shows reasonable performance.

C 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, as described in Algorithm 2. We do not use the greedy training algorithm. Instead, we autoregressively select k elements from the candidate set D c as described in line 12 from Algorithm 2 to construct the representative subset D s .

D 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 representative data points and (d) instance selection for few-shot classification. Only the shaded circles denote observable variable and the others are latent variables.

E Instance Selection Samples
In this section, we show more examples of our 1D and CelebA experiments on how the models select the set elements for the target task. 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 the set elements randomly, it can leave out important part of the 1D curve leading to wrong reconstructions. LTS also selects insufficient amount of set elements in some parts of the curves, resulting in suboptimal reconstructions. Figure 4b shows samples of reconstructed images for varying selection number of pixels. Additionally in Figure 10, we show the selected pixels of our model for both the classification and reconstruction task. For the attribute classification task, the model tends to 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.

F 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(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 in https://github.com/rusty1s/pytorch_cluster. Further, we provide qualitative results on a single dataset in Figure 11 where we show how our model picks 5 instances from the full set of 200 images face images.

F.1 Dataset Distillation: Classification
In Figure 12 we provide visualizations for the instance selection problem as applied to the few-shot classification task. Here, we go from a 20-shot to a 1-shot classification problem where the prototype is selected from the support using SSS. The selected subset is then used in place of the support set and used to classify new query instances.

G Model Specifications
In this section, we describe the main components of our Stochastic Subset Selection models - 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 D e , following DeepSets [45].
We parameterize ρ(·) with a 3 layered feedforward neural network h followed by sigmoid function as follows as described in Equation (2)  and feed it to set transformer [24] for set classification as follows: where rFF is a row-wise feedfoward layer which processes each instance independently and Multihead denotes Multihead Attention [42] 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) s .

G.1 Attention
We details on Attention and Multihead Attention here for completeness. For a more thorough exposition, we refer the reader to Vaswani et al. [42] and Lee et al. [24].
An attention module computes the following interactions using the dot product: where Q ∈ R n×dq are the n query vectors each with of dimension d q . K ∈ R nv×dq and V ∈ R nv×dq are the keys and values respectively. Interactions are modelled through QK and ω is an activation function such as Softmax or Sigmoid.
where O j = Att(QW Q j , KW K j , V W V j ; ω j ) The Multihead Attention module has learnable parameters λ where λ = {W Q j , W K j , W V j } h j=1 and W Q j , W K j ∈ R dq×d M q , W V j ∈ R dv×d M v and W O ∈ R hd M v ×d . In all our experiments, we use Sigmoid as the activation function.  Figure 12: Sample visualization of prototype selection for the miniImagenet dataset on the few-shot classification task. Each row represents a set that corresponds to the support from which a prototype is selected for the few-shot classification task.

H Function Reconstruction: Ablation
In the ablation studies in Section 4.4 of the main paper, we investigated the contribution of the Candidate Selection stage (Stage 1) and the AutoRegressive subset selection stage (Stage 2). Additionally, we replaced Stage 1 with random selection in SSS and compare the performance of these three variants with SSS. In Figure 13, we provide qualitative results on the function reconstruction tasks for these 3 variants as well as SSS. Figure 13 is meant to supplement the results presented in Section 4.4 of the paper. As can be seen from the figure, using the Candidate Selection stage only to select elements generally focuses on the difficult sections of the curves ignoring the other parts which are relatively easier 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 a sub-par 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 we randomly sample the candidate 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 13. 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.