Batch Active Learning at Scale

1 Introduction . Training highly effective models for complex tasks often hinges on the abundance of training data . Acquiring this data can easily become a bottleneck in cost , time , and computational resources . One major approach for addressing this problem is active learning , where labels for training examples are sampled selectively and adaptively to more efficiently train the desired model over several iterations . The adaptive nature of active learning algorithms , which allows for improved data-efficiency , comes at the cost of frequent retraining of the model and calling the labeling oracle . Both of these costs can be significant . For example , many modern deep networks can take days or weeks to train and require hundreds of CPU/GPU hours . At the same time , training human labelers to become proficient in potentially nuanced labeling tasks require significant investment from both the designers of the labeling task and the raters themselves . A sufficiently large set of queries should be queued in order to justify these costs . To address these overhead costs , previous works have developed algorithms for the batch active learning setting , where label requests are batched and model updates are made less frequently , reducing the number of active learning iterations . Of course , there is a trade-off , and the practical benefits of batch sampling come with the downside of less adaptivity and the risk of sampling redundant or otherwise less effective training examples within a batch . Batch active learning methods directly combat these risks in several different ways , for example , by incorporating diversity inducing regularizers or explicitly optimizing over the choice of samples within a batch to optimize some notion of information . However , as the size of datasets grows to include hundreds of thousands and even millions of labeled examples ( cf . Deng et al . [ 2009 ] , Krasin et al . [ 2017 ] , Van Horn et al . [ 2018 ] ) , we expect the active learning batch sizes to grow accordingly as well . The challenge with very large batch sizes is two-fold : first , the risks associated with reduced adaptivity continue to be compounded and , second , the batch 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . sampling algorithm must scale well with the batch size and not become a computational bottleneck itself . While previous works have evaluated batch active learning algorithms with batch-sizes of thousands of points ( e.g. , Ash et al . [ 2020 ] , Sener and Savarese [ 2018 ] ) , in this work , we consider the challenge of active learning with batch sizes one to two orders of magnitude larger . In this paper , we develop , analyze , and evaluate a batch active learning algorithm called ClusterMargin , which we show can scale to batch sizes of 100K or even 1M while still providing significantly increased label efficiency . The main idea behind Cluster-Margin is to leverage Hierarchical Agglomerative Clustering ( HAC ) to diversify batches of examples that the model is least confident on . A key benefit of this algorithm is that HAC is executed only once on the unlabeled pool of data as a preprocessing step for all the sampling iterations . At each sampling iteration , this algorithm then retrieves the clusters from HAC over a set of least confident examples and uses a round-robin scheme to sample over the clusters . The contributions of this paper are as follows : • We develop a novel active learning algorithm , Cluster-Margin , tailored to large batch sizes that are orders of magnitude larger than what have been considered in the literature . • We conduct large scale experiments using a ResNet-101 model applied to multi-label Open Images Dataset consisting of almost 10M images and 60M labels over 20K classes , to demonstrate significant improvement Cluster-Margin provides over the baselines . In the best result , we find that Cluster-Margin requires only 40 % of the labels needed by the next best method to achieve the same target performance . • To compare against latest published results , we follow their experimental settings and conduct smaller scale experiments using a VGG16 model on multiclass CIFAR10 , CIFAR100 , and SVHN datasets , and show Cluster-Margin algorithm ’ s competitive performance . • We provide an initial theoretical analysis , proving label complexity guarantees for a marginbased clustering sampler , which we then show is approximately equivalent to the ClusterMargin algorithm in specific settings . 1.1 Related Work . The remarkable progress in Deep Neural Network ( DNN ) design and deployment at scale has seen a vigorous resurgence of interest in active learning-based data acquisition for training . Among the many active learning protocols available in the literature ( pool-based , stream-based , membership query-based , etc . ) , the batch pool-based model of active learning has received the biggest attention in connection to DNN training . This is mainly due to the fact that this learning protocol corresponds to the way labels are gathered in practical large-scale data processing pipelines . Even restricting to batch pool-based active learning , the recent literature has become quite voluminous , and we can hardly do it justice here . In what follows , we briefly mention what we believe are among the most relevant papers to our work , with a special attention to scalable methods for DNN training that delivered state of the art results in recently reported experiments . In Sener and Savarese [ 2018 ] , the authors propose a CoreSet approach to enforce diversity of sampled labels on the unlabeled batch . There , the CoreSet idea was used as a way to compress the batch into a subset of representative points . No explicit notion of informativeness of the data in the batch is adopted . The authors reported an interesting experimental comparison on small-sized datasets . Yet , their Mixed Integer Programming approach to computing CoreSets becomes largely infeasible as the batch size grows and the authors suggest a 2-approximation algorithm as a solution . As we find empirically , it seems this lack of an informativeness signal , perhaps coupled with the 2-approximation , limits the effectiveness of the method in the large batch-size regime . Among the relevant papers in uncertainty sampling for batch active learning is Kirsch et al . [ 2019 ] , where the uncertainty is provided by the posterior over the model weights , and diversity over the batch is quantified by the mutual information between the batch of points and model parameters . Yet , for large batch sizes and standard acquisition functions , their method also becomes infeasible in practice ( we discuss this in more detail in the Section 3 ) .1 Another relevant recent work , and one 1More recently , i.e. , contemporaneously with this publication , a more efficient variant of BatchBALD has been proposed by Kirsch et al . [ 2021 ] . The authors propose a simple idea to turn the original BALD algorithm of Houlsby et al . [ 2011 ] into a batch active learning algorithm via a softmax function over the current uncertainties Algorithm 1 Hierarchical Agglomerative Clustering ( HAC ) with Average-Linkage Require : Set of clusters C , distance threshold , minimum cluster distance d = minA 6=B∈Cd ( A , B ) . 1 : if |C| = 1 or d > then 2 : return C. 3 : end if 4 : A , B ← Pair of distinct clusters in C which minimize d ( A , B ) = 1|A||B| ∑ a∈A , b∈B d ( a , b ) . 5 : if d ( A , B ) ≤ then 6 : C ← { A ∪B } ∪ C \ { A , B } . 7 : end if 8 : return HAC ( C , , d ( A , B ) ) . which we will compare to , is Ash et al . [ 2020 ] , where a sampling strategy for DNNs is proposed which uses k-MEANS++ seeding on the gradients of the final layer of the network in order to query labels that balance uncertainty and diversity . One potential downside to this approach is that the dimension of the gradient vector grows with the number of classes . In Wei et al . [ 2015 ] ( see also the more recent Killamsetty et al . [ 2020 ] ) , the authors propose a submodular sampling objective that trades-off model uncertainty with a diversity-inducing regularizer , such as a facility location objective . Using naive greedy optimization to solve such objectives does not immediately scale to extremely large batch sizes of hundreds of thousands or more ( such an implementation would require a linear number of function evaluations per greedy example added to the batch ) . More efficient “ lazier-than-lazy ” stochastic approximations ( e.g. , Mirzasoleiman et al . [ 2015 ] ) may be able to scale to very large batch sizes as they require only a linear number of function evaluations overall ( modulo a log ( 1/ ) factor , where is the approximation parameter ) . However , this may still be impractical if computing the marginal gain is expensive ( for example , if it grows approximately linearly with the pool size ) . Further recent works related to DNN training through batch active learning are Zhdanov [ 2019 ] , Shui et al . [ 2020 ] , Kim et al . [ 2020 ] , Ghorbani et al . [ 2021 ] . In Zhdanov [ 2019 ] , the authors trade off informativeness and diversity by adding weights to k-means clustering . The idea is similar in spirit to our proposed algorithm , though the way we sample within the clusters is very different ( see Section 2 below ) . Shui et al . [ 2020 ] proposes a unified method for both label sampling and training , and indicates an explicit informativeness-diversity trade-off in label selection . The authors model the interactive procedure in active learning as a distribution matching problem measured by the Wasserstein distance . The resulting training process gets decomposed into optimizing DNN parameters and batch query selection via alternating optimization . We note that such modifications in the training procedure are not feasible in settings where only data selection can be modified while the training routine is treated as a black-box ( very frequent in practice ) . The work of Kim et al . [ 2020 ] is based on the idea that uncertainty-based methods do not fully leverage the data distribution , while data distribution-based methods often ignore the structure of the learning task . Hence the authors propose to combine them in Variational Adversarial Active Learning method from Sinha et al . [ 2019 ] , the loss prediction module from Yoo and Kweon [ 2019 ] , and RankCGAN from Saquil et al . [ 2018 ] . Despite the good performance reported on small datasets , these techniques are not geared towards handling large batch sizes , which is the goal of our work . From this lengthy literature , we focus on the BADGE [ Ash et al. , 2020 ] and CoreSet algorithms [ Sener and Savarese , 2018 ] ( described in more detail in Section 3 ) as representative baselines to compare against since they are relatively scalable in terms of the batch size , do not require modification of the model training procedure , and have shown state-of-the-art results on several benchmarks . 2 Algorithm . In this section , we present the Cluster-Margin algorithm , whose pseudo-code is in Algorithm 2 . Throughout , we refer to the model being trained as f and denote its corresponding weights/parameters w. Given an unlabeled pool of examples X , in each sampling iteration , Cluster-Margin selects a diverse set of examples on which the model is least confident . We compute the confidence of a model in the pool . This idea yields a faster and simpler algorithm than BatchBALD , but it does not appear to explicitly model diversity , other than sampling without replacement from the corresponding Gibbs distribution . Algorithm 2 The Cluster-Margin Algorithm Require : Unlabeled pool X , neural network f , seed set size p , number of labeling iterations r , margin batch size km , target batch size kt ≤ km , HAC distance threshold . 1 : S ← ∅ , the set of labeled examples . 2 : Draw P ⊂ X ( |P | = p ) seed set examples uniformly at random and request their labels . Set S ← S ∪ P . 3 : Train f on P . 4 : Compute embeddings EX on the entire set X , using the penultimate layer of f . 5 : CX ← HAC ( EX , , 0 ) . 6 : for i = 1 , 2 , . . . , r do 7 : Si ← ∅ . 8 : Mi ← The km examples in X \ S with smallest margin scores . 9 : CMi ←Mapping of Mi onto CX . 10 : Sort CMi ascendingly by cluster size . Set C ′ Mi ← [ C1 , C2 , . . . , C|CMi | ] as the sorted array , and set j ← 1 as the index into C′Mi . 11 : while |Si| < kt do 12 : Select x , a random example in Cj . 13 : Si ← Si ∪ { x } . 14 : if j < |C′Mi | then 15 : j ← j + 1 . 16 : else 17 : j ← The index of the smallest unsaturated cluster in C′Mi . 18 : end if 19 : end while 20 : Request labels for Si and set S ← S ∪ Si . 21 : Train f on S. 22 : end for 23 : return S. on an example as the difference between the largest two predicted class probabilities , just as is done in the so-called “ margin ” uncertainty sampling variant [ Roth and Small , 2006 ] , and refer to the value as the margin score . The lower the margin score , the less confident the model is on a given example . In order to ensure diversity among the least confident examples , we cluster them using HAC with average-linkage ( Algorithm 1 ) . The batch of examples for which labels are requested is then chosen such that each cluster is represented in the batch . In the following , we describe the Cluster-Margin algorithm in detail , which can be decomposed into an initialization step , a clustering step , and a sequence of sampling steps . Initialization step . Cluster-Margin starts by selecting a seed set P of examples uniformly at random , for which labels are requested . A neural network is trained on P and each x ∈ X is then embedded with the penultimate layer of the network . Clustering step . HAC with average-linkage is run as a preprocessing step . Specifically , Algorithm 1 is run once on the embeddings of the examples in the entire pool , X , to generate clusters C. As described in Algorithm 1 , HAC with average-linkage repeatedly merges the nearest two clusters A , B , so long as the distance d ( A , B ) = 1|A||B| ∑ a∈A , b∈B d ( a , b ) ≤ , where is a predefined threshold . In order to achieve speedup in this step , we describe an approach in Appendix A.2 , where HAC is run only on the initial seed set P , and the examples in X \ P are projected onto the clusters of P . This reduces the clustering run-time from O ( |X|2 log |X| ) to O ( |P |2 log |P |+ |P ||X \ P | ) , without any observed loss in performance in our experiments . Sampling steps . In each labeling iteration , we employ a sampling step that selects a diverse set of low confidence examples . This process first selects M , a set of unlabeled examples with lowest margin scores , with |M | = km . Then , M is filtered down to a diverse set of kt examples , where kt is the target batch size per iteration ( kt ≤ km ) . Specifically , given the km lowest margin score examples , we retrieve their clusters , CM , from the clustering step and then perform the diversification process . To select a diverse set of examples from CM , we first sort CM ascendingly by cluster size . As depicted in Figure 1 , we then employ a round-robin sampling scheme whereby we iterate through the clusters in the sorted order , selecting one example at random from each cluster , and returning to the smallest unsaturated cluster once we have sampled from the largest cluster . This process is then repeated until kt examples have been selected . We sample from the smallest clusters first as they come from the sparsest areas of the embedded distribution and contain some of the most diverse points . We then leverage round-robin sampling to maximize the number of clusters represented in our final sample . Cluster-Margin is simple to implement and admits the key property of only having to run the clustering step once as preprocessing . As we will see , this is in contrast to recent active learning algorithms , such as BADGE and CoreSet , that run a diversification algorithm at each sampling iteration . HAC with average-linkage runs in time O ( n2 log n ) where n = |X| , but lends itself to massive speedup in practice from multi-threaded implementations [ Sumengen et al. , 2021 ] . At each iteration , ClusterMargin takes time O ( n log n ) to sample examples , whereas BADGE and CoreSet take time O ( ktn ) , which in the large batch size setting ( e.g. , kt = Ω ( √ n ) ) can be far more expensive in practice .

Summary: This paper proposes a large-batch active learning Clustering based algorithm, Cluster-Margin. The authors show that the algorithm is more efficient and more effective than comparable approaches (Random, Margin, CoreSet and BADGE) on several benchmark datasets (Open Images Dataset, CIFAR10, CIFAR100 and SVHN). The authors also provide initial theoretical guarantees justifying the relevance of the algorithm.