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: The paper introduces a new active learning acquisition function, Cluster-Margin, which clusters the samples with the lowest confidence scores (margin scores), and selects samples round-robin from the clusters to obtain a diverse set of low-confidence points. The experiments show that it works with very large batch sizes (100k--1M). Additionally, the paper provides an initial theoretical analysis of their algorithm.