Training Data Distribution Search with Ensemble Active Learning

1 INTRODUCTION . Deep Neural Networks ( DNNs ) have become the dominant approach for addressing supervised learning problems . They are trained using stochastic gradient methods , where ( i ) data is subsampled into mini-batches , and ( ii ) the network parameters are iteratively updated by the gradient of the parameter weights , relative to a loss function for the given labeled mini-batch of data . The architecture of the DNN , hyper-parameters of the training process , and distribution of the dataset used are all crucial ingredients that impact the final performance of the DNN . Given the remarkable success of this recipe for supervised learning , there is significant interest in automating the end-to-end process of applying DNNs to real-world problems ( Wong et al. , 2018 ; He et al. , 2018 ; 2019 ) . While there has been a considerable effort towards methods and frameworks that automate DNN architecture search ( Elsken et al. , 2018 ; Luo et al. , 2018 ; Liu et al. , 2018 ; Xie et al. , 2019 ; Dong & Yang , 2019 ) and training hyper-parameter search ( Golovin et al. , 2017 ; Mutny & Krause , 2018 ; Kandasamy et al. , 2019 ) ; the process of searching for the right training data distribution ( also called dataset curation ) is still performed by experts , requiring several heuristics and significant manual effort . With the rapid growth in the availability of labeled data for large-scale vision tasks , to the order of billions of samples ( Sun et al. , 2017 ; Zeki Yalniz et al. , 2019 ) , automating the training data subset search would make the application of DNNs much easier for non-experts , and potentially lead to datasets and models that outperform those that were curated by hand . With the cost of data storage also becoming increasingly important , the ability to automatically identify data to remove from the training distribution is appealing from a practical perspective . In this paper , we present a simple yet effective method to perform a training data subset search by using ensemble Active Learning ( AL ) . The typical goal of AL is to to select , from a large unlabeled dataset , the smallest possible training set to label in order to solve a specific task ( Cohn et al. , 1994 ) . We instead propose to use AL to build data subsets of a large labeled training dataset that give more accurate DNNs in less training time . We tackle several key issues that have not been addressed so far by state-of-the-art AL methods . The first is the difficulty in scaling the number of models for ensemble AL . While it seems intuitive that more ensembles can improve performance , existing studies show no gains in AL performance beyond 10 models , and even recommend the use of only 5 models ( Lakshminarayanan et al. , 2017 ; Ovadia et al. , 2019 ) . Implicit ensemble methods for AL have been found similarly ineffective , but only evaluated in settings with with < 10 models ( Beluch et al. , 2018 ) . In this study , we propose the use of implicit ensembles with hundreds of training checkpoints from different experimental runs , and empirically demonstrate the effectiveness of this approach . Second , we switch to a large-scale experimental setting compared to what is typically used for AL experiments . For example , Beluch et al . ( 2018 ) and Sinha et al . ( 2019 ) only acquire 40k and 64k samples at a time respectively on ImageNet , never use more than 30 % of the dataset , and do not compare to the full dataset performance . In contrast , for ResNet-18 training on ImageNet , we acquire up to 500k samples at a time and use 80 % of the dataset , improving top-1 accuracy by 0.5 % over a model trained with the entire dataset . Third , we demonstrate the transferability of the subsets obtained with our approach between DNNs with very different network capacities , ranging from 10-layer ResNets ( He et al. , 2016 ) to 121-layer DenseNets ( Huang et al. , 2016 ) . Given the rapid pace of model development , this shows that our data subsets remain valuable even after a particular model is surpassed by newer architectures . This has been difficult to achieve in prior AL studies ( Lowell et al. , 2018 ) , and not investigated in detail for large-scale vision tasks . To summarize , our contributions are as follows : ( i ) we propose a simple approach to scale up ensemble AL methods to hundreds of models with a negligible computational overhead at train time ; ( ii ) we evaluate several methods to reduce the size of existing large datasets with AL ; and ( iii ) we conduct a detailed empirical study on three popular image classification benchmarks , studying the impact of key design choices , and the robustness of the selected subsets to changes in model architecture . Our study has significant practical implications for DNN training with stochastic gradient methods , as well as for storage of large training datasets . 2 BUILDING DATA SUBSETS WITH ACTIVE LEARNING . Consider a distribution p ( x , y ) over inputs x and labels y . In a Bayesian framework , the predictive uncertainty of a particular input x∗ after training on a dataset L is denoted as P ( y = k|x∗ , L ) . The predictive uncertainty will result from data ( aleatoric ) uncertainty and model ( epistemic ) uncertainty ( Kendall & Gal , 2017 ) . A model ’ s estimates of data uncertainty are described by the posterior distribution over class labels y given a set of model parameters θ . This is typically the softmax output in a classification DNN . Additionally , the model uncertainty is described by the posterior distribution over the parameters θ given the training data L ( Malinin & Gales , 2018 ) : P ( y = k|x∗ , L ) ︸ ︷︷ ︸ Predictive = ∫ P ( y = k|x∗ , θ ) ︸ ︷︷ ︸ Data p ( θ|L ) ︸ ︷︷ ︸ Model dθ . ( 1 ) We see that , uncertainty in the model parameters p ( θ|L ) induces a distribution over the softmax distributions P ( y = k|x∗ , θ ) . The expectation is obtained by marginalizing out the parameters θ . Unfortunately , obtaining the full posterior p ( θ|L ) using Bayes ’ rule is intractable . If we train a single DNN , we only obtain a single sample from the distribution p ( θ|L ) . Ensemble uncertainty estimation techniques approximate the integral from Eq . 1 by Monte Carlo estimation , generating multiple samples using different members of an ensemble ( Lakshminarayanan et al. , 2017 ) : P ( y = k|x∗ , L ) ≈ 1 E ∑ e∈E P ( y = k|x∗ , θe ) , θe ∼ q ( θ ) . ( 2 ) where q ( θ ) represents an approach used for building the ensemble . The strength of the ensemble approximation depends on both the number of samples drawn ( E ) , and how the parameters for each model in the ensemble are sampled ( i.e. , how closely the distribution q ( θ ) matches p ( θ|L ) ) . Algorithm 1 Building Data Subsets with Active Learning . Initialize parameters of acquisition model { θ ( e ) a } Ee=1 and subset model θs . initialization Initialize data subset S Compute acquisition function α ( x , Ma ) ∀ x ∈ L . acquisition Append Ns samples with the maximum acquisition function to S Use S to optimize parameters θs of subset model . optimization Uncertainty estimation acts as the core component of our approach to build data subsets , for which the pseudo code is presented in Algorithm 1 . The algorithm is a modification of the typical AL setting , and involves the following : 1 . A labeled dataset , consisting of Nl labeled pairs , L = { ( xjl , y j l ) } Nl j=1 , where each xj ∈ X is a data point and each yj ∈ Y is its corresponding label . 2 . An acquisition model , Ma : X → Y . For our ensemble AL approach , the acquisition modelMa takes the form of a set of E different DNNs with parameters { θ ( e ) a } Ee=1 . 3 . A data subset , S = { ( xjs , yjs ) } Ns j=1 , where S is a subset of L selected using an acquisition function α ( x , Ma ) . 4 . A subset model , Ms : X → Y , with parameters θs , trained on S. As detailed in Algorithm 1 , we decouple the data selection and final optimization of AL into two different models . The acquisition model selects a subset of data using ensemble uncertainty estimation , which is then used to optimize the parameters of the subset model . 2.1 INITIALIZATION SCHEMES . We consider three different initialization schemes for the parameters of the acquisition and subset models : pretrain , compress and build up . The pretrain scheme uses the entire dataset L for pretraining both the acquisition and subset models . During optimization , the subset model is then finetuned on the data subset S. In the compress scheme , the acquisition model is pretrained on L but the subset model is randomly initialized and trained from scratch on S. The acquisition model therefore accesses all the data and then ’ compresses ’ the dataset for the subset model . For the pretrain and compress schemes , the subset S is initialized with an empty set . Finally , in the build up scheme , we follow an iterative AL loop . Specifically , we start by initializing S with a randomly selected subset of the data to train an acquisition model . After performing acquisition , instead of training a single subset model , we optimize an ensemble of E networks . This ensemble is used as an acquisition model for a subsequent iteration . Our goal is to finally reach a subset of Ns samples . As observed by Chitta et al . ( 2018a ) , exponentially growing the dataset size offers practical benefits in an AL loop setting . We therefore follow this approach , by initializing S with Ns8 random samples , and iterating two further times at Ns 4 and Ns 2 samples before obtaining a final subset of size Ns . 2.2 ACQUISITION FUNCTIONS . In our experiments , we empirically evaluate four acquisition functions of the form α : x → R for AL– entropy , mutual information , variation ratios and error count . We choose these well-known acquisition functions to maintain simplicity and scalability . For a detailed theoretical analysis of these acquisition functions , we refer the reader to Gal ( 2016 ) . Bayesian AL approaches argue that the acquisition function must target samples with high model uncertainty rather than data uncertainty . Intuitively , the reasoning behind this is that model uncertainty is a result of not knowing the correct parameters , which can be reduced by adding the right data and retraining the model . However , data uncertainty exists even for a model with the most optimal parameters , and can not be reduced by adding more data ( Kendall & Gal , 2017 ) . In the case of classification , the predictive uncertainty for a sample P ( y = k|x∗ , L ) from Eq . 1 is a multinomial distribution , which can be represented as a vector of probabilities p over each of the K classes . We can obtain the predictive uncertainty for a sample as its entropy ( Shannon , 1948 ) : H ( p ) = −pT log p. ( 3 ) However , once we marginalize out θ from Eq . 1 it is impossible to tell if this predictive uncertainty is a result of model uncertainty or data uncertainty . One acquisition function that explicitly looks for large disagreement between the models ( i.e. , model uncertainty ) , is mutual information ( Houlsby et al. , 2011 ; Smith & Gal , 2018 ) , also called BALD and Jensen-Shannon Divergence : J ( p ) ︸ ︷︷ ︸ Model = H ( p ) ︸ ︷︷ ︸ Predictive − 1 E ∑ e∈E H ( p ( e ) ) ︸ ︷︷ ︸ Data . ( 4 ) where p ( e ) denotes the prediction of an individual member of the ensemble before marginalization . Since entropy is always positive , the maximum possible value for J ( p ) isH ( p ) . However , when the models make similar predictions , 1E ∑E e=1H ( pe ) → H ( p ) , and J ( p ) → 0 , which is its minimum value . This shows that J encourages samples with high disagreement to be selected during the data acquisition process . An alternate way to look at the metric is that from the predictive uncertainty , we subtract away the expected data uncertainty , leaving an approximate of the model uncertainty ( Depeweg et al. , 2017 ) . Variation ratios ( Gal , 2016 ) is another acquisition function that looks for disagreement . It is defined as the fraction of members in the ensemble that do not agree with the majority vote M = Mode e∈E ( argmax k∈K p ( e ) k ) : V ( p ) = 1− 1 E ∑ e∈E ( argmax k∈K p ( e ) k =M ) , ( 5 ) where K is the number of classes . This is the simplest quantitative measure of variation , and prior applications in literature show that it works well in practice ( Gal et al. , 2017 ; Beluch et al. , 2018 ; Chitta et al. , 2018b ) . Our work focuses on classification due to its prevalence in the large-scale setting , but we note that V shares similarities to variance and standard deviation , which are used as acquisition functions in the Bayesian setting for regression ( Tsymbalov et al. , 2018 ) . Our final acquisition function , error count , is similar to variation ratios , but checks for disagreement with the ground truth label y rather than the mode of predictions M : E ( p ) = 1− 1 E ∑ e∈E ( argmax k∈K p ( e ) k = y ) . ( 6 ) This is typically not used in AL experiments as it can not be computed without the ground truth labels . In our setting , where the labels are available , we use this function as a baseline which , in principle , prioritizes mistakes made by the network when selecting data .

Summary: This work makes use of uncertainty estimation methods from active learning to select a subset of training data that produces models with similar (or better) performance compared to models trained on the full training set. It proposes a way to improve the Monte Carlo estimation of model uncertainty by including multiple checkpoints that are generated "for free" during a training run, thereby increasing the number of samples from 5-10 in previous work to 100 in this work. It compares several initialization schemes for the subset model using mutual information as the acquisition function, finds that a "build-up" approach (based on Chitta et. al 2018a) works best, and uses that for the rest of the studies. It then compares several acquisition functions, using the build-up approach, finds that variation ratio performs best, and uses that for the rest of the studies. Next, it compares the Top-1 accuracy on ImageNet obtained by evaluating the ensemble models produced by different ensembling schemes, and finds that ensembling 20 checkpoints from 5 training runs with different random seeds work best. Then, it uses acquisition models that use ensembles from each ensembling scheme to select subsets of the ImageNet data to be used for training the subset model, and then compares the performance of the subset models. Finally, it demonstrates this method of selecting a subset of the training data works even if the subset is used to train a model with a different architecture from the acquisition model.