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 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 endto-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 . 1 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 it 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-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. , 2018b ; b ) and the dataset level , such that we selects only a 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 set ( either each instance or the entire dataset ) or for a 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 a 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 can not 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 , we propose to learn a sampler that learns to sample the most important samples for a given task , by training it jointly with the target task and additionally meta-learn a sampler over a distribution of datasets for instance selection in the classification task . 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 a 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 a 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 on downstream tasks . Our model can benefit from a wide range of practical applications . For example , when sending an image to an edge device with low computing power , instead of sending the entire image , we could send a subset of pixels with their coordinates , which will reduce both communication and inference cost . Similarly , edge devices may need to perform inference on a 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-parametric 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 reconstruction 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 an 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 a 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 an instance and instance selection from a dataset , on 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-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 ( 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 as Stochastic Deep Network ( De Bie et al. , 2018 ) 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/regression tasks ; however , they have run time complexities of O ( n2 ) or higher . Subset sampling There exist some works which have been proposed to handle large sets . Dovrat et al . ( 2019 ) proposed to learn to sample a subset from a set by generating k virtual points , then matching them back to a 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 et al . ( 2018 ) proposed to distill the knowledge of a large dataset to a 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. , 2017a ; c ; Li et al. , 2018b ; Eldar et al. , 1997 ; Moenning & Dodgson , 2003 ) propose farthest point sampling , which selects k points from a set by ensuring that the selected samples are far from each other on a given metric space . Image Compression Due to the huge demand for image and video transfer over the internet , a 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 a pair of encoder and decoder , where the encoder will transfer the image into a compact matrix to reduce the memory footprint and communication cost , while the decoder is 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 as 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 ) ( Kingma & Welling , 2013 ) . However , while VAE learns a compact representation of a data point , our model learns a compact representation of a 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 a given sample . Our method also falls in this category . Active Learning Active learning methods are aimed at selection of data points for labeling given a small labelled set . This domain is 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-parametric 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 APPROACH . 3.1 PRELIMINARIES . In this work , we consider data of the type D = { d1 , . . . , dn } where individual di ’ s are possibly represented as input xi and target yi . D is the complete set and we assume that within D , there exists a subset Ds = { si , . . . , sk } ⊂ D such that k n and that for an arbitrarily defined loss function ` ( . , D ) that we are interested in optimizing over the full set D , Ds can be used as a proxy for D such that ` ( . , D ) ≈ ` ( . , Ds ) . In what follows , we present a method that learns the conditional distribution p ( Ds|D ) of the subsetDs via a 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 Ds , Ep ( Ds|D ) [ ` ( . , Ds ) ] . When the set D itself follows a distribution of sets as in the meta-learning framework , then the objective becomes ED [ Ep ( Ds|D ) [ ` ( . , Ds ) ] ] . In essence , we seek to construct a subset Ds that is optimally representative of the full set D w.r.t ` ( . ) .

Summary: 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.