IMF: Integrating Matched Features Using Attentive Logit in Knowledge Distillation

1 Introduction

Although deep learning methods have been shown to achieve great performance in various tasks over computer vision and natural language processing [Dosovitskiy et al., 2020; Brown et al., 2020; Ramesh et al., 2021], they generally require large datasets, model parameters, and extensive computations. Therefore, deploying complex models on mobile devices is one of the major challenges due to the device's limited computational capabilities. Thus, many researchers have proposed approaches to overcome such limitations. In particular, knowledge distillation (KD) [Hinton et al., 2015] has been proposed to effectively transfer the knowledge of a large model ("teacher") to a smaller model ("student").

∗Corresponding author

The existing KD methods can be categorized into one of two groups: 1) softmax representation distillation [Hinton et al., 2015; Park et al., 2019; Passalis and Tefas, 2018; Tian et al., 2019; Yang et al., 2021; Lee et al., 2022] and 2) latent (intermediate) feature matching [Romero et al., 2014; Zagoruyko and Komodakis, 2016; Yim et al., 2017; Tung and Mori, 2019]. The first approaches have explored how to generate and transfer better final representations of the model by minimizing the difference of the softmax representation between the teacher and student model using KL divergence proposed by Hinton et al [Hinton et al., 2015]. On the other hand, the second approach has focused on improving latent feature map matching by mimicking the intermediate representation that works as a strong indicator to help learn the terminal representation [Gou et al., 2021].

Most KD approaches have focused on enhancing the performance of the student model without changing its underlying structure. However, such approaches inevitably have to select the matching student model that satisfies the accuracy, inference time, and a number of model parameter requirements. To reduce the inference time, one can choose the small models for convenience. However, due to the unsatisfactory performance from severe capacity gap or model structures [Cho and Hariharan, 2019], it is needed to re-select and retrain the student, which can be inefficient and cumbersome. Moreover, not much research has been conducted in the area of KD to explore the performance vs. computational cost trade-off, focusing on the flexible and adaptive student model architecture. To alleviate the issues above, we introduce the flexible and adaptive KD architecture to improve performance while maintaining efficiency.

In this paper, we propose Integrating Matched Features using Attentive Logit in Knowledge Distillation (IMF), to improve the student performance by combining the outputs from each intermediate feature distiller (IFD) that the teacher model trains. Moreover, IMF is designed to be an adaptive architecture so that the student model can be flexibly configured to improve the performance, while lowering the complexity via the different number of distilled IFDs in branch networks. In particular, each IFD takes the student's intermediate feature map as input and produces its latent representation by exploiting the element-wise normalized attentive (ENA) layer to learn and mimic the teacher's knowledge better for the higher performance, achieving higher performance. Moreover, combined outputs of IFD are also trained to mimic the teacher to boost the performance further. Hence, with the smaller final student's model capacity compared to the original student model, IMF can achieve higher performance. Through extensive experiments, we demonstrate that our approach outperforms the best performing state-of-the-art baselines, and shows robustness in distilling knowledge from various teacher models to the student models.

2 RelatedWork

2.1 Knowledge Distillation

Hinton et al. [Hinton et al., 2015] first proposed KD to boost the performance of a student network by transferring the knowledge from a teacher network, showing that the student can learn not only the hard-label information but also the softlabel information from the softmax output of the teacher. On the other hand, feature map-based approaches [Romero et al., 2014; Zagoruyko and Komodakis, 2016] are proposed, to reduce the distance between the respective feature maps of a teacher and student. In terms of direct matching features between teacher and student network, relational knowledge distillation (RKD) [Park et al., 2019] considered the input examples as mutual relations and transferred these relations instead of distilling softmax representation. Furthermore, contrastive representation distillation (CRD) was proposed by Tian et al. [Tian et al., 2019], which utilizes contrastive learning for the student to capture distance-based information (data point) from the teacher's representation with an additional layer during training; however, in a regression task, CRD cannot be easily applied because positive and negative pairs are difficult to define and obtain.

