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