Computer Science
Computer Science-2021.gz:143183
Computer Science
f89307abc6a06127e8d1cc4f7fbbc8d427e5ff2b
KDExplainer: A Task-oriented Attention Model for Explaining Knowledge Distillation

Knowledge distillation (KD) has recently emerged as an efficacious scheme for learning compact deep neural networks (DNNs). Despite the promising results achieved, the rationale that interprets the behavior of KD has yet remained largely understudied. In this paper, we introduce a novel task-oriented attention model, termed as KDExplainer, to shed light on the working mechanism underlying the vanilla KD. At the heart of KDExplainer is a Hierarchical Mixture of Experts (HME), in which a multi-class classification is reformulated as a multi-task binary one. Through distilling knowledge from a free-form pre-trained DNN to KDExplainer, we observe that KD implicitly modulates the knowledge conflicts between different subtasks, and in reality has much more to offer than label smoothing. Based on such findings, we further introduce a portable tool, dubbed as virtual attention module (VAM), that can be seamlessly integrated with various DNNs to enhance their performance under KD. Experimental results demonstrate that with a negligible additional cost, student models equipped with VAM consistently outperform their non-VAM counterparts across different benchmarks. Furthermore, when combined with other KD methods, VAM remains competent in promoting results, even though it is only motivated by vanilla KD. The code is available at https://github.com/zju-vipa/KDExplainer.


Introduction
Recent progress in deep neural networks (DNNs) has significantly benefited many if not all tasks in artificial intelligence, ranging from computer vision to natural language processing. The encouraging results, however, come with an enormous cost: state-of-the-art DNNs have been scaled up to hundreds and even thousands of layers with millions of parameters, making it demanding to train and deploy even on GPU clusters, let alone on edge devices like mobile phones. * Corresponding author Many research efforts have thus been made to craft lightweight DNNs applicable to in-the-wild scenarios. Representative schemes include weight pruning [Han et al., 2015], model quantization [Jacob et al., 2018], and knowledge distillation (KD) , among which KD has recently emerged as one of the most flourishing topics in the field. The goal of KD is to extract knowledge from a wellbehaved but cumbersome teacher model, often known as dark knowledge, and to learn a compact student model capable to handle the task of the teacher but with fewer parameters. Since the pioneering work of , a variety of dark knowledge has been explored, including hint [Romero et al., 2015], attention [Zagoruyko and Komodakis, 2017], and instance relationship [Liu et al., 2019].
Despite the above forms of dark knowledge showcase promising results, the naive soft target  is found still among the most competitive ones [Tian et al., 2020]. Nevertheless, few attempts have been dedicated to explaining the rationale of soft targets on the student learning. A common belief is that soft labels reveal richer information like category similarities than the widely-used one-hot vector, so that the student obtains more supervision signals for learning. Recently, the work [Yuan et al., 2020] argues that the soft target is intrinsically a type of learned label smoothing [Szegedy et al., 2016] regularization. [Furlanello et al., 2018] conjectures that soft targets resemble importanceweighting, where weights correspond to confidences of the teachers in the correct prediction. [Tang et al., 2020], on the other hand, dissects the effects of soft targets into three main factors: label smoothing, importance weighting, and category similarities. Albeit the inspiring insights provided by these works, the underlying mechanism of how category similarities regularize the student model learning has remained largely under-studied to date.
In this paper, we take a closer look at the role of the soft label through the lens of a novel attention model, which we term as KDExplainer. KDExplainer is an interpretationfriendly student model, which distinguishes itself from conventional student DNNs from two main aspects, as illustrated in Figure 1. First, KDExplainer takes the form of a Hierarchical Mixture of Experts (HME), where each expert is expected to specialize in specific subtasks and learn adaptive features. Second, KDExplainer casts the multi-class classification task as a multi-task problem, in which each task is formulated as a binary classification problem. Such a design explicitly shapes KDExplainer as a neural tree, for which the inherent working mechanism is well understood, in the aim to interpret KD.
We then carry out KD from a free-form pre-trained DNN to the dedicated KDExplainer, through which process the rationale of KD is highlighted thanks to the interpretable essence of the HME architecture. Interestingly, we find that the KD objective promotes lower entropy of the attention distribution, indicating that soft labels, in reality, encourage specialization of different experts and hence play a role in modulating the knowledge conflicts for solving different subtasks. To further understand the connection and difference between KD and label smoothing, we train a KDExplainer using label smoothing, and discover that the derived attention distribution exhibits no significant differences with those by vanilla training without KD. This phenomenon marks that soft labels indeed have more to offer, including feature specialization, than label smoothing.
Inspired by these observations, we further introduce a portable component, termed as virtual attention module (VAM), to enhance the performance of conventional DNNs under KD. The key idea of VAM is to coordinate the knowledge conflicts for discriminating different categories, achieved via lowering the entropy of the attention distribution. VAM can be readily integrated with existing DNNs while bringing negligible additional computation cost. Moreover, since VAM is naturally orthogonal to KD, it can be seamlessly combined with various KD schemes, rendering it a handy module to facilitate KD.
Our contributions are therefore summarized as follows.
• We propose a novel attention model of interpretable nature, KDExplainer, to understand the role of soft labels in training the student.
• Through KDExplainer, we observe that soft labels implicitly modulate the knowledge conflicts between different subtasks by promoting feature specialization, and offer more regularization than only label smoothing.
• Understanding the KD rationale via KDExplainer further motivates us to design a portable and compact module, VAM, readily applicable to various DNNs and KDs.
Extensive experimental results across benchmarks demonstrate that VAM not only consistently improves the performance of vanilla KD under various experimental settings, but also can be readily integrated with other state-of-the-art KD schemes to further promote their results.  [Liu et al., 2019;Park et al., 2019;Tung and Mori, 2019]. By imitating the teacher to behave in a similar way, the student achieves comparable performance even with much fewer parameters. Attention Mechanism Inspired by human cognition, attention mechanisms focus on relevant regions of input data to solve the desired task rather than ingesting the entire input. Attention-based neural networks have been broadly adopted in natural language models for machine translation [Bahdanau et al., 2015], image caption generation [Xu et al., 2015], and unsupervised representation learning [Devlin et al., 2019]. Attention mechanisms also achieve great success in vision models [Mnih et al., 2014;Bello et al., 2019]. Except the performance boost, attention mechanism also provides an important way to explain the workings of neural models Wiegreffe and Pinter, 2019;Xu et al., 2015]. Unlike most prior works, our focus here is to utilize an attention mechanism to interpret KD, which has been largely overlooked in previous literature.

