Deep Active Learning by Leveraging Training Dynamics

1 INTRODUCTION . Training deep learning ( DL ) models usually requires a large amount of high-quality labeled data ( Zhang et al. , 2017 ) to optimize a model with a massive number of parameters . The acquisition of such annotated data is usually time-consuming and expensive , making it unaffordable in the fields that require high domain expertise . A promising approach for minimizing the labeling effort is active learning ( AL ) , which aims to identify and label the maximally informative samples , so that a high-performing classifier can be trained with minimal labeling effort ( Settles , 2009 ) . Under classical statistical learning settings , theories of active learning have been extensively studied from the perspective of VC dimension ( Hanneke et al. , 2014 ) . As a result , a variety of methods have been proposed , such as ( i ) the version-space-based approach , which requires maintaining a set of models ( Cohn et al. , 1994 ; Balcan et al. , 2009 ) , and ( ii ) the clustering-based approach , which assumes the data within the same cluster have pure labels ( Dasgupta & Hsu , 2008 ) . However , the theoretical analyses for these classical settings may not hold for over-parameterized deep neural networks where the traditional wisdom is ineffective ( Zhang et al. , 2017 ) . For example , marginbased methods select the labeling examples in the vicinity of the learned decision boundary ( Balcan et al. , 2007 ; Balcan & Long , 2013 ) . However , in the over-parameterized regime , every labeled example could potentially be near the learned decision boundary ( Karzand & Nowak , 2019 ) . As a result , theoretically , such analysis can hardly guide us to design practical active learning methods . Besides , empirically , multiple deep active learning works , borrowing observations and insights from the classical theories and methods , have been observed unable to outperform their passive learning counterparts in a few application scenarios ( Kirsch et al. , 2019 ; Ash et al. , 2020 ) . On the other hand , the analysis of neural network ’ s optimization and generalization performance has witnessed several exciting developments in recent years in terms of the deep learning theory ( Jacot et al. , 2018 ; Arora et al. , 2019b ; Lee et al. , 2019 ) . It is shown that the training dynamics of deep neural networks using gradient descent can be characterized by the Neural Tangent Kernel ( NTK ) of infinite ( Jacot et al. , 2018 ) or finite ( Hanin & Nica , 2020 ) width networks . This is further leveraged to characterize the generalization of over-parameterized networks through Rademacher complexity analysis ( Arora et al. , 2019b ; Cao & Gu , 2019 ) . We are therefore inspired to ask : How can we design a practical and generic active learning method for deep neural networks with theoretical justifications ? To answer this question , we firstly explore the connection between the model performance on testing data and the convergence speed on training data for the over-parameterized deep neural networks . Based on the NTK framework ( Jacot et al. , 2018 ; Arora et al. , 2019b ) , we theoretically show that if a deep neural network converges faster ( “ Train Faster ” ) , then it tends to have a better generalization performance ( “ Generalize Better ” ) , which matches with the existing observations ( Hardt et al. , 2016 ; Liu et al. , 2017 ; Lyle et al. , 2020 ; Ru et al. , 2020 ; Xu et al. , 2021 ) . Motivated by the aforementioned connection , we first introduce Training Dynamics , the derivative of training loss with respect to iteration , as a proxy to quantitatively describe the training process . On top of it , we formally propose our generic and theoretically-motivated deep active learning method , dynamicAL , which will query labels for a subset of unlabeled samples that maximally increases the training dynamics . In order to compute the training dynamics by merely using the unlabeled samples , we leverage two relaxations Pseudo-labeling and Subset Approximation to solve this non-trivial subset selection problem . Our relaxed approaches are capable of effectively estimating the training dynamics as well as efficiently solving the subset selection problem by reducing the complexity from O ( N b ) to O ( b ) . In theory , we coin a new term Alignment to measure the length of the label vector ’ s projection on the neural tangent kernel space . Then , we demonstrate that higher alignment usually comes with a faster convergence speed and a lower generalization bound . Furthermore , we extend the previous analysis to an active learning setting , where the i.i.d . assumption may not hold with the help of the maximum mean discrepancy ( Borgwardt et al. , 2006 ) . Finally , we show that alignment is positively correlated with our active learning goal , training dynamics , which implies that maximizing training dynamics will lead to a better generalization performance . Regarding experiments , we have empirically verified our theory by conducting extensive experiments on three datasets , CIFAR10 ( Krizhevsky et al. , 2009 ) , SVHN ( Netzer et al. , 2011 ) , and Caltech101 ( Fei-Fei et al. , 2004 ) using three types of network structure : vanilla CNN , ResNet ( He et al. , 2016 ) , and VGG ( Simonyan & Zisserman , 2015 ) . We first show that result of the subset selection problem delivered by the subset approximation is close to the global optimum solution . Furthermore , under active learning setting , our method not only outperforms the other baselines but also scales well on large deep learning models . The main contributions of our paper can be summarized as follows : • We propose a theory-driven deep active learning method , dynamicAL , inspired by the results of “ train faster , generalize better ” . To this end , we introduce the Training Dynamics , as a proxy to describe the training process and embody the inspiration . • We demonstrate that the convergence speed of training and the generalization performance is ( positively ) strongly correlated under the ultra-wide condition ; we also show that maximizing the training dynamics will lead to a lower generalization error in the scenario of active learning . • Our method is easy to implement . We conduct extensive experiments to evaluate the effectiveness of dynamicAL and empirically show that our method consistently outperforms other methods in a wide range of active learning settings . 2 BACKGROUND . 2.1 NOTATIONS . We use the random variable x ∈ X to represent input data feature and y ∈ Y as label where K is the number of classes and [ K ] : = { 1 , 2 , ... , K } . We are given a data source D with unknown distribution p ( x , y ) . We further denote the concatenation of x as X = [ x1 , x2 , ... , xM ] > and that of y as Y = [ y1 , y2 , ... , yM ] > . We consider a deep learning classifier hθ ( x ) = argmax σ ( f ( x ; θ ) ) : x→ y parameterized by θ ∈ Rp , where σ ( · ) is the softmax function and f is a neural network . Let ⊗ be the Kronecker Product and IK ∈ RK×K be an identity matrix . 2.2 ACTIVE LEARNING . The goal of active learning is to improve the learning efficiency of a model with a limited labeling budget . In this work , we consider the pool-based AL setup , where a finite data set S = { ( xl , yl ) } Ml=1 with M points are i.i.d . sampled from the p ( x , y ) as the ( initial ) labeled set . The AL model receives an unlabeled data set U sampled from p ( x ) and request labels according to p ( y|x ) for any x ∈ U in each query round . There are R rounds in total and for each round , a query set Q consisting of b unlabeled samples can be queried . The total budget size B = b×R . 2.3 NEURAL TANGENT KERNEL . The Neural Tangent Kernel ( Jacot et al. , 2018 ) has been widely applied to analyze the dynamics of neural networks . If the neural network is sufficiently wide , properly initialized , and trained by gradient descent with infinitesimal step size ( a.k.a . gradient flow ) , then the neural network is equivalent to kernel regression predictor with a deterministic kernel Θ ( · , · ) , called Neural Tangent Kernel ( NTK ) . When minimizing the mean squared error loss , at the iteration t , the dynamics of the neural network f has a closed-form expression : df ( X ; θ ( t ) ) dt = −Kt ( X , X ) ( f ( X ; θ ( t ) ) − Y ) ( 1 ) where θ ( t ) denotes the parameter of the neural network at iteration t , Kt ( X , X ) ∈ R|X |×K×|X|×K is called the empirical NTK and Ki , jt ( x , x′ ) = ∇θf i ( x ; θ ( t ) ) > ∇θf j ( x′ ; θ ( t ) ) , for two samples x , x′ ∈ X and i , j ∈ [ K ] . The time-variant kernel Kt ( · , · ) is equivalent to the ( time-invariant ) NTK with high probability , that is , if the neural network is sufficiently wide and properly initialized , then : Kt ( X , X ) = Θ ( X , X ) ⊗ IK ( 2 ) The final learned neural network at iteration t , is equivalent to the kernel regression solution with respect to the NTK ( Lee et al. , 2019 ) . For any input x and training data { X , Y } we have , f ( x ; θ ( t ) ) ≈ Θ ( x , X ) > Θ ( X , X ) −1 ( I − e−ηΘ ( X , X ) t ) Y ( 3 ) where η is the learning rate , Θ ( x , X ) is the NTK matrix between input x and all samples in training data X . 3 METHOD . In section 3.1 , we introduce the notion of training dynamics which can be used to describe the training process . Then , in section 3.2 , based on the training dynamics , we proposed dynamicAL . In section 3.3 , we discuss the connection between dynamicAL and existing deep active learning methods . 3.1 TRAINING DYNAMICS . In this section , we introduce the notion of training dynamics . The cross-entropy loss over labeled set S is defined as : L ( S ) = ∑ ( xl , yl ) ∈S ` ( f ( xl ; θ ) , yl ) = − ∑ ( xl , yl ) ∈S ∑ i∈ [ K ] yil log σ i ( f ( xl ; θ ) ) ( 4 ) where σi ( f ( x ; θ ) ) = exp ( f i ( x ; θ ) ) ∑ j exp ( f j ( x ; θ ) ) . We first analyze the dynamics of the training loss , with respect to iteration t , on one labeled sample ( derivation is in Appendix A.1 ) : ∂ ` ( f ( x ; θ ) , y ) ∂t = − ∑ i ( yi − σi ( f ( x ; θ ) ) ) ∇θf i ( x ; θ ) ∇ > t θ ( 5 ) For neural networks trained by gradient descent , if learning rate η is small , then∇tθ = θt+1 − θt = −η ∂ ∑ ( xl , yl ) ∈S ` ( f ( xl ; θ ) , yl ) ∂θ . Taking the partial derivative of the training loss with respect to the parameters , we have ( the derivation of the following equation can be found in Appendix A.2 ) : ∂ ` ( f ( x ; θ ) , y ) ∂θ = ∑ j∈ [ K ] ( σj ( f ( x ; θ ) ) − yj ) ∂f j ( x ; θ ) ∂θ ( 6 ) Therefore , we can further get the following result for the dynamics of training loss , ∂ ` ( f ( x ; θ ) , y ) ∂t = −η ∑ i ( σi ( f ( x ; θ ) ) − yi ) ∑ j ∑ ( x l ′ , y l ′ ) ∈S ∇θf i ( x ; θ ) > ∇θf j ( xl′ ; θ ) ( σj ( f ( xl′ ; θ ) ) − y j l ′ ) ( 7 ) Furthermore , we define di ( X , Y ) = σi ( f ( X ; θ ) ) − Y i and Y i is the label vector of all samples for i-th class . Then , the training dynamics ( dynamics of training loss ) over training set S is denoted by G ( S ) ∈ R : G ( S ) = −1 η ∑ ( xl , yl ) ∈S ∂ ` ( f ( xl ; θ ) , yl ) ∂t = ∑ i ∑ j di ( X , Y ) > Kij ( X , X ) dj ( X , Y ) ( 8 )

Summary: The authors address the problem of active learning in the context of deep learning. Instead of querying new examples based on the decision boundary (which in nonlinear models can be tricky or even ill-defined), as it is usually done for linear models, they rely on the train-faster--generalize-better paradigm. Thus, they propose to optimize the "training dynamics", which is the time derivative of the loss function in the ultra-wide limit. By using pseudo labeling and subset approximation, their method allows for fast selection of the examples that should be added to the dataset. They justify their active learning strategy through an analysis in the limit of very wide models, through the neural tangent kernel, which is particularly convenient since the NTK does not depend on time in the ultra-wide limit.  They study a quantity that they call "alignment" (which measures the correlation between input and output in the NTK space) and show that a higher alignment is related to better generalization (a better bound) and faster training (--> larger training dynamics). They then show that this is also true in the active-learning setting, where the data is no more iid, with the help of some empirical evidence. Finally, they test their method and compare it with other active learning strategies, which they systematically outperform.