Understanding the Success of Knowledge Distillation -- A Data Augmentation Perspective

1 INTRODUCTION . Deep neural networks ( DNNs ) are the best performing machine learning method in many fields of interest ( LeCun et al. , 2015 ; Schmidhuber , 2015 ) . How to effectively train a deep network for classification has been a central topic for decades . In the past several years , efforts have mainly focused on better architecture design ( e.g. , batch normalization ( Ioffe & Szegedy , 2015 ) , residual blocks ( He et al. , 2016 ) , dense connections ( Huang et al. , 2017 ) ) and better loss functions ( e.g. , label smoothing ( Szegedy et al. , 2016 ; Müller et al. , 2019 ) , contrastive loss ( Hinton , 2002 ) , largemargin softmax ( Liu et al. , 2016 ) ) than the standard cross-entropy ( CE ) loss . Knowledge distillation ( KD ) ( Hinton et al. , 2014 ) is a training framework that falls in the second group . In KD , a stronger network – called teacher – is introduced to guide the learning of the original network – called student – by minimizing the discrepancy between the representations of the two networks , LKD = ( 1− α ) LCE ( y , p ( s ) ) + ατ2DKL ( p ( t ) /τ , p ( s ) /τ ) , ( 1 ) where DKL represents KL divergence ( Kullback , 1997 ) ; α ∈ ( 0 , 1 ) is a factor to balance the two loss terms ; LCE denotes the cross-entropy loss ; y is the one-hot label and p ( t ) , p ( s ) stands for the teacher ’ s output and student ’ s output , respectively ( which are probability distributions over the classes ) ; τ is a temperature constant ( Hinton et al. , 2014 ) to smooth predicted probabilities . KD allows us to train smaller , more efficient neural networks without compromising on accuracy , which facilitates deploying deep learning in resource constrained environments ( e.g. , on mobile devices ) . The effectiveness of KD has been seen in many tasks ( Chen et al. , 2017 ; Wang et al. , 2020 ; Jiao et al. , 2019 ; Wang & Yoon , 2021 ) . Meanwhile , many works have investigated the reason behind its success , such as class similarity structure ( Hinton et al. , 2014 ) and regularization ( Yuan et al. , 2020 ) . However , few works have paid attention to its interplay with the input image data augmentation ( DA ) , a technique to obtain more data through various transformations ( Shorten & Khoshgoftaar , 2019 ) . In this paper , we will show that data augmentation is also an important dimension to explain the success of KD . Moreover , our findings show we can achieve much better performance simply using the original KD loss equipped with a stronger data augmentation scheme . Our proposed algorithms are inspired from interesting observations shown in Fig . 1 , where we plot the student test error curves when the model is trained for different numbers of epochs using KD loss vs. CE loss1 . Three data augmentation scenarios are examined : not using DA at all ( Without DA ) , only using the horizontal flip ( Flip ) , and using both the horizontal flip and random crop ( Flip+Crop ) . We have the following observations . ( 1 ) Within each plot , KD loss delivers lower test error than CE loss . ( 2 ) When DA is used ( comparing the middle or right plot to the left ) , both CE and KD curves are improved . ( 3 ) When DA is used ( comparing the middle or right plot to the left ) , the optimal number of training epochs is postponed and the postponement is greater for KD than CE ( the optimal number of epochs is postponed from 180 to 480 for KD versu from 60 to 120 for CE ) . ( 4 ) When a stronger DA is employed ( comparing the right plot to the middle ) , the optimal number of epochs is further postponed with even lower test error . The first two obser- vations are well-recognized by existing works ( they simply reiterate the effectiveness of KD and DA , respectively ) , while the last two observations are new discoveries of this work , concerning the interaction between KD and DA . In other words , KD and DA , as two common techniques to improve the performance of DNNs , are actually not independent . This paper explains the interplay between KD and DA and leverages it for stronger KD performance using only the standard KD loss . Specifically , we explain why KD is able to exploit DA more than the CE loss . Owing to the random transformations in data augmentation , the input data are not fixed over different epochs . Different views of each image are presented over the training process . When KD loss is used , the teacher maps these different views to different targets . As illustrated in Fig . 2 ( a ) , these targets with different probability structures can reveal more information of the data , thus helping the student more . In contrast , when CE loss is adopted , the target is fixed regardless of the different views of the input . The extra information is thus lost . This observation inspired us to develop two stronger data augmentation techniques ( TLmixup , TLCutmix ) that are tailored for KD . We further tap into the idea of active learning to make TLCutmix even better , that is , our TLCutmix+pick method . In summary , these are the main contributions of this work : • We make novel observations of the interaction between KD and DA . We explain how DA methods are more suited to be exploited by KD which uses teacher outputs as labels instead of the ground-truth one-hot labels used by CE loss . • Inspired by the above , we propose to enhance the original KD loss with stronger data augmentation schemes ( by adapting mixup ( Zhang et al. , 2018 ) and CutMix ( Yun et al. , 2019 ) to KD ) . It is shown that these methods are more reasonably applied in the KD case than in the CE case . • We further propose an even stronger data augmentation method specifically for knowledge distillation using the idea of active learning ( Settles , 2011 ) . • We show empirically better results simply by using the original KD loss combined with the proposed DA scheme , compared to state-of-the-art KD methods , which adopt more advanced distillation losses . 1Note , data points on one curve in Fig . 1 are from independent experiments of different total number of epochs ; the learning rate schedule is proportionally scaled based on the total number of epochs . 2 RELATED WORK . Knowledge distillation : The general idea of knowledge distillation is to guide the training of a student model through a stronger pretrained teacher model ( or an ensemble of models ) . It was pioneered by Buciluǎ et al . ( Buciluǎ et al. , 2006 ) and later refined by Hinton et al . ( Hinton et al. , 2014 ) , who coined the term . Since its debut , knowledge distillation has seen extensive application in vision and language tasks ( Chen et al. , 2017 ; Wang et al. , 2020 ; Jiao et al. , 2019 ; Wang & Yoon , 2021 ) . Many variants have been proposed regarding the central question in knowledge distillation , that is , how to define the knowledge that is supposed to be transferred from the teacher to the student . Examples of such knowledge definitions include feature distance ( Romero et al. , 2015 ) , feature map attention ( Zagoruyko & Komodakis , 2017 ) , feature distribution ( Passalis & Tefas , 2018 ) , activation boundary ( Heo et al. , 2019 ) , inter-sample distance relationship ( Park et al. , 2019 ; Peng et al. , 2019 ; Liu et al. , 2019 ; Tung & Mori , 2019 ) , and mutual information ( Tian et al. , 2020 ) . Over the past several years , the progress has been made primarily at the output end ( i.e. , through a better loss function ) . In contrast to previous works , our goal in this paper is to improve the KD performance at the input end with the help of data augmentation . We will show this path is as effective and also has much potential for future work . Data augmentation : Deep neural networks are prone to overfitting , i.e. , building input-target connection using undesirable or irrelevant features ( like noise ) in the data . Data augmentation is a prevailing technique to curb overfitting ( Shorten & Khoshgoftaar , 2019 ) . In classification tasks , data augmentation aims to explicitly provide data with label-invariant transformations ( such as random crop , horizontal flip , color jittering ) in the training so that the model can learn representations robust to those nuisance factors . Recently , some advanced data augmentation methods were proposed , which not only transform the input , but also transform the target based on certain corresponding relations . For example , mixup ( Zhang et al. , 2018 ) linearly mixes two images with the labels mixed by the same linear interpolation ; manifold mixup ( Verma et al. , 2019 ) is similar to mixup but conducts the mix operation in the feature level instead of pixel level ; CutMix ( Yun et al. , 2019 ) pastes a patch cut from an image onto another image with the label decided by the area ratio of the two parts . Now that the input and target are transformed simultaneously , the key is to maintain a semantic correspondence between the new input and new target . Although these methods have been proven effective , one lingering concern is about the reasoning behind them . Specifically , it is easy to come up with examples where the semantic correspondence is poorly kept ( see Fig . 5 for examples on CutMix ) . Unlike these methods , which focus on general classification using the cross-entropy loss , our work investigates the interplay between data augmentation and knowledge distillation loss and the proposed new data augmentation is specifically for knowledge distillation . One recent work ( Das et al. , 2020 ) also conducts empirical study about the impact of data augmentation on knowledge distillation . However , their exploration is very different from ours : they first apply data augmentation ( e.g. , mixup/CutMix ) to the teacher then conduct the distillation step as usual ( no extra data augmentation in this step ) ; our investigation is the exact opposite to their setup : we train the teacher as usual ( not applying mixup/CutMix ) , then in the distillation step we employ a more advanced data augmentation ( e.g. , mixup/CutMix ) . Interestingly , they conclude that the teacher trained with mixup/CutMix hurts the student ’ s generalization ability , while we consistently see a performance boost by using a stronger DA in the distillation step during student training . 3 PROPOSED METHOD . 3.1 INTERPLAY OF KD , DA , AND TRAINING ITERATIONS . We first introduce a framework to explain the phenomenon that KD can exploit DA more simply by training for more iterations . Training for more iterations means presenting more examples to the network . Over iterative training , the presented examples are not exactly the same among different epochs because of the random transformations of data augmentation . Different versions of an image produced by data augmentation can be regarded as multi-views of that image ( Wu et al. , 2018 ; Tian et al. , 2020 ) . We term this kind of data difference as input view diversity . As depicted in Fig . 2 ( a ) , when the CE loss is employed , different views of an image are mapped to a single point in the target space ( the hard label ) . In contrast , when the KD loss is used , different views of the data are mapped to a group of points in the target space through the teacher , which can reveal richer information around that class . By richer , we specifically mean two sources . First is the class structure information provided by the soft labels instead of the hard labels , i.e. , the well-known dark knowledge ( Hinton et al. , 2014 ) . Second is the information provided by the different input views from data augmentation . Concretely , in Fig . 2 ( a ) , although the three views of the input share the same main class “ dog ” , the target probability vectors are different : compared with the first view , the second one has more lawn in it , thus the teacher has larger predicted probability in the “ lawn ” class ; similarly , the third view has more sunshine , thus larger probability in the “ sunset ” class . These subtle changes in class-wise probabilities are beneficial for the student to learn . However , if the CE loss is used , all the three views are mapped to the same one-hot label , not putting the extra information to good use . More training iterations keep producing new data views to the student , making it afford more training epochs without overfitting . In contrast , for CE , as the target is fixed for different views , little new information is added in longer training . Thus , the student can only afford fewer extra epochs . Denote the optimal number of training epochs asN , lowest test error asE . The synergistic interplay between KD and DA can be summarized in the following hypotheses , N ( DA ) KD −N ( w/o DA ) KD > N ( DA ) CE −N ( w/o DA ) CE , N ( DA+ ) KD > N ( DA ) KD and E ( DA+ ) KD < E ( DA ) KD , ( 2 ) where “ DA ” refers to a data augmentation method ; “ DA+ ” refers to another data augmentation method stronger than “ DA ” ; “ w/o DA ” means not using any data augmentation . These inequalities are empirically verified in our experiments ( Fig . 1 , Fig . 3 ) . The first inequality suggests the advantage of KD over CE : given the same DA scheme , KD loss can make more use of extra training iterations . The second suggests we can obtain better accuracy simply by training for more epochs using a stronger DA method . This leads us to investigate stronger DA methods for KD , as follows .

Summary: This paper aims to explore the advantage of using DA for KD for training deep neural networks and try to explain the interplay between DA and KD regarding the label advantage.   The main idea delivered in the paper regarding the interplay is that instead of fixing labels for augmented samples when using CE, the teacher in KD can act on different augmented to provide the better label to train the student. The paper also comes up with the hypothesis of the training interactions between KD and DA (e.g. KD loss makes use of more training iterations and the “stronger” DA method leads to better accuracy if training more interactions. The experiments are conducted with DA methods of Flip, Flip + Crop, TLmixup, TLCutMix. The paper also proposes TLCutMix+pick, which improves the baseline TLCutMix with an active learning approach.   Experiments conducted on CIFAR-100, Tiny ImageNet and ImageNet show the improvements of using DA over the baselines and the improvement of the proposed TLCutMix+pick.