Knowledge Distillation with Soft Targets
Vanilla KD [Hinton et al., 2015] distills the "dark knowledge" from the teacher via aligning the soft targets where z s and z t are respectively the logits from the student and the teacher models. y is its associated one-hot label vector, p denotes the softmax function that produces the category probabilities given the logits, and τ is a non-negative temperature hyperparameter used to smooth the distributions. As for p i (z; τ ), we have Also, L CE is the conventional cross-entropy loss, and D KL is the Kullback-Leibler divergence between the categorial distributions predicted from the teacher and the student models. α is a hyperparameter to trade off the two objective terms.

The Proposed KDExplainer
We propose the KDExplainer, a task-oriented attention model, as the student model to uncover the rationale underlying KD. KDExplainer makes two main modifications based on exiting popular DNNs: (1) dissecting the effects of class similarity in soft targets, KDExplainer reformulates the multi-class classification problem as an ensemble of multiple binary classification problems; (2) KDExplainer remodels the student model as a Hierarchical Mixture of Experts (HME) and introduces a task-oriented attention mechanism as the gating function to make the student model more friendly for human interpretation. The overview of the proposed KDExplainer is shown in Figure  classic DNNs have some common characteristics: these models are usually composed by several blocks, each of which is a stack of convolutional layers, batch normalization layers [Ioffe and Szegedy, 2015], and nonlinear activation layers. The number of filters usually keeps fixed in the same block and changes across different blocks. Formally, we use B i to denote the i-th block, and then a standard DNN can be formulated as F DN N = B C • B L • · · ·•B 2 •B 1 , where symbol • denotes the function composition operation. B C is the classification block that consists of the fully connected layer and the softmax layer. KDExplainer roughly follows the design of the existing DNN, but divides each block B i (except the first block B 1 ) into N i equal-sized sub-blocks, i.e.,B i = {B i 1 , B i 2 , ..., B i Ni } for any i > 1. We view each sub-block B i j as an expert, and introduce a taskoriented attention module before each expert as the gating function to select a combination of the outputs from previous experts. Therefore, the whole model can be viewed as an HME model.

