Comprehensive Knowledge Distillation with Causal Intervention
Year: 2022

Abstract:
Knowledge distillation (KD) addresses model compression by distilling knowledge from a large model (teacher) to a smaller one (student). The existing distillation approaches mainly focus on using different criteria to align the sample representations learned by the student and the teacher, while they fail to transfer the class representations. Good class representations can benefit the sample representation learning by shaping the sample representation distribution. On the other hand, the existing approaches enforce the student to fully imitate the teacher while ignoring the fact that the teacher is typically not perfect. Although the teacher has learned rich and powerful representations, it also contains unignorable bias knowledge which is usually induced by the context prior (e.g., background) in the training data. To address these two issues, in this paper, we propose comprehensive, interventional distillation (CID) that captures both sample and class representations from the teacher while removing the bias with causal intervention. Different from the existing literature that uses the softened logits of the teacher as the training targets, CID considers the softened logits as the context information of an image, which is further used to remove the biased knowledge based on causal inference. Keeping the good representations while removing the bad bias enables CID to have a better generalization ability on test data and a better transferability across different datasets against the existing state-of-the-art approaches, which is demonstrated by extensive experiments on several benchmark datasets1.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Comprehensive, Interventional Distillation
text: 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).

heading: 3.1 Comprehensive Representation Distillation
text: 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 cannot 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.

heading: 3.2 Interventional Distillation
text: Although a teacher has learned good representations, it is typically not perfect. Comprehensive distillation enables the student to inherit the superior representations from the teacher while it also introduces the bias to the student. To address this issue, we use causal intervention to remove the bias.

