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: This paper proposes a 2-stage sampling method that can find its use in either feature selection or instance (training examples) subsampling. At a high level, the authors frame the problem as a subset selection problem with the goal to minimize the performance difference for a downstream task between using the original set and the downsampled set. The proposed method consists of two stages: 1) select a subset (called candidate set) using Bernoulli variables as the binary mask, and the Bernoulli distribution's parameter depends on the corresponding example and the original set, 2) select elements from the candidate set iteratively to form the final selected subset, and the selection here is based on a softmax whose parameters depend on the interactions between the remaining candidate set and the selected subset so far. To support the merits of the proposed method, the authors provide experiment results on feature selection and dataset distillation tasks.