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 points out the fact that in-depth analysis of data augmentation (DA) in the field of knowledge distillation is rarely performed. - In addition, the authors present a new DA strategy that simultaneously utilizes the data generated through the existing DA and the data generated through the stronger DA to improve the performance of the knolwedge distillation algorithm. - And it shows that additional performance improvement is possible with a learning algorithm inspired by active learning. - Through this, the authors show that it is possible to effectively improve performance even with the DA technique, which has not been well fused in the past. - The proposed augmentation techniques effectively converge with existing knowledge distillation algorithms to show that they can lead to further performance improvements, and as a result, state-of-the-art performance is achieved in various datasets and various teacher-student pairs.