Task Conditioned Stochastic Subsampling

1 INTRODUCTION . Deep learning algorithms operate on large volume of high-dimensional dense inputs such as the pixels of an image ( Deng et al. , 2009 ; Krizhevsky et al. , 2009 ; Liu et al. , 2015 ) . Training or evaluating a model with such data is computationally expensive and several works ( Balın et al. , 2019 ; Huijben et al. , 2019 ; Yoon et al. , 2018 ) have proposed subsampling techniques to subsample such dense inputs . Subsampling methods have the potential to drastically reduce the data acquisition effort and reduce the inference time of algorithms that operate on dense inputs . Additionally , subsampling techniques have found applications in medical research for the purpose of interpretation ( Ribeiro et al. , 2016 ) . However , these methods have a major drawback in that they are only applicable when each feature ( such as a pixel ) of the input is 1-dimensional . This restriction is imposed by the way the subsampling methods are designed : the input ( e.g . an image ) , is flattened to a single vector and a model predicts a binary mask for each feature ( e.g . pixels ) . This becomes problematic when we consider a 3- channel image . Flattening out the image results in ambiguities as to which pixels to select since channels are treated independently . For instance , in order to perform subsampling on CIFAR10 images , INVASE ( Yoon et al. , 2018 ) and DPS ( Huijben et al. , 2019 ) convert the images into single channel images ( e.g . grey-scaled images ) before subsampling pixels . In more extreme cases such as subsampling ( different from instance-wise feature selection ) of training instances , each feature is itself an image ( each possibly multi-channel ) and hence those subsampling techniques can not be used . Finally , as we show in our experiments , most of these methods fail under extremely low subsampling rates and their performance is similar to random sampling in this setting . In order to tackle these limitation , we propose to consider each feature or instance as an element of a set . We formulate the subsampling problem as selecting a subset of features , or instances that minimizes the performance degradation of an arbitrary model on an arbitrary task such as image classification , regression or instance subsampling for tasks such as few-shot classification . As a result , there are several advantages compared to the previous works . First , we can handle multi-dimensional feature , such as RGB pixel value , if we use set functions parameterized with expressive neural networks ( Zaheer et al. , 2017 ; Lee et al. , 2018 ) . Second , a subsampling model with set function can process arbitrary number of elements . Thus the model is robust to a wide range of subsampling rates at test time even when trained with fixed sampling rate . Lastly , the set-based formulation unifies the feature and instance subsampling tasks under a single framework . However , it is prohibitively expensive to process all the set elements ( e.g . all the pixels in an image ) with expressive set function such as Set Transformer ( Lee et al. , 2018 ) due to the quadratic computational cost with respect to the number of elements . Hence we propose an efficient two-stage selection method . In the first stage , as shown in Fig . 1- ( a ) , we learn the Bernoulli sampling rate for individual samples and efficiently screen out less important ones resulting in a subset which we call the candidate set . The second stage is more fine-grained and designed to select a smaller subset from the candidate set by considering the relative importance of the samples in the candidate set using a conditionally dependent Categorical distribution through an iterative procedure as shown in Fig . 1- ( b ) . Once optimized , the resulting subsampling model can perform stochastic subsampling of a given input with linear time complexity . We call the resulting model Stochastic SubSampling ( SSS ) which is a general subsampling framework and can be applied to both feature and instance selection . 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 experimental results show that SSS is able to subsample with minimal degradation on the target task performance , largely outperforming the relevant baselines . We summarize our contribution as follows . • We reformulate the feature/instance subsampling problem by treating all the features/instances as members of a set . This allows us to apply set-based functions to the subsampling problem and extend its range of applicability . • We propose a set-based two-stage stochastic subsampling method that learns to efficiently subsample a set with minimal performance degradation on a target task . • We validate the efficacy and generality of our method on various datasets for feature selection in the input space ( e.g . pixels ) and instance selection from a dataset and show that it significantly outperforms the relevant baselines . 2 RELATED WORK . Set Functions Recently , extensive research efforts have been made in the area 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 ( Ravanbakhsh et al. , 2016 ; Qi et al. , 2017b ; Zaheer et al. , 2017 ; Sannai et al. , 2019 ) . However , these models have limited expressive power . Yet , approaches such as Set Transformer ( Lee et al. , 2018 ) consider the pairwise interactions among set elements and hence can capture more complex statistics of set distributions . They often show better performance but require at least O ( n2 ) time complexity . Feature Selection Recent interest in deep learning based subsampling methods has produced many works mostly applied to feature selection . In Balın et al . ( 2019 ) , continuous approximation of the Concrete Distribution ( Maddison et al. , 2016 ) is used for global feature selection where a fixed set of features are sampled across an entire dataset . In Chen et al . ( 2018 ) , instance-wise feature selection is used for interpretation of deep learning models applied to medical data . Dovrat et al . ( 2019 ) propose learning to subsample by generating virtual points , then matching them back to the original input . Several works ( Qi et al. , 2017a ; c ; Li et al. , 2018b ; Eldar et al. , 1997 ; Moenning & Dodgson , 2003 ) also propose farthest point sampling , which selects k points from an input by ensuring that the selected samples are far from each other on a metric space . However our work is most similar with the recent works of Yoon et al . ( 2018 ) and Huijben et al . ( 2019 ) which learn a subsampling model conditioned on an given task . However these models have limitations both in terms of their range of applicability and poor performance under extremely low subsampling rates . Our method on the other hand is flexible , more general and applicable to a wide range of subsampling problems . Image Compression Due to the huge demand for data transfer over the internet , some works attempt to compress images with minimal distortion . These models ( Toderici et al. , 2017 ; Rippel & Bourdev , 2017 ; Mentzer et al. , 2018 ; Li et al. , 2018a ) typically consist of an encoder and decoder , where the encoder transforms the image with a compact matrix and the decoder reconstructs the image . These methods , while highly successful for the image compression problem , are less flexible than ours . Our model can be applied to any type of set structured data while the aforementioned models mainly work for images represented in tensor form . Furthermore , our method can be applied both at the instance and dataset level . Active Learning Active learning aims to select data points for labeling given a small labeled set . This domain is different from ours since active learning does not consider the label information ( for the selected data points ) but our method does utilize label information . Also , our motivation is quite different . We focus on optimal subsampling conditioned on an arbitrary task and this greatly differs from the goal of active learning . Methods such as ( Sener & Savarese , 2017 ; Coleman et al. , 2019 ; Wei et al. , 2015 ) all tackle the data selection problem in the active learning setting . 3 APPROACH . 3.1 PRELIMINARIES . We consider a set D = { d1 , . . . , dn } as an input with D ∼ p ( · ) for some unknown data distribution , where individual di ’ s either represent a pair of input xi and target yi or a feature such as the pixel value of an image . We assume there exists a subset Ds = { s1 , . . . , sk } ⊂ D such that ` ( · , D ) ≈ ` ( · , Ds ) 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ξ ( Ds|D ) of the subset Ds via a two-stage subsampling method consisting of candidate selection and autoregressive subset selection as illustrated in Fig . 2 . In general , we minimize the loss function Ep ( D ) [ Epξ ( Ds|D ) [ ` ( · , Ds ) ] ] with respect to ξ the parameters of the subsampling model . 3.2 STOCHASTIC SUBSAMPLING . To select Ds , 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 , di ’ 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 Dc , 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 . 3.3 CANDIDATE SELECTION . We formulate the candidate selection problem as a random Bernoulli process where the parameters of the Beronulli distribution are conditioned on the set representation of D and the individual elements di ∈ D. Specifically , we first encode the dataset D to a single representation De as follows : De = 1 n n∑ i=1 g ( di ) , ( 1 ) 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 ( Zaheer et al. , 2017 ) except that we do not perform message-passing between the set elements . We then concatenate every di with De , denoted as di . This ensures that each element of D has a global view of all the other elements in the set at a coarse level . For every di ∈ D , we can sample a mask zi as follows : pξ ( zi|di , D ) = Ber ( zi|ρ ( di ) ) , ρ ( di ) = σ ( h ( di ) ) ∑n j=1 σ ( h ( dj ) ) ( 2 ) where h is a neural network that outputs the logits for the probability that di is in the candidate set Dc and σ ( · ) is the Sigmoid function , and Ber denotes the Bernoulli distribution . zi is a binary random variable where zi = 1 indicates that di is an element in Dc . We concatenate all zi ’ s to obtain a single vector Z = ( z1 , . . . , zn ) . 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 Dc . Since sampling from the Bernoulli distribution is not differentiable , during training , we use the continuous relaxations of the Bernoulli distribution ( Maddison et al. , 2016 ; Jang et al. , 2016 ; Gal et al. , 2017 ) to sample zi ∼ Ber ( ρ ( di ) for di as shown in the black boxes in Fig . 2- ( a ) . Although pairwise interactions are not considered in this stage , the ablation studies show that learning pξ ( zi|di , D ) leads to selecting highly informative samples compared to random selection . Constraining the size of Dc For computational efficiency , we want to restrict the size of Dc to save computational cost when constructing Ds . Hence we introduce a sparse Bernoulli prior r ( Z ) and minimize the KL divergence along with the target downstream task loss ` ( · , Ds ) w.r.t ξ as follows : Ep ( D ) [ Epξ ( Ds|D ) [ ` ( · , Ds ) ] + βKL [ pξ ( Z|D ) ||r ( Z ) ] ] , where pξ ( Z|D ) = n∏ i=1 pξ ( zi|di , D ) ( 3 ) where β > 0 is a hyperparmeter used to control the sparsity level in Z .

Summary: The paper proposes a Stochastic SubSampling method (SSS) that is helpful in reducing the size of input dataset while preserving reasonable performance on a target task. SSS is two-stage and set-based, which can be jointly optimized with arbitrary downstream tasks. Extensive experiments verify that SSS do perform well in the sense of efficiency and generalization a variety of tasks and datasets.