Recently, Yang et al. [Yang et al., 2021] devised a novel distillation via softmax regression representation learning (SR) to enable direct feature matching at the penultimate layers of the teacher and student. First, they decoupled the latent feature and the softmax classifier. Then, for successful representation learning, they employed a frozen teacher classifier to train the student's penultimate layer feature in the teacher's classifier, and they achieved state-of-the-art performance. Also, weight soft label distillation (WSL) was proposed by Zhou et al. [Zhou et al., 2021] to apply a biasvariance trade-off to the KD method. They observed that the bias-variance trade-off varies sample-wisely during training with soft label information, which affects the negative results in KD. Especially, they filtered out larger trade-off samples and achieved higher performance. Our approach greatly differs from the above methods as we effectively combine intermediate features with attentive logits from features of student's blocks and inference only with those blocks.

2.2 Branch Network

Before the advent of skip connections [He et al., 2016], Inception [Szegedy et al., 2015] was the state-of-the-art method to achieve the best performance on ImageNet [Deng et al., 2009] by using auxiliary classifiers. In particular, an auxiliary classifier can encourage discrimination in class labels in the shallow layers and effectively provide regularization to intermediate features, directly injecting cross entropy loss to the intermediate classifiers. Not only are branch networks used to enhance classification performance, but also Teerapittayanon et al. proposed using branch networks for an early exit in the inference process [Teerapittayanon et al., 2016]. Based on an early layer of a CNN network which is sufficient for classifying many data points, they achieved compressed inference time with comparable performance in several CNN models.

Self-distillation [Zhang et al., 2019], one of the newly proposed KD concepts as mentioned, enhanced the performance of DNNs without the supervision of the teacher by transferring the knowledge of deep layers to shallow layers via branch networks; however, each branch network, which is much smaller than teacher models, cannot provide refined information because of its limited student's capacity [Ji et al., 2021]. Therefore, its performance enhancement is much less compared to the KD methods using teacher networks. The early exit and aforementioned methods used branch networks for performance improvement and inference time reduction, where we similarly adapt branch networks as feature distillers.

However, in our methods, we first propose a novel framework, in which the branch networks densely distill the logit of the teacher to the branch networks, and takes the intermediate multi-scale feature maps of the student network as inputs. Moreover, the branch networks are used as predictors without extra computation, while discarding the existing classifier and the last MLP block in the inference phase. By changing the forwarding path, the IMF flexibly controls the model size to have the same or even less number of parameters than that of the original student networks, while achieving higher performance. Exploiting several branch networks, we propose a novel ENA layer that provides adjustable gradients to each branch network from ID loss. We compare our approach to all aforementioned baseline approaches for our experiments - KD, RKD, CRD, SR, PKT, FN, AT, and WSL.

3 Our Approach

Intermediate Feature Distiller (IFD). We first denote ti and si as the ith layer's feature map of the teacher and student network, respectively. And, zt and zs represent the teacher's and student's logit, respectively as well. Then, the softmax function is defined as follows: σ(zi) = exp(zi/τ) Σj (exp(zj /τ )) , i = 1, ...,M, where zi is the ith class logit in M classes, and τ is the temperature. The IFD is trained with the teacher's knowledge (softmax representation), forwarded by each student's intermediate layers as inputs. Next, we designed several branch networks to sufficiently classify the input instances and utilized the trained IFD block as an auxiliary network to improve the overall performance, as shown in Fig. 1.

More formally, we define F i IFD as the ith IFD network to produce the output from si, and s′i as F i IFD (si) for simplicity. Next, we define SiIF as the softmax output of ith IFD in Eq 1, as follows:

Then, we first design a na¨ıve version of intermediate feature matching loss LIF to demonstrate our training concept. Here, each IFD is individually trained, taking into account the classification object. Specifically, in LIF , we consider not only distillation loss from the teacher but also cross entropy loss for hard label information. Finally, we formulate the intermediate feature matching loss as follows:

where Nblock is the number of blocks, and α is the weight variable between DKL and LCE, and we empirically set α = 0.2.

The intuition behind our design of LIF is that the teacher provides a proper supervision at each stage of a student. Generally, shallower latent features of the student tend to result in weaker representational power of the student. These features are much different from the teacher representation, making training the student model difficult. On the other hand, our framework uses a branch network by training IFDs, where shallow layers and IFDs can easily learn the teacher representation without the risk of losing knowledge.

Element-wise Normalized Attentive (ENA) Layer. In addition, we add the element-wise normalized attentive (ENA) layer to the last layer of each IFD block. We designed the network such that each layer in the network can attend to the representation by itself. Therefore, each IFD can emphasize its class logits by attaching the attentive layer with elementwise multiplication.