heading: 3.2.1 Structural Causal Model
text: In knowledge distillation, the pretrained teacher with the context information in training data can be considered as the prior knowledge for training the student. We illustrate the causalities among prior knowledge K, sample X , and prediction Y in Figure 4(a), where A→ B denotes that A is the causer of B. We describe the causal relationships among these variables in the following.
K →X: the context prior in K determines where the object appears in an image X , e.g., the context prior in the training dataset in Figure 1 puts the dog object in green grasses instead of rooms.
K → J ← X: J is the context-based representation of X by using the context bases in K. This relationship exists due to the fact that even for the same image, its context representation under different dataset contexts or with different pretrained teachers can differ substantially.
X → Y ← J: Besides the regular X → Y , the prediction is also affected by the prior knowledge K through mediation J . For example, in Figure 1, the cats in the test dataset are misclassified to dogs, since the context prior in K misleads the model to focus on the grass feature in X .
Therefore, the prior knowledge K is a confounder of X and Y . The existing approaches that directly learn P (Y |X) from the teacher bring the bias to the student model. We propose to model P ((Y |do(X)) with causal intervention [33] to remove the bias.

heading: 3.2.2 Interventional Distillation via Backdoor Adjustment
text: After determining the cause-effect relationships, we use causal intervention P (Y |do(X)) instead of P (Y |X) as the classifier, which pursues the true causality from X to Y by removing the effects of confounder K. Physical intervention, i.e, collecting samples with objects in all possible contexts evenly, is impossible [57]. Thanks to backdoor adjustment, we can model P (Y |do(X)) by cutting off K → X , which is achieved by stratifying the confounder into pieces K = {k1, k2, ..., k|K|}, so that K is not a confounder of X and Y anymore as shown in Figure 4(b). The de-confounded student is expressed as:
P (Y |do(X)) = |K|∑ i=1 [P (Y |X, J = g(X, ki))P (ki)] (4)
where g() is a function which we define later for generating context representation J from X and ki.
As there are m classes which can be considered as m different context items [54], we set each item ki of the prior knowledge to a class ci, i.e., K = {ci}mi=1. The m context base vectors are set to the class centers. Since different samples in a class have different probabilities of containing the object ci, we use the the weighted average of sample features as the class center.
c̄i = ∑k j=1 P (ci|xj)hj∑k j=1 P (ci|xj)
(5)
where k is the total number of samples in class ci. P (ci|xj) is set to the teacher learned probability. With context base vectors, we define the sample-specific context representation J as a linear combination of the context base vectors. As the logits learned by the teacher contain sample-to-class similarities, we use the softened logits to approximate the context coefficients to provide context information. The coefficient learned by the teacher for sample X on base c̄i is written as ati = σ( T (X) τ )[i], where τ is temperature to soften the logits and σ is the softmax function. J can thus be expressed as: J = g(X, ci) = aic̄i. Since the teacher has learned appropriate context information, we enforce the student to learn the context information from the teacher, which leads to the final interventional distillation objective:
Linv = P (Y |do(X)) +K(asi , ati) = m∑ i=1 [P (Y |fs(X)&(asi c̄i))P (ci)] +K(asi , ati) (6)
where & denotes the concatenation operation; K is a metric to force the student to learn the context information from teacher. The learned context information is then used in the first term (for causal intervention) to make the final prediction through a linear classifier. we set K to KL-divergence in this paper. We simply set the weight for K() to 1 as we find that it works very well. P (ci) is set to the percentage of the samples in class ci, e.g., in balanced datasets, P (ci) = 1m . We summarize the interventional distillation in Figure 5.
The ∑
operation in (6) makes the forward cost of the final linear classifier linearly increase with the number of classes. This issue can be addressed by adopting the normalized weighted geometric mean [50] as an approximation:
Linv ≈ P (Y |fs(X)& m∑ i=1 [P (ci)asi c̄i]) +K(asi , ati) (7)
The complete objective of CID. CID aims to achieve comprehensive knowledge distillation while removing the bad bias with causal intervention. Thus, its final objective is written as:
LCID = Linv + αM(hTs W,ht) + β m∑ i=1 ||CSi − CTi ||2 (8)

heading: 4 Experiments
text: In this section, we first conduct ablation studies and then compare CID with SOTA approaches.

heading: 4.1 Ablation Studies
text: The ablation studies are conducted on CIFAR-100 by using WRN-40-2 and WRN-16-2 as the teacher and the student, respectively.
Effects of different components of CID. We use w/o SR, w/o CR, and w/o INV to denote CID without the sample representation distillation, without the class representation distillation, and without the intervention P (Y |do(X)) by using the regular P (Y |X), respectively. As shown in Figure 6, the performances drop significantly without any of these terms. Specifically, as expected, sample representations play the most important role with performance gain 1.51%, since sample representations are directly involved in making final predictions and are further used in class representation distillation and causal intervention for removing biases. The improvement of the class representation distillation is 0.65%, which suggests that good class presentations benefit the sample representation learning and thus benefit the performance. On the other hand, by using the proposed P (Y |do(X)) instead of P (Y |X) to pursue the true causality from X to Y , the student obtains 0.70% improvement, which demonstrates the effectiveness and necessity of the interventional distillation.
Effects of NM_MSE. The proposed NM_MSE is modified from MSE by using a personalized weight for each sample to remove the feature norm biases of the teacher so that each sample contributes equally to the student model. Despite its simplicity, it is observed in Figure 6 that NM_MSE (i.e., CID) substantially outperforms MSE (i.e., CIDmse), which demonstrates the superiority of NM_MSE.
Which layer’s representations should be transferred? CID transfers the last layer’s feature vectors of the teacher to the student with the motivation that these features are directly involved in making the final prediction. We check the effects of distilling the representations in different layers. We report the results in Figure 7. It is observed that the sample representations in the last layer are more effective than those in the other layers, and even better than the combination of the representations in all the layers, The reason can be that as the representations in the intermediate layers are not directly used for the final prediction, enforcing the student to imitate these representations hurts the learning ability and flexibility of the student which has a small capacity.

heading: 4.2 Comparison Settings with SOTA Approaches
text: We compare CID with SOTA approaches across varieties of (a) benchmark datasets (i.e., CIFAR-10 [24], CIFAR-100 [24], Tiny ImageNet 2, and ImageNet [11]), (b) network architectures (i.e., ResNet
2https://tiny-imagenet.herokuapp.com
[16], WRN [55], VGG [41], and MobileNet [39]), (c) settings (i.e., the teacher and the student share the architecture or use different architectures), (d) cases (i.e., regular cases, across-dataset cases, and data-limited cases). The competitors include FitNet [37], AT [56], SP [46], CC [35], PKT [31], AB [18], VID [2], RKD [30], CRD [45], SRRL [21], and CSKD [7]. Since CID uses the logits of teacher to provide context information for each sample, for a fair comparison, the KD [19] loss is added to all the competitors. We omit "+KD" for simplicity when denoting these competitors plus KD, e.g., "FitNet+KD" is abbreviated to "FitNet". On the other hand, since CID does intervention with the assistance of a linear layer, we also try to add an extra linear layer to the competitors, but we find that it hurts their performances due to overfitting, which we report in Appendix.

heading: 4.3 Model Compression
text: CIFAR-10: Table 1 reports the comparison results on CIFAR-10. We adopt the teacher and the student with the same architecture or different architectures. It is observed that CID consistently outperforms all the baselines significantly across different architectures on both settings, while there is no baseline consistently as the second best, since these baselines show their advantages in different architectures. These results demonstrate the superiority of CID.
CIFAR-100: We further report the comparison results on CIFAR-100 in Table 2. For a fair comparison, we adopt the architectures from the SOTA approaches (CRD [45] and SRRL [21]). As shown in Table 2, for compressing the large models to the smaller ones, CID obtains the best performances in different settings, which demonstrates the effectiveness of CID for model compression. The superior performances of CID are due to its ability to distill comprehensive knowledge and remove biases.
Tiny ImageNet: We further evaluate CID in more challenging datatset Tiny ImageNet. Table 3 shows that CID beats all the SOTA approaches substantially in terms of both Top-1 and Top-5 accuracies on the challenging dataset, which validates the usefulness and superiority of CID on different datasets.
ImageNet: To investigate the performance of CID on large scale datasets, we follow CRD by adopting ImageNet and using ResNet-34 and ResNet-18 as the teacher and the student, respectively. As shown
in Table 4, CID outperforms these competitors significantly, which demonstrates the applicability and effectivness of CID on large scale datasets.

heading: 4.4 Transferability Comparison
text: An important goal of representation learning is to learn general representations which can be transferred to different datasets. We investigate the across-dataset generalization ability of CID. For a fair comparison, we follow the settings of CRD. Specifically, we freeze the feature encoder of the student and train a linear classifier on STL-10 [9] or TinyImageNet. WRN-40-2 and WRN-16-2 are adpoted as the teacher and the student, respectively.
The transferability comparison results are reported in Table 5. It is clearly observed that CID beats the prior work substantially on both datasets, which demonstrates its superior generalization ability on new data. The reason is that when transferring the knowledge from one dataset to another, the inherited bias from the teacher can be a disaster to the new dataset. The ability of CID to remove the biased knowledge mitigates this issue, thus leading to a better generalization on new datasets.

heading: 4.5 Data-Limited Distillation Performances
text: In reality, it happens that when a powerful model is released, only a few data samples are publicly accessible due to the privacy or confidentiality issues in various domains such as medical and industrial
domains. It is thus necessary for distillation approaches to work on these practical cases. We compare CID with the existing approaches in the data-limited scenario on CIFAR-100 by using VGG-13 and VGG-8 as the teacher and the student, respectively. As shown in Table 6, CID outperforms all the baselines by a large margin in all the three cases with 20%, 40%, and 60% training data. We also notice that the advantage of CID is more obvious when fewer training data are available, e.g., the improvement of CID over the second best method is about 2% on 20% or 40% training data cases, which is much higher than the improvement on 60% training data. The reason is that when fewer data samples are available, these samples are severely inadequate to represent the real data distribution so that the biases become more serious. While the existing approaches fail to handle this issue, CID is able to address it with the interventional distillation, which leads to a better performance.

heading: 5 Conclusion, Limitations, and Broader Impact
text: Conclusion. In this paper, we have proposed comprehensive, interventional distillation (CID) that captures both sample and class representations while removing the bias by using softened logits as the context information based on causal intervention. To our best knowledge, CID is the first framework along the line of using causal inference to address KD-based model compression. To this end, CID is able to keep the good representations and remove the bad bias. Extensive experiments demonstrate that CID has a better generalization ability on test data and a better transferability across different datasets against the existing SOTA approaches.
Limitations. A major assumption in CID is that the training data used by the teacher and the student are from the same distribution. The assumption is typically satisfied in knowledge distillation literature as almost all the existing work uses the same data to train the teacher and the student. On the other hand, when the assumption is violated, new biases will be introduced from the new data. CID is not designed to solve this problem and we leave this question to the future work. Also, when the training data used by the student and the teacher differ substantially, the teacher may not be able to supervise the student anymore. The role of the teacher needs to be changed in this case, which we leave for the future work.
Broader Impact. There is an increasing interest in implementing DNNs on portable devices such as smart phones and watches, while DNNs need a large amount of memory and computation, which highly limits their deployments on these resource-limited devices. CID can be used to address this issue by compressing large models (teachers) to small and fast ones (students). The advantage of CID over the other distillation approaches is that it not only enables the students to inherit comprehensive knowledge from the teachers but also removes the bad biased knowledge, which leads to a better generalization and transferability. More essentially, in real world, collecting data is very expensive, while using sparse data points to train a student induces severe biases, which poses challenges to the existing distillation approaches. CID is able to address this problem with the interventional distillation. So far, no negative impact has been observed.

Review summaries:
This paper focuses on comprehensive knowledge distillation to transfer the class representation. The proposed method aims to improve the performance of the student model by considering the transferring of class representation. It is a novel view, especially with a causal intervention approach. The experiments are conducted to demonstrate the advantages of the proposed method.

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. 

Authors have proposed a novel framework (CID) for knowledge distillation that is able to capture sample and class representations and reduce knowledge bias induced by the context prior through causal intervention. There are two major differences between the proposed approach and previous work: first, CID is able to transfer class representations. Second, it uses soften logits as sample context information. Authors have provided extensive experimental results to show the effectiveness of their approach.

The paper addresses an interesting concern : how can we transfer knowledge from a large model (master) to a tiny one (student) taking into account possible bias in the learning process of the teacher The framework considered is that of  image classification, where the context contained is the image can yield some bias in the classification task (for e.g. garden when distinguishing a dog from a cat) The paper considered two kinds of knowledge transfer * one may want to transfer representation of each class * one may also want to transfer the classifier in itself  The approach proposed in this paper to circumvent the bias problem is based on causal  intervention.  The claim is that there is an underlying graph where the output of the classifier may have as parent in the graph, both the true representation of the object to identify and the context The idea is to perform causal intervention to remove the influence of the context in the classification Numerical experiments are performed in the image context. The task is object clasiffication and several classical datasets are considered