Task-oriented Attention
The proposed attention is "task-oriented" as the parameters of the attention module are trained on the whole dataset, which is largely different from the attention mechanism in existing literature [Bahdanau et al., 2015;Bello et al., 2019] where the attention weights are determined by instances. For each sub-block B i j , we use a trainable parameter vector to learn the attention distribution over previous experts. Let output feature maps from previous blocks be F i−1 1 , ..., F i−1 Ni−1 . At training phase, the input of sub-block B i j is computed by where a i j,k is the attention weight of the k-th expert, T is a temperature hyper-parameter shared by all attention modules. Note that the classification block B C is divided into K blocks, where K equals the number of classes in the classification problem. The multi-class classification problem thus turns into an ensemble of binary classification problems in KDExplainer, as shown in Figure 1.

Explaining KD by KDExplainer
Recall that KDExplainer is a substitute of the conventional student model, using KDExplainer to understand the effects of KD is straightforward: analyzing differences between the experimental results of KDExplainer trained with and without KD. As KDExplainer reformulates the multi-class classification as multiple binary classification tasks, each binary classification task encounters the imbalanced data problem. Let p s k = [p s k,0 , p s k,1 ] be the probability predictions of KDExplainer for the k-th classification task, and y k = [y k,0 , y k,1 ] be the ground truth in the form of a one-hot vector. When the proposed KDExplainer is trained with KD, the objective function is where L WCE is the weighted cross-entropy loss L WCE = −w 0 y k,0 log p s k,0 − w 1 y k,1 log p s k,1 .
(5) w 0 and w 1 are the weights balancing the cost incurred by the positive and the negative samples, which alleviates the negative effects caused by imbalanced data.q t k andp s k are the softened category probabilities from the teacher and the student models, respectively. Note that all teacher models involved in this paper are still conventional DNNs for multi-class classification, so we convert the K-class probability prediction  p t = [p t 0 , p t 1 , ..., p t K ] from the teacher to K two-dimensional probability vectors as followŝ , for any k ∈ {1, 2, ..., K}.
If KDEplainer is trained without KD, the second term in Eqn. 4 is removed.
KDExplainer enjoys an appealing property thanks to the elaborate design: after training, it becomes a tree-like multibranch network (i.e., neural tree) if only the maximum in the attention weights is kept: The derived neural trees provide us with stronger support for interpretation and analysis of knowledge distillation.

Virtual Attention Mechanism
KDExplainer is tailored to understand the effects of soft targets in KD. However, it may suffer from slight accuracy sacrifice compared to conventional student DNNs due to the ad hoc design that multi-class classification turns into multiple binary classification tasks. Here we propose a virtual attention mechanism that is easily integrated into existing student models with few architecture modifications, to retain the capacity of these student models for higher KD performance. As shown in Figure 2, VAM views all convolution filters in a layer as several "virtual" filter blocks, each filter block akin to the expert in the KDExplainer. Note that "virtual" means that VAM does not really split the convolution filters into any blocks. The only modification is that VAM slightly changes the conventional convolution operation by incorporating an attention mechanism.
Formally, we use K i ∈ R C i out C i in SS to denote the convolution filters in the i-th layer, where C i in and C i out denote the input and the output channels. Here for simplicity, we assume all filters are square, where both the width and the length are S. Assume the input tensor of the i-th convolution layer is F i ∈ R C i in HW , where H and W are the height and the width of the feature map. After the i-th layer and the (i + 1)-th layer, the features turn into F i+1 ∈ R C i out HW and F i+2 ∈ R C i+1 out HW , respectively. To incorporate the attention module into the DNN, we view convolution filters in each layer as a set of filter blocks. Assume the filters in the i-th layer are divided into M virtual groups, such that K After the i-th layer, the feature maps produced by each block can be calculated as follows where denotes the convolution operation. The whole feature tensor can be denoted by where denotes feature concatenation. Assume filters in the (i + 1)-th layer are divided into N virtual blocks, such that To align each filter block in the (i + 1)th layer with the output from each block of the i-th layer, we further divide each block K i+1 j into M groups along the input channel, such that K i+1 in /M )SS for any 1 ≤ l ≤ M . Similar to KDExplainer, we introduce a task-oriented attention module a i+1 j = a i+1 j,1 , a i+1 j,2 , ..., a i+1 j,M before each filter block K i+1 j , thus the output of the block can be computed by The attention weights are computed by a softmax function over trainable parameters v = v i+1 j,1 , ..., v i+1 j,M , i.e., As soft targets are found to encourage lower entropy of the attention distribution (seen in Section 5.1), we introduce another regularization term based on the conventional KD objective in Eqn. 1 where A denotes all the involved attention distributions, and H is the sum of their entropy