As an example, Fig. 2 illustrates that the final results can be changed by ENA layers. Specifically, in Fig. 2, let us assume that the red bars represent the correct labels (e.g., Panda) and the blue bars represent incorrect labels (e.g., Gibbon). Simple aggregation from 3 IFD logits would yield an incorrect result higher than that of the red bars (as shown at the bottom). However, multiplying class-wise weight parameters from the ENA layer provides more weight to the confident class (Panda) in the first IFD block. Therefore, the IFD layer combined with the ENA layer presents optimal inferences for individual classes according to the depth of the network.

Furthermore, the ENA requires only negligible additional parameters and computational cost into an IFD during training and inference time. For example, in CIFAR100 [Krizhevsky and Hinton, 2009], which has 100 classes, 100 parameters are added to the output logits of each IFD layer. Therefore, using four blocks, as shown in Fig. 1, we can achieve performance enhancement with only additional 400 parameters in total. Lastly, to prevent gradient exploding in the early training phase, we apply z-score normalization to each ENA layer at every training iteration. By introducing new attentive layers, we can further refine Eq. 2 to the following equation:

where ⊙ is hadamard product and vENA is a vector for emphasizing their softmax representation.

Integrated Distillation (ID) Loss. For integrating each well-trained IFD block, we introduce integrating distillation (ID) loss LID. Motivated by the early exit method that illustrates that easy tasks can be effectively classified in early layers with a relatively small student capacity, the method can provide fast inference time. Therefore, we focus on exploiting early layers to improve classification performance and inference time rather than exploiting the fixed student model commonly used in traditional KD methods. Since attaching the ENA layers boosts the confidence of each IFD block to an instance, LID is a key component of combining the results from the separated branch networks. With the ENA layer, ID loss can distill and distribute the refined information from the teacher model to the IFD block through the integrated logit from several small branch networks. Therefore, we define LID as follows:

where D is a set of IFD's index for integrated logit as a final prediction, as shown in Fig. 1. In D, it is not compulsory to include all IFDs that are trained with LIF .

Finally, we can construct our overall optimization objective as follows:

where α and β are the weight variable set as 10 and 30, respectively, and LCE is a cross entropy loss between zs and class label.

Inference. Traditional KD methods do not alter the structure of the student network. In contrast, IMF is not restricted to the architecture of the student network, as the architecture of the student network can be adjusted to achieve higher performance. In the inference stage, IMF does not use the entire student network but uses only a part of the network. As illustrated in Fig. 1, a partial network of the student is integrated with D to become an inference model to produce the final outputs. We can control the number of IFD layers of D to make the model lighter or to produce the highest performance. In our experiment, we can empirically obtain the best performance when we use Nblock -1 IFD layers in inference. Therefore, we can adjust the number of IFDs not to exceed the computational cost of the student.

Theoretical Support. We provide the following theoretical results for characterizing the interaction and validity of LIF , and the further need for LID and the ENA layer to achieve higher performance from the gradient perspective.

First, we can obtain ∂LID ∂s′i as follows:

Hence, the gradient from LID depends on each ENA layer and all outputs of branch networks used in inference. Since LID is necessary for LIF , all IFD outputs f i (s′i ) would be matched to the value between the teacher's logit and hard label if it is theoretically optimized.

However, due to the limited capacity and different path lengths of IFD (student) blocks, this solution cannot be achievable. Therefore, we introduce LID and the ENA layer (vENA) in IMF, where, in the Eq. 6, the gradient from LID is affected by ENA layer and all output of the IFD blocks (∑i∈D viENA ⊙ s′i ). It enables the proposed ENA layer, and LID to be attentive to each value of the branch networks.

4 Experimental Results

Training details. Backbone architecture and training settings for experiments are similar to the recent research [Tian et al., 2019]. In our method, we conduct a grid search to choose the α and β values in Eq. 5 from {10, 20, 30, 40}. The IFD block has the same structure in all experiments and model architectures. Specifically, we used a block structure of DepthConv(3 × 3) - PointConv(1 × 1) - BN & ReLU. As the layers became deeper and the resolution decreased, we incrementally reduced the number of blocks by one. Only the channel size of the conv layer was adjusted, based on the size of the original student network. Also, for a fair comparison, we repeated each experiment three times and calculated the mean of the scores.

