Comprehensive Knowledge Distillation with Causal Intervention

1 Introduction . The superior performances of deep neural networks ( DNNs ) are accompanied with large amounts of memory and computation requirements , which seriously restricts their deployment on resourcelimited devices . An effective and widely used solution to this issue is knowledge distillation [ 19 , 37 ] that compresses a large network ( teacher ) to a compact and fast network ( student ) by knowledge transfer . To this end , the student obtains a significant performance boost . The original knowledge distillation ( KD ) [ 19 ] uses the softened logits generated by a teacher as the targets to train a student . Ever since then , substantial efforts including [ 37 , 45 ] have been made on aligning the sample representations learned by the student with those learned by the teacher using different criteria . However , almost all the existing approaches [ 45 , 49 , 58 ] have overlooked the class representations . Good class representations are beneficial to sample representation learning , since they can shape the sample representation distribution . To address this issue , we propose comprehensive distillation to incorporate the class representations learned by the teacher into the distillation process . On the other hand , as the teacher has learned rich and powerful representations , the existing approaches enforce the student to fully mimic the behavior of the teacher . However , fully imitating 1Code : https : //github.com/Xiang-Deng-DL/CID 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the representations of the teacher may not be optimal , since the bias contained in the teacher is also transferred to the student . The bias is usually caused by the context prior in the training data . For example , as shown in Figure 1 , the dogs in the training dataset are usually on green grasses and the cats are in a room , which misleads the trained classifier to classify the cats on green grasses in the test dataset as dogs and the dogs in a room as cats due to the bias induced by the context . Similar cases can also happen on the attributes of data samples , e.g. , when the colors of the dogs in the training dataset are mostly black , the black cats in the test dataset may be wrongly classified as dogs . Transferring this kind of the bias contained in the pretrained teacher to the student hurts the student . Since the biased knowledge in the teacher is caused by the training data , we assume that the training data used by the teacher and those used by the student are from the same distribution . This is not a strong assumption in knowledge distillation literature as almost all the existing work uses the same dataset when training a teacher and a student , which obviously satisfies the assumption . Contrary to this assumption , when the training data for the teacher and the student are from different distributions , two issues arise . First , the teacher may not be able to teach the student anymore due to the data distribution discrepancy . Second , new biases will be introduced in the distillation process from the new training dataset . We leave these questions for the future work . Under the above assumption , we formulate the causal relationships [ 32 ] among the pretrained teacher , the samples , and the prediction in a causal graph as shown in Figure 4 ( a ) . More details are given in Section 3.2.1 . We then use the softened logits learned by the teacher as the context information of an image to remove the biased knowledge based on backdoor adjustment [ 14 ] . To this end , we propose a simple yet effective framework ( i.e. , CID ) to achieve comprehensive distillation and bias removal . We summarize our contributions and the differences from the existing approaches as follows : • We propose a novel knowledge distillation framework , i.e. , CID , which captures comprehensive representations from the teacher while removing the bias with causal intervention . To our best knowledge , this is the first work to study how to use causal inference to address KD-based model compression . • CID is different from the existing approaches in two aspects . First , CID is able to transfer the class representations which are largely ignored by the existing literature . Second , CID uses softened logits as sample context information to remove biases with causal intervention , which differs from the existing literature that uses the softened logits as the training targets to train a student . Keeping the good knowledge while removing the bad bias enables CID to have a better generalization on test data and a better transferability on new datasets . • Extensive experiments on several benchmark datasets demonstrate that CID outperforms the state-of-the-art approaches significantly in terms of generalization and transferability . 2 Related Work . Knowledge Distillation . Hinton et al . [ 19 ] propose the original KD that trains a student by using the softened logits of a teacher as targets . Compared to one-hot labels , the logits provide extra information learned by the teacher [ 19 , 13 ] . However , KD fails to transfer the powerful representations learned by the teacher . Ever since then , many efforts have been made on aligning the sample representations learned by a student and a teacher . FitNet [ 37 ] aligns the sample representations learned by a student with those learned by a teacher through regressions . AT [ 56 ] distills sample feature attention from a teacher into a student . CRD [ 45 ] maximizes the mutual information between sample representations learned by a student and a teacher . SRRL [ 21 ] aligns the sample representations of a teacher and a student by using the teacher ’ s classifier . CC [ 35 ] and SP [ 46 ] transfer the sample correlation over the whole dataset to the student , which may contain redundant and irrelevant information as pointed out in [ 7 ] . Similarly , other approaches [ 52 , 20 , 22 , 25 , 42 , 48 , 17 , 8 , 2 , 23 , 1 , 40 , 45 , 49 , 21 , 12 ] use different criteria to align the sample representations . We notice that almost all these approaches only transfer the sample representations while largely ignoring the class representations which can benefit the sample representation learning by shaping the sample representation distribution . Causal Inference . Causal inference [ 33 , 34 , 38 ] aims to explore the cause-effect relationships between different variables . It can not only be used to interpret a particular phenomenon [ 6 , 26 ] , but also serve as a tool to address problems by determining and using the causal effects [ 3 , 5 , 29 ] . Recently , it has been introduced to machine learning [ 4 ] and has been used in different applications , including but not limited to domain adaptation [ 15 , 27 ] , imitation learning [ 10 ] , image captioning [ 51 ] , scene graph generation [ 44 ] , visaul dialog [ 36 ] , few-shot learning [ 54 ] , imbalance classification [ 43 ] , semantic segmentation [ 53 , 57 ] , VQA [ 28 ] , and unsupervised learning [ 47 ] . In this work , we provide an interventional framework for knowledge distillation to remove the biased knowledge in the teacher . 3 Comprehensive , Interventional Distillation . In this section , we first describe the comprehensive distillation which takes both the sample and class representations into account . We then present interventional distillation to remove the biased knowledge with causal intervention and thus achieve comprehensive , interventional distillation ( CID ) . 3.1 Comprehensive Representation Distillation . CID considers both sample and class representations and thus achieves comprehensive distillation . Which layer ’ s sample representations are transferred ? Many approaches [ 37 , 45 ] transfer the intermediate or the last few layers ’ feature representations . In contrast , CID only distills the feature vectors in the last layer ( before the classifier ) , since only these vectors are directly involved in making the final prediction . We empirically validate this point in Section 4.1 . The criterion to align representations . We slightly modify the mean square error ( MSE ) as the criterion to align the sample representations . MSE is used in FitNet [ 37 ] for transferring the sample representations . However , we find that MSE has a disadvantage that it biases towards the samples that have large-norm features . For example , for MSE ( Va , Vb ) =||Va − Vb||2 where Va and Vb are the variable and target , respectively , when Va = [ 0.02 , 0.02 ] and Vb = [ 0.01 , −0.01 ] , the gradient is only 2 ( Va − Vb ) = [ 0.02 , 0.06 ] . It is observed that Va and Vb have totally different directions and are orthogonal , but the gradient is very small due to the small norms of the two vectors . In contrast , for large vectors Va = [ 10 , 20 ] and Vb = [ 9.5 , 19.5 ] , the gradient is [ 1 , 1 ] which is much larger than [ 0.02 , 0.06 ] , even if the two vectors are very close in terms of both directions and norms . The inherent disadvantage of MSE makes the samples with different feature norms contribute differently to the student , which induces biases . As shown in Figure 2 , the norms of the sample features learned by the teacher are across a wide range and have a noticeable variance . One natural idea to address this issue is to use the MSE of the normalized vectors , i.e. , || Va||Va||2 − Vb ||Vb||2 || 2 . However , this loss only aligns the directions of the two vectors and the minimum point is not necessarily Va = Vb anymore , e.g. , Va = [ 1 , 1 ] and Vb = [ 100 , 100 ] are a solution to this loss due to the same vector direction although they are different substantially . To address this issue , we propose Normalized MSE ( NM_MSE ) : M ( Va , Vb ) = MSE ( Va , Vb ) ||Vb||2 = ||Va − Vb||2 ||Vb||2 ( 1 ) where ||.||2 denotes the square of L2 norm . NM_MSE can be considered as a sample-wise weighted MSE and the weights 1||Vb||2 for different samples are negatively related to their target feature vector norms , which mitigates the bias and makes different samples equally contribute to the student . Although this is a slight modification to the original MSE , we empirically find that it performs much better than MSE on benchmark datasets . Class representations . The existing approaches only enable the student to capture the sample representations from the teacher while overlooking the class representations . CID incorporates this part into the distillation process . CID uses the class shapes to represent the class representations . As shown in Figure 3 , the shape of a class is a graph with all the samples in the class as the nodes . The weight between two nodes is defined as the representation similarity between the two nodes . We adopt the cosine similarity and thus the class representation of class i is expressed as : Ci = [ hi1 ||hi1||2 , hi2 ||hi2||2 , .. , hik ||hik||2 ] T [ hi1 ||hi1||2 , hi2 ||hi2||2 , .. , hik ||hik||2 ] ( 2 ) where hij denotes the feature vector of sample j in class i ; k is the total number of samples in class i ; superscript T means transpose ; ||.||2 denotes L2 norm . The class representation is different from the sample-class relation defined in [ 7 ] which only captures sample-to-class-center similarity while failing to capture sample-to-sample relations and thus can not well represent a class . The objective for comprehensive representation distillation . CID transfers comprehensive knowledge consisting of both sample and class representations from a teacher to a student , and thus the objective for comprehensive distillation is written as : Lrep ( S , T ) = αM ( hTs W , ht ) + β m∑ i=1 ||CSi − CTi ||2 ( 3 ) whereM is the proposed NM_MSE ; hs ∈ Rms and ht ∈ Rmt are the sample representations learned by student S and teacher T , respectively ; W ∈ Rms×mt is a linear transformation for converting hs to the space with dimension mt ; CSi and C T i are the ith class representations learned by S and T , respectively ; m is the total number of classes ; α and β are two balancing weights .

Summary: This paper proposes an improved knowledge distillation process named Comprehensive, Interventional Distillation (CID). It adds a class representation term to the loss function of a knowledge distillation model to enforce the student and the teacher to learn not only similar sample representations but also class representations, where a class is represented by the similarity matrix among all of its samples. It further modifies the loss function term of the student to remove the bias learned by the teacher from the context information in training data, based on casual intervention. Experimental results on four real datasets confirm the effectiveness of the proposed model in improving the classification accuracy of the student model.