Experimental results
To understand the effects of vanilla KD, KDExplainer is trained with and without soft targets as the optimization objective. Furthermore, as label smoothing [Szegedy et al., 2016] is recently also viewed as a type of KD [Yuan et al., 2020], we also train the KDExplainer with label smoothing to understand its effects. In Figure 3, we visualize the derived neural trees by keeping only the maximum attention weight in every attention module as Eqn. 7. It can be seen that KD significantly encourages sparse connectivity between blocks, as its retained blocks are much fewer than normal training and label smoothing. It can be also verified by the attention distribution shown in the histograms in internal nodes, where KD produces sharper (i.e., lower entropy) attention distribution than normal training and label smoothing. Furthermore, KD produces more human-interpretable branching architecture than normal training and label smoothing. For example, in the derived trees, man-made vehicle categories are attached to the same branch, while animal categories are attached to other branches. This is a reasonable organization as similar categories or tasks share similar decision patterns, and thus they can share the same network branch. Less related categories or tasks, on the other hand, depend on different patterns to make their decision, such that they should be solved separately. The results of normal training and label smoothing somewhat violate this principle, thus their derived trees are much larger than that of KD.
In Figure 4, we provide the accuracies of all the KD-Explainers. Under all experimental settings, KD significantly outperforms label smoothing and vanilla training. Label smoothing marginally improves the performance in most cases compared to vanilla training, but it sometimes leads to performance degradation. For example, it causes performance drops by 0.51% and 0.14% of WRN-40-1 and WRN-16-2 on CIFAR-100. These results prove that by properly organizing categories into different branches, KD helps modulate the knowledge conflicts between different tasks and thus achieves higher accuracy. Label smoothing, on the other hand, provides some regularization on the model learning, but it seems to have play no role in addressing conflicts between decision patterns of different categories. In the supplementary material, we provide more results and analyses to explain the differences between KD and label smoothing.

Improving KD with VAM
Experimental settings Experiments are conducted on CIFAR-10, CIFAR-100, and Tiny-ImageNet [Le and Yang, 2015]. We incorporate VAM into six widely-used DNNs as student models, including ResNet18, resnet20 1 , VGG8, WRN-16-2, WRN-40-1, and ShuffleNetV1 [Zhang et al., 2018]. For all involved student models, we adopt ResNet50 as their teacher model to provide the soft logits. During the training phase of the student model, the learning rate is initially 0.05 (attention module 0.01), gets decayed at epoch 150, 180, 210 by a factor of 0.1, and ceases at 240. The temperature hyper-parameter is set to 1 for all attention modules and 4 for KD loss. The trading-off factor α is set to 0.9. The number of channels in each virtual block is 8 for VGG8, WRN-16-2, and WRN-40-1, 4 for resnet20, 16 for ResNet18, and 10 for ShuffleNetV1. For more details, please refer to the supplementary material.   [Kim et al., 2018b], SP [Tung and Mori, 2019], and CRD [Tian et al., 2019]. Results are listed in Table 2. It can be seen that combined with other KD methods, VAM still yields performances consistently superior to the plain DNN though it is motivated by only vanilla KD. For resnet20, VAM achieves 0.20% ∼ 1.18% performance boosts on CIFAR-100 in four KD methods. For WRN-16-2, VAM improves 0.18% ∼ 0.81% on CIFAR-100. Since VAM brings negligible additional overhead, it is an economical way to further improve the performance of existing KD methods. Please refer to the supplementary material for more experimental results, including how the block size and the hyper-parameter γ affect the final performance.

Conclusion and Future Work
In this paper, we propose KDExplainer to shed light on the working mechanism underlying soft targets during KD. We find that KD implicitly modulates the knowledge conflicts between different subtasks, and effectively brings about more benefits as compared to label smoothing. Based on these observations, we propose a portable module, VAM, to further improve the results of vanilla KD alongside other state-ofthe-art ones. Extensive experimental results exhibit that the proposed VAM significantly enhances the KD performance at a negligible additional cost. In our future work, we will extend the proposed VAM to other tasks and systematically evaluate its effectiveness beyond the scope of KD.