Comparison. Since IMF uses the flexible architecture with the intermediate IFDs, it is not easy to directly compare IMF with other baselines with the fixed student structures. Instead, to provide a fair comparison as much as possible during inference, we keep our entire computational cost to be the same as the total computation cost of the compared student models in baselines. In other words, the final student model in IMF would not be equivalent to the original student model because IMF-trained student model has a structural difference compared with its original model. However, we can clearly compare the performance of each approach by fixing the total computational cost.

4.1 Performance Evaluation

CIFAR-100. We compared top-1 accuracy for 13 teacherstudent pairs. Furthermore, we calculated the number of parameters of a student model to verify that our method improves performance from a trade-off between the model's size and performance. As shown in Table 1, our IMF outperforms all other methods across all backbone architecture pairs. Especially, our IMF is the only model that outperforms the teacher model in VGG, ResNet110, and WRN.

In addition, our method has fewer parameters than the original student model since we do not utilize an entire structure of the student model in the inference stage. For example, in the ResNet32x4-ResNet8x4 pair, our IMF method has 15.5% of the teacher's parameters but achieves 4.91% performance improvement compared to the original student model. Moreover, as shown in Table 2, IMF also consistently outperforms other KD methods in different teacher-student model architectures.

CIFAR-10. To demonstrate that IMF achieves performance gain in other image classification tasks, we experiment with CIFAR-10. As shown in Table 3, our IMF obtains the performance improvement of up to 1.3% compared to all other competing methods, even though all results are above 90%. Moreover, IMF shows better performance than the teacher networks in VGG19-VGG8 and VGG19-VGG11 pairs, as same with the CIFAR-100 experimental results. Therefore, in the more simple dataset, our IMF outperforms other competing works.

ImageNet. To evaluate IMF performance in a large-scale dataset, we demonstrated experimental results on the ImageNet dataset. In this experiment, we selected the ResNet family as the same base architecture style. As shown in Table 4, IMF outperforms baseline methods with a fewer number of parameters. In detail, IMF's top-1 performance increases by 2.63% and 0.94%, compared to the baseline student as well as the state-of-the-art, SR, respectively. Similar to CIFAR-100, the number of parameters in our model during the inference phase is less than that of the student model, as shown in Table 7.

Facial Keypoints Detection. We conducted experiments not only on the popular classification datasets but also on the regression dataset. Our method is compared with KD, AT, PKT, FN, and RKD with one same and two different network architecture styles, as shown in Table 6. We achieve higher performance than all the other baselines in the different architecture styles, except for the VGG19-VGG8 pair, where our method shows the second-highest performance. However, the overall performance of IMF is much better than others.

4.2 Attributions of IMF

For ablation, we explore the effectiveness of each proposed loss and ENA layer through an ablation study on CIFAR-100.

Effect of LID. First, we started the experiment using only IFD, as shown in Config. A in Table 5. By only using LIF , accuracies significantly increased compared to that of the original student model. Additionally, we conducted Config. B experiment by adding LID to Config. A to integrate the IFD blocks. Interestingly, we observed that LID is helpful in increasing the performances in relatively small student networks (i.e., ResNet20, ResNet32, WRN-16-2, and ShuffleNetV2). Therefore, we found that these models do not have sufficient capacity to learn the representations to classify the images. To compensate for this weakness, LID integrates the small branch networks so that the capacity of the merged model is much superior to the capacity of individual models.

Effect of ENA Layer. To demonstrate the effectiveness of weight to the confident classes from several IFD blocks, we explore the effect of the ENA layer. As mentioned above, we define Config B. as a combination of Config. A and LID. This shows impressive results that the accuracies in ResNet pairs increase, compared with the Config. A. As shown in the second and third rows in Table 5, all performances are much higher in Config. C, which attaches the ENA layer to Config. B, except in the VGG19-VGG8 pair. Therefore, as demonstrated in Section 3, ENA layer provides a weighted gradient to each branch network from ID loss, which shows the best-generalized performance over the teacher-student model pairs.

Effect of LIF . Lastly, we evaluate the performance in Config. D of Table 5, which removes LIF from Config. C. We can observe that this combination still outperforms students' performance. However, compared with Config. C, it shows performance degradation in all pairs except for ResNet110-ResNet32. Interestingly, compared with ResNet20 and ResNet32 student networks, the accuracies decreased greater in VGG networks (more than 3%), which have much more parameters than ResNet20 and ResNet32, demonstrating the benefits of LIF .

4.3 Robustness and Generalization to Various Teacher Models

