Knowledge distillation via softmax regression representation learning

1 INTRODUCTION . Recently , there has been a great amount of research effort to make Convolutional Neural Networks ( CNNs ) lightweight so that they can be deployed in devices with limited resources . To this end , several approaches for model compression have been proposed , including network pruning ( Han et al. , 2016 ; Lebedev & Lempitsky , 2016 ) , network quantization ( Rastegari et al. , 2016 ; Wu et al. , 2016 ) , knowledge transfer/distillation ( Hinton et al. , 2015 ; Zagoruyko & Komodakis , 2017 ) , and neural architecture search ( Zoph & Le , 2017 ; Liu et al. , 2018 ) . Knowledge distillation ( Buciluǎ et al. , 2006 ; Hinton et al. , 2015 ) aims to transfer knowledge from one network ( the so-called “ teacher ” ) to another ( the so-called “ student ” ) . Typically , the teacher is a high-capacity model capable of achieving high accuracy , while the student is a compact model with much fewer parameters , thus also requiring much less computation . The goal of knowledge distillation is to use the teacher to improve the training of the student and push its accuracy closer to that of the teacher . The rationale behind knowledge distillation can be explained from an optimization perspective : there is evidence that high capacity models ( i.e . the teacher ) can find good local minima due to over-parameterization ( Du & Lee , 2018 ; Soltanolkotabi et al. , 2018 ) . In knowledge distillation , such models are used to facilitate the optimization of lower capacity models ( i.e . the student ) during training . For example , in the seminal work of ( Hinton et al. , 2015 ) , the softmax outputs of the teacher provide extra supervisory signals of inter-class similarities which facilitate the training of the student . In other influential works , intermediate representations extracted from the teacher such as feature tensors ( Romero et al. , 2015 ) or attention maps ( Zagoruyko & Komodakis , 2017 ) have been used to define auxiliary loss functions used in the optimization of the student . Training a network whose output feature representation is rich and powerful has been shown crucial for achieving high accuracy for the subsequent classification task in recent works in both unsupervised and supervised learning , see for example ( Chen et al. , 2020 ; He et al. , 2020 ) and ( Kang et al. , 2020 ) . Hence , in this paper , we are advocating for representation learning-based knowledge distillation by optimizing the student ’ s penultimate layer output feature . If we are able to do this effectively , we expect ( and show experimentally ) to end up with a student network which can generalize better than one trained with logit matching as in the KD paper of ( Hinton et al. , 2015 ) . Main contributions : To accomplish the aforementioned goal we propose two loss functions : The first loss function , akin to ( Romero et al. , 2015 ; Zagoruyko & Komodakis , 2017 ) , is based on direct feature matching but focuses on optimizing the student ’ s penultimate layer feature only . Because direct feature matching might be difficult due to the lower representation capacity of the student and , more importantly , is detached from the classification task at hand , we also propose a second loss function : we propose to decouple representation learning and classification and utilize the teacher ’ s pre-trained classifier to train the student ’ s penultimate layer feature . In particular , for the same input image , we wish the teacher ’ s and student ’ s feature to produce the same output when passed through the teacher ’ s classifier , which is achieved with a simple L2 loss ( see Fig . 1 ) . This softmax regression projection is used to retain from the student ’ s feature the information that is relevant to classification , but since the projection matrix is pre-trained ( learned during the teacher ’ s training phase ) this does not compromise the representational power of the student ’ s feature . Main results : Our method has two advantages : ( 1 ) It is simple and straightforward to implement . ( 2 ) It consistently outperforms state-of-the-art methods over a large set of experimental settings including different ( a ) network architectures ( WideResNets , ResNets , MobileNets ) , ( b ) teacherstudent capacities , ( c ) datasets ( CIFAR-10/100 , ImageNet ) , and ( d ) domains ( real-to-binary ) . 2 RELATED WORK . Knowledge transfer : In the work of ( Hinton et al. , 2015 ) , knowledge is defined as the teacher ’ s outputs after the final softmax layer . The softmax outputs carry richer information than one-hot labels because they provide extra supervision signals in terms of the inter-class similarities learned by the teacher . In a similar fashion to ( Hinton et al. , 2015 ) , intermediate representations extracted from the teacher such as feature tensors ( Romero et al. , 2015 ) or attention maps ( Zagoruyko & Komodakis , 2017 ) have been used to define loss functions used to facilitate the optimization of the student . Trying to match the whole feature tensor , as in FitNets ( Romero et al. , 2015 ) , is hard and , in certain circumstances , such an approach may adversely affect the performance and convergence of the student . To relax the assumption of FitNet , Attention Transfer ( AT ) was proposed in ( Zagoruyko & Komodakis , 2017 ) where knowledge takes the form of attention maps which are summaries of the energies of the feature tensors over the channel dimension . An extension of ( Zagoruyko & Komodakis , 2017 ) using Maximum Mean Discrepancy of the network activations as a loss term for distillation was proposed in ( Huang & Wang , 2017 ) . Cho & Hariharan ( 2019 ) showed that very accurate networks are “ too good ” to be good teachers and proposed to mitigate this with early stopping of the teacher ’ s training . Recently , the work of ( Heo et al. , 2019a ) studied the location within the network at which feature distillation should be applied and proposed margin ReLU and a specifically designed distance function that transfers only the useful ( positive ) information from the teacher to the student . More recently , Li et al . ( Li et al. , 2020a ) proposed to supervise the blockwise architecture search by the architecture knowledge distilled from a teacher model . Another NAS based method was proposed in ( Guan et al. , 2020 ) , in which a student-to-teacher loss is used to find the aggregation weights that match the learning ability of the student . Passalis et al . ( 2020 ) claimed that traditional KD ignores information plasticity during the training process , and proposed to model the information flow through the various layers of the teacher . Feature relationship transfer : Another line of knowledge distillation methods focus on exploring transferring the relationship between features , rather than the actual features themselves . In ( Yim et al. , 2017 ) , feature correlations are captured by computing the Gram matrix of features across layers for both teacher and student and then applying an L2 loss on pairs of teacher-student Gram matrices . The limitation of this work is the high computational cost , which is addressed to some extent in ( Lee et al. , 2018 ) by compressing the feature maps by singular value decomposition . Park et al . ( 2019 ) proposed a relational knowledge distillation method which computes distance-wise and angle-wise relations of each embedded feature vector . This idea is further explored in ( Peng et al. , 2019 ) and ( Liu et al. , 2019 ) . In ( Peng et al. , 2019 ) , Taylor series expansion is proposed to better capture the correlation between multiple instances . In ( Liu et al. , 2019 ) , the instance feature and relationships are considered as vertexes and edges respectively in a graph and instance relationship graph is proposed to model the feature space transformation across layers . Inspired by the observation that semantically similar inputs should have similar activation patterns , ( Tung & Mori , 2019 ) proposed a similarity-preserving knowledge distillation method which guides the student to mimic the teacher with respect to generating similar or dissimilar activations . More recently , ( Jain et al. , 2020 ) proposed to matching the student output with the teacher ’ s by distilling the knowledge through a quantized visual words space . Li et al . ( 2020b ) proposed the local correlation exploration framework to represent the relationships of local regions in the feature space which contains more details and discriminative patterns . Finally , a similar connection between distillation and representation learning was very recently made in ( Tian et al. , 2020 ) which uses contrastive learning for knowledge distillation . We note that our loss is not related to the one used in ( Tian et al. , 2020 ) , is simpler , and as shown in Section 5 , outperforms it for all of our experiments , often by a significant margin . 3 METHOD . We denote by T and S the teacher and student networks respectively . We split these networks into two parts : ( i ) A convolutional feature extractor fNet , Net = { T , S } , the output of which at the i-th layer is a feature tensor F iNet ∈ RC i Net×H i×W i , where CiNet is the output feature dimensionality , and Hi , W i the output spatial dimensions . We also denote by hNet = ∑HL h=1 ∑WL w=1 F L Net ∈ RC L Net the last layer feature representation learned by fNet . ( ii ) A projection matrix WNet ∈ RC L Net×K which projects the feature representation hNet into K class logits ziNet , i = 1 , . . . , K , followed by the softmax function s ( ziNet ) = exp ( ziNet/τ ) ∑ j exp ( z j Net/τ ) with temperature τ ( τ = 1 for Cross Entropy loss ) which put together form a softmax regression classifier into K classes . Knowledge Distillation ( KD ) ( Hinton et al. , 2015 ) trains the student with the following loss : LKD = − K∑ k=1 s ( zkT ) log s ( z k S ) , ( 1 ) so that the discrepancy between the teacher ’ s and student ’ s classifiers is directly minimized . FitNets ( Romero et al. , 2015 ) match intermediate feature representations . For the i-th layer , the following loss is defined : LFit = ∥∥F iT − r ( F iS ) ∥∥2 , ( 2 ) where r ( . ) is a function for matching the feature tensor dimensions . In our work , we propose to minimize the discrepancy between the representations hT and hS . To accomplish this goal , we propose to use two losses . The first one is an L2 feature matching loss : LFM = ‖hT − hS‖2 , ( 3 ) where for notational simplicity we dropped the dependency on r ( . ) . Hence , LFM loss is a simplified FitNet loss which focuses only on the final representation learned . The intuition for this is that this feature is directly connected to the classifier and hence imposing the student ’ s feature to be similar to that of the teacher could have more impact on classification accuracy . Moreover , it might be questionable why one should optimize for other intermediate representations as in ( Romero et al. , 2015 ) especially when the student is a network of lower representational capacity . In Section 4 : Where should the losses be applied ? , we confirm that LFM alone has a positive impact but feature matching in other layers is not helpful . We found LFM to be effective but only to limited extent . One disadvantage of LFM and , in general , of all feature matching losses e.g . ( Romero et al. , 2015 ; Zagoruyko & Komodakis , 2017 ) , is that it treats each channel dimension in the feature space independently , and ignores the inter-channel dependencies of the feature representations hS and hT for the final classification . This is in contrast to the original logit matching loss proposed by Hinton et al . in ( Hinton et al. , 2015 ) which directly targets classification accuracy . To alleviate the aforementioned problem , in this work , we propose a second loss for optimizing hS which is directly linked with classification accuracy . To this end , we will use the teacher ’ s pre-trained Softmax Regression ( SR ) classifier . Let us denote by p the output of the teacher network when fed with some input image x . Let us also feed the same image through the student network to obtain feature hS ( x ) . Finally let us pass hS ( x ) through the teacher ’ s SR classifier to obtain output q . See also Fig . 1 . Our loss is defined as : LSR = −p log q . ( 4 ) At this point , we make the following two observations : ( 1 ) If p = q ( and since the teacher ’ s classifier is frozen ) , then this implies that hS ( x ) = hT ( x ) which shows that indeed Eq . ( 4 ) optimizes the student ’ s feature representation hS ( hT is also frozen ) . ( 2 ) The loss of Eq . ( 4 ) can be written as : LSR = −s ( W ′ThT ) log s ( W ′ThS ) . ( 5 ) Now let us now write KD loss in a similar way : LKD = −s ( W ′ThT ) log s ( W ′shS ) . ( 6 ) By comparing Eq . ( 5 ) with Eq . ( 6 ) , we see that the only difference in our method is that the frozen , pre-trained teacher ’ s classifier is used for both teacher and the student . On the contrary , in KD , WS is also optimized . This gives more degrees of freedom to the optimization algorithm , in particular , to adjust the weights of both the student ’ s feature extractor fS and the student ’ s classifier WS in order to minimize the loss . This has an impact on the learning of the student ’ s feature representation hS which , in turn , hinders the generalization capability of the student on the test set . We confirm this hypothesis with the experiment of Section 4 : Transferability of representations . Finally , we note that we found that , in practice , an L2 loss between the logits : LKD = ‖W ′ThT −W ′ThS‖ 2 = ‖hT − hS‖2WT , ( 7 ) works slightly better than the cross-entropy loss . The comparison between different types of losses for LSR is given in the appendix . Overall , in our method , we train the student network using three losses : L = LCE + αLFM + βLSR , ( 8 ) where α and β are the weights used to scale the losses . The teacher network is pretrained and fixed during training the student . LCE is the standard loss based on ground truth labels for the task in hand ( e.g . cross-entropy loss for image classification ) . Note that this results in a very simple algorithm for training the student , summarized in Algorithm 1 . Algorithm 1 Knowledge distillation via Softmax Regression Representation Learning Input : Teacher network T , Student network S , input image x , ground truth label y , weights α , β . 1 . Input x to S to obtain feature hS and class prediction ŷ . Calculate cross entropy loss LCE = H ( ŷ , y ) ; 2 . Input x to T to obtain feature hT . Calculate distillation losses from Eqs . ( 3 ) and ( 7 ) ; 3 . Update S by optimizing Eq . ( 8 ) Output : the updated S .

Summary: This paper proposes a new knowledge distillation method by enhancing the student network's representation learning. The proposed framework is very simple: use the teacher's final fully-connected layer (or, 'projection matrix $W$' in the paper) to obtain both the teacher's and the student's logit output and optimize the student network by minimizing the gap between two logits. Experiments cover various network architectures on CIFAR and ImageNet datasets and show improvements over the other competitive knowledge distillation methods.