KDExplainer: A Task-oriented Attention Model for Explaining Knowledge Distillation

1 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) [Hinton et al., 2015], 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 [Hinton et al., 2015], 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 [Hinton et al., 2015] 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.

2 RelatedWork

Knowledge Distillation. Knowledge distillation has attracted increasing attention thanks to its important role in deploying deep networks to low-capacity edge devices [Yu et al., 2017]. The main idea is leveraging the dark knowledge encoded in a bulky teacher to craft a lightweight student model with performance on par with the teacher. Over the last several years, most works devote themselves to the exploration of different forms of the dark knowledge, including soft targets [Hinton et al., 2015; Chen et al., 2020], features [Romero et al., 2015; Ye et al., 2019; Shen et al., 2021; Shen et al., 2019], attention [Zagoruyko and Komodakis, 2017; Yang et al., 2020b; Liu et al., 2021], factors [Kim et al., 2018; Yang et al., 2020a], activation boundary [Byeongho Heo, 2019], and instance relationship [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; Yang et al., 2020c]. Except the performance boost, attention mechanism also provides an important way to explain the workings of neural models [Li et al., 2016; 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.

3 KDExplainer

3.1 Knowledge Distillation with Soft Targets

Vanilla KD [Hinton et al., 2015] distills the "dark knowledge" from the teacher via aligning the soft targets

where zs and zt 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 pi(z; τ ), we have

Also, LCE is the conventional cross-entropy loss, and DKL 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.

3.2 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 1. KDExplainer is designed based

𝑩𝑩ସଶ

𝑩𝑩ସଷ

𝑩𝑩ସସ

Figure 1: A conceptual illustration. Top: a conventional student network for knowledge distillation. Bottom: the proposed KDExplainer. To explain the effects of soft targets as the distillation objective, we use the KDExplainer as the student model for knowledge distillation.

on existing widely-used networks such as ResNet [He et al., 2016], Wide Residual Network [Zagoruyko and Komodakis, 2016], and VGG [Simonyan and Zisserman, 2014]. These classic DNNs have some common characteristics: thesemodels 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 Bi to denote the i-th block, and then a standard DNN can be formulated as FDNN = BC ◦ BL ◦ · · ·◦B2 ◦B1 , where symbol ◦ denotes the function composition operation. BC 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 Bi (except the first block B1 ) into Ni equal-sized sub-blocks, i.e., ˆBi = {Bi1 ,Bi2 , ...,BiNi } for any i > 1. We view each sub-block Bi 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 "taskoriented" 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 Bi j , we use a trainable parameter vector vij = [vij,1 , vij,2 , ..., vij,Ni-1 ] to learn the attention distribution over previous experts. Let output feature maps from previous blocks be {Fi-1 1 , ..., Fi-1 Ni-1 }.

At training phase, the input of sub-block Bi j is computed by

where aij,k is the attention weight of the k-th expert, aij,k = exp(vij,k /T ) ∑m exp(vij,m/T) . T is a temperature hyper-parameter shared by all attention modules. Note that the classification block BC 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 psk = [psk,0 , psk,1 ] be the probability predictions of KDExplainer for the k-th classification task, and yk = [yk,0, yk,1] be the ground truth in the form of a one-hot vector. When the proposed KDExplainer is trained with KD, the objective function

where LWCE is the weighted cross-entropy loss

w0 and w1 are the weights balancing the cost incurred by the positive and the negative samples, which alleviates the negative effects caused by imbalanced data. ˆqtk and ˆpsk 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 pt = [pt0 , pt1 , ..., ptK ] from the teacher to K two-dimensional probability vectors as follows

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.

4 Virtual AttentionMechanism

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 Ki ∈ RCioutCi in SS to denote the convolution filters in the i-th layer, where C i in and C iout 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 Fi ∈ RCi inHW , 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 Fi+1 ∈ RCioutHW and Fi+2 ∈ RCi+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 Ki = {Ki1 ,Ki2 , ...,KiM },Kij ∈ R(Ciout /M)Ci in SS for 1 ≤ j ≤ M. 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 Ki+1 = {Ki+1 1 ,Ki+1 2 , ...,Ki+1 N },Ki+1 j ∈ R(Ci+1 out /N)Ci+1 in SS for 1 ≤ j ≤ N. 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 Ki+1 j intoM groups along the input channel, such that Ki+1 j = {Ki+1 j,1 ,Ki+1 j,2 , ...,Ki+1 j,M }, Ki+1 j,l ∈ R(Ci+1 out /N)(Ci+1 in /M)SS for any 1 ≤ l ≤ M. Similar to KDExplainer, we introduce a task-oriented attention module ai+1 j = [ai+1 j,1 , ai+1 j,2 , ..., ai+1 j,M ] before each filter block Ki+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 = [vi+1 j,1 , ..., vi+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

5 Experiments

5.1 Explaining KDwith KDExplainer

Experimental settings. Experiments are conducted on CIFAR-10 and CIFAR-100 [Krizhevsky et al., 2009]. We implement the proposed KDExplainers based on four DNN architectures: ResNet18 [He et al., 2016], VGG8 [Simonyan and Zisserman, 2014], WRN-16-2, and WRN-401 [Zagoruyko and Komodakis, 2016]. For all models, we use ResNet50 as their teacher model for KD. The initial learning rate is 0.1 and decayed every 30 epochs. The training ceases at 300 epochs. For more details, please refer to supplementary materials.

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.

Experimental analysis. In Figure 4, we provide the accuracies of all the KDExplainers. 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 WRN40-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.

5.2 Improving KDwith VAM

Experimental settings. Experiments are conducted on CIFAR-10, CIFAR-100, and Tiny-ImageNet. We incorporate VAM into six widely-used DNNs as student models, including ResNet18, resnet20 1 , VGG8, WRN-16-2, WRN-401, 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-401, 4 for resnet20, 16 for ResNet18, and 10 for ShuffleNetV1. For more details, please refer to the supplementary material.

1 Fowlloing [Tian et al., 2020], ResNet and resnet represent cifarand ImageNet-style networks, respectively.

Improving vanilla KD. As VAM is motivated by vanilla KD, we first validate its effectiveness upon vanilla KD. To give a more comprehensive view of VAM, we make comparisons between six model variants. Note that we simply use VAM to denote the DNN incorporated with the proposed virtual attention module. (M1) Plain+LCE: plain DNN trained with only LCE; (M2) VAM+LCE: VAM trained with LCE; (M3) VAM+LCE+H: VAM trained with LCE and H (Eqn. 13); (M4) Plain+LCE+DKL: plain DNN trained with LCE and DKL; (M5) VAM+LCE+DKL: VAM trained with LCE and DKL; (M6) VAM+LCE+DKL+H: VAM trained with LCE, DKL and H. Experimental results are shown in Table 1. In general, DNN incorporated with VAM yields consistently superior performance to the plain DNN without VAM. For example, M5 consistently outperforms M4 in almost all our experiments, which validates its effectiveness under KD. Furthermore, when optimized with low entropy constraint H, VAM produces better performance (M3>M2, M6>M5) under almost all settings. As H is motivated by the results from KDExplainer, it indicates that the proposed KDExplainer indeed provides us with general and valuable insights into KD. Improving state-of-the-art KD. Although VAM is motivated by vanilla KD, it is also straightforward to be combined with other KD methods. Here we evaluate VAM combined with other KD methods, including FitNets [Romero et al., 2014], FT [Kim et al., 2018], SP [Tung and Mori, 2019], and CRD [Tian et al., 2020]. 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 CIFAR100 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.

6 Conclusion and FutureWork

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.

Acknowledgments

This work is funded by the National Key R&D Program of China (Grant No: 2018AAA0101503) and the Science and technology project of SGCC (State Grid Corporation of China): fundamental theory of human-in-the-loop hybridaugmented intelligence for power grid dispatch and control.