The previous work [Cho and Hariharan, 2019] shows that the knowledge distillation may not be effective when the student's capacity is too low to successfully mimic the teacher, even the teacher shows better performance. Therefore, finding an adequate teacher and student network pair considering model architecture and capacity is a critical research problem. In this section, we conduct extensive experiments to explore whether our proposed method is robust and generalizable to different teacher-student architecture pairs with varying model capacity. We chose ResNet20 as a student model, while using three ResNet and five VGG families with various sizes as teacher models.

As shown in Fig. 3, ResNet20 with KD [Hinton et al., 2015] shows the best performance with additional supervision from ResNet56 as a teacher. On the other hand, it has a slight performance decrease in the ResNet110 as a teacher, even the latter teacher model is more powerful, as shown by the work [Cho and Hariharan, 2019]. Moreover, in the VGG family as teacher networks, there are marginal performance improvements in the three smallest models, VGG8, VGG11, and VGG13, but we observe a gradual decrease in the performance of the teacher network as the capacity increases.

On the other hand, interestingly, our proposed method achieves performance improvement not only in ResNet but also in VGG as a teacher, even though the teachers have the different architectures and large model capacity. Specifically, compared to the original KD, which shows the secondworst performance in ResNet20-VGG16 pair, IMF shows the highest accuracy, which is higher than the performance of ResNet20 reported in the Table 1. Moreover, compared to KD [Hinton et al., 2015], IMF shows lower variance but higher mean accuracy. Therefore, our proposed IMF is robust and generalizable to different model pairs and capacities, which can reduce the computational time and cost to find the best teacher and student pair.

4.4 Parameters and FLOPs

Our method is flexible in transforming the fixed model architecture into smaller student networks. While our final model can be even smaller than the original student network, our final model was able to achieve higher performance. In particular, we adapted our method to the student model to obtain approximately equal or fewer parameters as well as FLoating point Operations (FLOPs). In this section, we calculated the FLOPs and model parameters of IMF in detail.

We use ResNet32x4-ResNet8x4 pair as teacher and student on CIFAR-100. In this case, we use the three IFD blocks for integrated logit. In Table 7, we reported the parameters and FLOPs of the student model and our IMF. As shown in Table 7, we can observe that Block1:3 has much fewer parameters than that of Block1:4. That is, if we early exit on the third block, we can add branch networks (i.e., IFD1, IFD2, and IFD3) to the student model instead of forwarding the fourth block and MLP layer in the inference phase. Accordingly, it is notable to observe that the final parameters and FLOPs of IMF can be approximately similar to or less than that of student models.

5 Limitation and FutureWork

In this work, we focused on image classification and regression tasks, and presented the theoretical analysis to characterize the validity of our proposed methods. More research would be needed to determine the number of branch networks to find the optimal numbers for the depth of the original backbone network and IFD blocks as well as the branch networks. For future work, we plan to incorporate additional KD methods and validate our findings.

6 Conclusion

We proposed a novel flexible knowledge distillation framework, IMF, integrating matched features using attentive logit. In IMF, the student is directly trained by teacher information by utilizing an intermediate feature distiller (IFD) in a branch network. During the inference time, we use the combined logits from learned IFDs instead of the entire student model. We conducted extensive experiments with image classification as well as facial keypoint detection tasks over 4 datasets. In image classification, our method significantly outperforms 7 baselines in 10 same and 4 different backbone architecture styles with an equal or less number of parameters. Furthermore, we also demonstrate that our method is effective in keypoint detection, compared with 5 baselines. IMF not only achieved a significant performance gain regarding classification tasks as well as a regression task, without requiring an extra computational cost.

Acknowledgments

This work was partly supported by Institute for Information & communication Technology Planning & evaluation (IITP) grants funded by the Korean government MSIT: (No. 2022-0-01199, Graduate School of Convergence Security at Sungkyunkwan University), (No. 2022-0-01045, Self-directed Multi-Modal Intelligence for solving unknown, open domain problems), (No. 2022-0-00688, AI Platform to Fully Adapt and Reflect Privacy-Policy Changes), (No. 2021-0-02068, Artificial Intelligence Innovation Hub), (No. 2019-0-00421, AI Graduate School Support Program at Sungkyunkwan University), and (No. 2023-00230337, Advanced and Proactive AI Platform Research and Development Against Malicious deepfakes).