Hierarchical Self-supervised Augmented Knowledge Distillation

1 Introduction

Orthogonal to efficient network architecture designs [Yang et al., 2019; Zhu et al., 2019; Yang et al., 2020], Knowledge Distillation (KD) [Hinton et al., 2015] aims to transfer knowledge from a pre-trained high-capacity teacher network to a light-weight student network. The student's performance can often be improved significantly, benefiting from the additional guidance compared with the independent training. The current pattern of KD can be summarized as two critical aspects: (1) what kind of knowledge encapsulated in teacher network can be explored for KD; (2) How to effectively transfer knowledge from teacher to student.

The original KD [Hinton et al., 2015] minimizes the KLdivergence of predictive class probability distributions between student and teacher networks, which makes intuitive sense to force the student to mimic how a superior teacher generates the final predictions. However, such a highly abstract dark knowledge ignores much comprehensive information encoded in hidden layers. Later works naturally proposed to transfer feature maps [Romero et al., 2015] or their refined information [Zagoruyko and Komodakis, 2017; Heo et al., 2019; Ahn et al., 2019] between intermediate layers of teacher and student. A reasonable interpretation of the success of feature-based distillation lies in that hierarchical feature maps throughout the CNN represent the intermediate learning process with an inductive bias of the final solution. Beyond knowledge alignment limited in the individual sample, more recent works [Peng et al., 2019;

∗Corresponding author

Figure 1: Difference of self-supervised knowledge between SSKD and our method. (a) SSKD applies contrastive learning by forcing the image and its transformed version closed against other negative images in the feature embedding space. It defines the contrastive relationship as knowledge. (b) Our method unifies the original task and self-supervised auxiliary task into a joint task and defines the self-supervised augmented distribution as knowledge.

Tian et al., 2020] leverage cross-sample correlations or dependencies in high-layer feature embedding space. Inspired by the recent success of self-supervised visual representational learning [Chen et al., 2020], SSKD [Xu et al., 2020] introduces an auxiliary self-supervised task to extract richer knowledge. As shown in Fig. 1a, SSKD proposes transferring cross-sample self-supervised contrastive relationships, making it achieve superior performance in the field of KD.

However, forcing the network to learn invariant feature representations among transformed images using a self-supervised pretext task with random rotations from 0◦ , 90◦ , 180◦ , 270◦ utilized in SSKD may destroy the original visual semantics (e.g. 6 v.s. 9). It would increase the difficulty of representation learning for semantic recognition tasks. As validated in Table 1, applying random rotation as an additional data augmentation degrades the classification performance, especially on more challenging TinyImageNet.

To effectively learn knowledge from self-supervised representation learning without interfering with the original fullysupervised classification task, we use a unified task by combining the label spaces of the original task and self-supervised task into a joint label space, as shown in Fig. 1b. This task is partly inspired by the previous seminal self-supervised representation learning [Gidaris et al., 2018; Lee et al., 2020]. We further introduce these prior works to explore more powerful knowledge for distillation. To verify the effectiveness of the self-supervised augmented label, we also conduct initial exploratory experiments on standard image classification in Table 1. We find that the performance can be significantly improved by SAL, which can be attributed to learned better feature representations from an extra well-combined selfsupervised task. The good performance further motivates us to define the self-supervised augmented distribution as a promising knowledge for KD.

Another valuable problem lies in how to transfer the probabilistic knowledge between teacher and student effectively. Vanilla KD aligns probability distributions only in the final layer but ignores comprehensive knowledge. Feature-based distillation methods provide one-to-one matching between the same convolutional stages of teacher and student. However, matched feature maps may have different semantic abstractions and result in a negative supervisory effect [Passalis et al., 2020]. Compared with feature information, the probability distribution is indeed a more robust knowledge for KD, especially when existing a large architecture gap between teacher and student [Tian et al., 2020].

However, it is difficult to explicitly derive comprehensive probability distributions from hidden layers over the original architecture. Therefore a natural idea is to append several auxiliary classifiers to the network at various hidden layers to generate multi-level probability distributions from hierarchical feature maps. It allows us to perform comprehensive one-to-one matching in hidden layers in terms of probabilistic knowledge. Moreover, it is also noteworthy that the gap of abstraction level of any matched distributions would be easily reduced due to the delicately-designed auxiliary classifiers.

We guide all auxiliary classifiers attached to the original network to learn informative self-supervised augmented distributions. Furthermore, we perform Hierarchical Selfsupervised Augmented Knowledge Distillation (HSAKD) between teacher and student towards all auxiliary classifiers in a one-to-one manner. By taking full advantage of richer selfsupervised augmented knowledge, the student can be guided to learn better feature representations. Note that all auxiliary classifiers are only used to assist knowledge transfer and dropped during the inference period. The overall contributions are summarized as follows:

• We introduce a self-supervised augmented distribution that encapsulates the unified knowledge of the original classification task and auxiliary self-supervised task as the richer dark knowledge for the field of KD.
• We propose a one-to-one probabilistic knowledge distillation framework by leveraging the architectural auxiliary classifiers, facilitating comprehensive knowledge transfer and alleviating the mismatch problem of abstraction levels when existing a large architecture gap.
• HSAKD significantly refreshes the results achieved by previous SOTA SSKD on standard image classification benchmarks. It can also learn well-general feature representations for downstream semantic recognition tasks.

2 RelatedWork

Knowledge Distillation. The seminal KD [Hinton et al., 2015] popularizes the pattern of knowledge transfer with a soft probability distribution. Later methods further explore feature-based information encapsulated in hidden layers for KD, such as intermediate feature maps [Romero et al., 2015], attention maps [Zagoruyko and Komodakis, 2017], gram matrix [Yim et al., 2017] and activation boundaries [Heo et al., 2019]. More recent works explore cross-sample relationship using high-level feature embeddings with various definitions of edge weight [Park et al., 2019; Peng et al., 2019; Tung and Mori, 2019]. The latest SSKD [Xu et al., 2020] extracts structured knowledge from the self-supervised auxiliary task. Beyond knowledge exploration, [Ahn et al., 2019; Tian et al., 2020] maximize the mutual information between matched features. [Yang et al., 2021] further extend this idea into online KD in a contrastive-based manner. To bridge the gap between the teacher and student, [Passalis et al., 2020] introduce teacher assistant models for smoother KD. However, extra teacher models increase the complexity of the training pipeline. Therefore, we choose to append several welldesigned auxiliary classifiers to alleviate the knowledge gap and facilitate comprehensive knowledge transfer.

Self-supervised Representational Learning (SRL). The seminal SRL popularizes the pattern that guides the network to learn which transformation is applied to a transformed image for learning feature representations. Typical transformations can be rotations [Gidaris et al., 2018], jigsaw [Noroozi and Favaro, 2016] and colorization [Zhang et al., 2016]. More recently, [Misra and Maaten, 2020; Chen et al., 2020] learn invariant feature representations under self-supervised pretext tasks by maximizing the consistency of representations among various transformed versions of the same image. Both SSKD and our HSAKD are related to SRL. SSKD uses the latter SRL pattern to extract knowledge. In contrast, HSAKD combines the former classification-based pattern of SRL with the fully-supervised classification task to extract richer joint knowledge.

3 Method

3.1 Self-supervised Augmented Distribution

We present the difference between the original class probability distribution and self-supervised augmented distribution using a conventional classification network of CNN. A CNN can be decomposed into a feature extractor Φ(·; µ) and a linear classifier g(·;w), where µ and w are weight tensors. Given an input sample x ∈ X, X is the training set, z = Φ(x;µ) ∈ Rd is the extracted feature embedding vector, where d is the embedding size. We consider a conventional N-way object classification task with the label space N = {1, · · · , N}. The linear classifier attached with softmax normalization maps the feature embedding z to a predictive class probability distribution p(x; τ ) = σ(g(z; w)/τ ) ∈ RN over the label space, where σ is the softmax function, weight matrix w ∈ RN×d , τ is a temperature hyper-parameter to scale the smoothness of distribution.

We introduce an additional self-supervised task to augment the conventional supervised object class space. Learning such a joint distribution can force the network to generate more informative and meaningful predictions benefiting from the original and auxiliary self-supervised task simultaneously. Assuming that we define M various image transformations {tj}Mj=1 with the label space M = {1, · · · ,M}, where t1 denotes the identity transformation, i.e. t1(x) = x. To effectively learn composite knowledge, we combine the class space from the original supervised object recognition task and self-supervised task into a unified task. The label space of this task is K = N ⊗M, here ⊗ is the Cartesian product. |K| = N ∗ M, where |·| is the cardinality of the label collection, * denotes element-wise multiplication.

Given a transformed sample ˜x ∈ {tj(x)}Mj=1 by applying one transformation on x, ˜ z = Φ(˜x;µ) ∈ Rd is the extracted feature embedding vector, q( ˜x; τ ) = σ(g( ˜ z; w)/τ ) ∈ RN∗M is the predictive distribution over the joint label space K, where weight tensor w ∈ R(N∗M)×d . We use p ∈ RN to denote the normal class probability distribution and q ∈ RN∗M to denote the self-supervised augmented distribution.

3.2 Auxiliary Architecture Design

It has been widely known that feature maps with various resolutions encode various patterns of representational information. Higher-resolution feature maps often present more fine-grained object details, while lower-resolution ones often contain richer global semantic information. To take full advantage of hierarchical feature maps encapsulated in a single network, we append several intermediate auxiliary classifiers into hidden layers to learn and distill hierarchical selfsupervised augmented knowledge.

For ease of notation, we denote a conventional classification network as f(·), which maps a input sample tj(x), j ∈ M to the vanilla class probability distribution p(tj(x); τ) = σ(f(tj(x))/τ) ∈ RN over the original class space N. Modern CNNs typically utilize stage-wise convolutional blocks to gradually extract coarser features as the depth of the network increases. For example, popular ResNet-50 for ImageNet classification contains consecutive four stages, and extracted feature maps produced from various stages have different granularities and patterns. Assuming that a network contains L stages, we choose to append an auxiliary classifier after each stage, thus resulting in L classifiers {cl(·)}Ll=1 , where cl(·) is the auxiliary classifier after l-th stage. cl(·) is composed of stage-wise convolutional blocks, a global average pooling layer and a fully connected layer. Denoting the extracted feature map after l-th stage as Fl, we can obtain the self-supervised augmented distribution inferred by cl(·) as ql(tj(x); τ) = σ(cl(Fl))/τ) ∈ RN∗M over the joint class space K. The overall design of auxiliary classifiers over a 3-stage network is illustrated in Fig. 2 for example. The detailed design formulation of various auxiliary classifiers for a specific network can be found in our released codes.

3.3 Training the Teacher Network

We denote the teacher backbone network as f T (·) and L auxiliary classifiers as {cTl (·)}Ll=1 . We conduct an end-to-end training process for preparing the teacher network. On the one hand, we train the f T (·) with normal data x by the conventional Cross-Entropy (CE) loss to fit the ground-truth label y ∈ N, and pT (x; τ) = σ(f T (x)/τ) ∈ RN is the predictive class probaility distribution. On the other hand, we aim to train L auxiliary classifiers {cTl (·)}Ll=1 for learning hierarchical self-supervised augmented distributions. Given an input sample tj(x), we feed the feature maps {F T l,j }Ll=1 generated from backbone f T (·) to {cTl (·)}Ll=1 , respectively. The predictive self-supervised augmented distribution inferred by the lth classifier cTl is qT l (tj(x); τ) = σ(cTl (F T l,j ))/τ) ∈ RN∗M . We train all auxiliary classifiers using CE loss with selfsupervised augmented labels across {tj(x)}Mj=1 as Eq. (1).

Where τ = 1 and Lce denotes the Cross-Entropy loss. For a bit abuse of notation, we use kj to denote the self-supervised augmented label of tj(x) in joint class space K. The overall loss for training a teacher is shown in Eq. (2).

Note that the two losses in Eq. 2 have different roles. The first loss aims to simply fit the normal data for learning general classification capability. The second loss aims to generate additional self-supervised augmented knowledge by the existing hierarchical features derived from the backbone network. This method facilitates richer knowledge distillation benefiting from the self-supervised task beyond the conventional fully-supervised task.

3.4 Training the Student Network

We denote the student backbone network as f S (·) and L auxiliary classifiers as {cSl (·)}Ll=1 . We conduct an end-to-end training process under the supervision of the teacher network. The overall loss includes the task loss from ground-truth labels and the mimicry loss from the pre-trained teacher.

Task Loss

We force the f S (·) to fit the normal data x as the task loss:

Where pS (x; τ) = σ(f S (x)/τ) ∈ RN is the predictive class probaility distribution. We also had tried to force L auxiliary classifiers {cSl (·)}Ll=1 to learn the self-supervised augmented distributions from the joint hard label of the original and selfsupervised tasks by LSce SAD as the additional loss:

Where qS l (tj(x); τ) = σ(cSl (F S l,j ))/τ) ∈ RN∗M , and F S l,j is the extracted feature map from the l-th stage of f S (·) for the input tj(x). However, we empirically found that introducing loss (4) into the original task loss (3) damages the performance of student networks, as validated in Section 4.2.

Mimicry Loss

On the one hand, we consider transferring hierarchical selfsupervised augmented distributions generated from L auxiliary classifiers of the teacher network to corresponding L auxiliary classifiers of the student network, respectively. The transfer performs in a one-to-one manner by KL-divergence loss DKL. The loss is formulated as Eq. (5), where τ 2 is used to retain the gradient contributions unchanged [Hinton et al., 2015].

Benefiting from Eq. (5), one can expect that the student network gains comprehensive guidances by unified selfsupervised knowledge and the original class full-supervised knowledge. The informative knowledge is derived from multi-scale intermediate feature maps encapsulated in hidden layers of the high-capacity teacher network.

On the other hand, we transfer the original class probability distributions generated from the final layer between teacher and student. Specifically, we transfer the knowledge derived from both the normal and transformed data {tj(x)}Mj=1 , where t1(x) = x. This loss is formulated as Eq. (6).

We do not explicitly force the student backbone f S (·) to fit the transformed data in task loss for preserving the normal classification capability. But mimicking the side product of predictive class probability distributions inferred from these transformed data from the teacher network is also beneficial for the self-supervised representational learning of student network, as validated in Section 4.2.

Overall loss. We summarize the task loss and mimicry loss as the overall loss LS for training the student network.

Following the wide practice, we set the hyper-parameter τ = 1 in task loss and τ = 3 in mimicry loss. Besides, we do not introduce other hyper-parameters.

4 Experiments

4.1 Experimental Settings

We conduct evaluations on standard CIFAR-100 and ImageNet [Deng et al., 2009] benchmarks across the widely applied network families including ResNet [He et al., 2016], WRN [Zagoruyko S, 2016], VGG [Simonyan and Zisserman, 2015], MobileNet [Sandler et al., 2018] and ShuffleNet [Zhang et al., 2018; Ma et al., 2018]. Some representative KD methods including KD [Hinton et al., 2015], FitNet [Romero et al., 2015], AT [Zagoruyko and Komodakis, 2017], AB [Heo et al., 2019], VID [Ahn et al., 2019], RKD [Park et al., 2019], SP [Tung and Mori, 2019], CC [Peng et al., 2019], CRD [Tian et al., 2020] and SOTA SSKD [Xu et al., 2020] are compared. For a fair comparison, all comparative methods are combined with conventional KD by default, and we adopt rotations {0◦ , 90◦ , 180◦ , 270◦ } as the self-supervised auxiliary task as same as SSKD. We use the standard training settings following [Xu et al., 2020] and report the mean result with a standard deviation over 3 runs. The more detailed settings for reproducibility can be found in our released codes.

4.2 Ablation Study

Effect of loss terms. As shown in Fig. 3 (left), applying hierarchical self-supervised augmented knowledge transfer through multiple auxiliary classifiers by loss Lkl q substantially boosts the accuracy upon the original task loss Ltask. We further compare Lkl p and Lkd upon Ltask + Lkl q to demonstrate the efficacy of transferring class probability distributions from additional transformed images. We find that Lkl p results in better accuracy gains than Lkd, which suggests that transferring probabilistic class knowledge from those transformed images is also beneficial to feature representation learning. Finally, beyond the mimicry loss, we also explore whether the self-supervised augmented task loss LSce SAD can be integrated into the overall task loss to train student networks. After adding LSce SAD upon the above losses, the performance of the student network is dropped slightly. We speculate that mimicking soft self-supervised augmented distribution from the teacher by Lkl q is good enough to learn rich self-supervised knowledge. Extra learning from hard one-hot distribution by LSce SAD may interfere with the process of self-supervised knowledge transfer.

Effect of auxiliary classifiers. We append several auxiliary classifiers to the network with various depths to learn and transfer diverse self-supervised augmented distributions extracted from hierarchical features. To examine this practice, we first individually evaluate each auxiliary classifier. As shown in Fig. 3 (right), we can observe that each auxiliary classifier is beneficial to performance improvements. Moreover, the auxiliary classifier attached in the deeper layer often achieves more accuracy gains than that in the shallower layer, which can be attributed to more informative semantic knowledge encoded in high-level features. Finally, using all auxiliary classifiers can maximize accuracy gains.

4.3 Comparison with State-Of-The-Arts

Results on CIFAR-100 and ImageNet. We compare our HSAKD with SOTA representative distillation methods across various teacher-student pairs with the same and different architectural styles on CIFAR-100 in Table 2 and on ImageNet in Table 3. Interestingly, using LTce SAD as an auxiliary loss can improve the teacher accuracy of the final classification. Compared to the original 'Teacher', 'Teacher*' achieves an average gain of 4.09% across five teacher networks on CIFAR-100 and a top-1 gain of 2.17% for ResNet-34 on ImageNet. For downstream student networks, 'Teacher*' leads to an average improvement of 1.15% on CIFAR-100 and an improvement of 0.23% on ImageNet than 'Teacher'. These results indicate that our proposed loss LT can improve the performance of a given network and produce a more suitable teacher for KD to learn a better student. Moreover, our HSAKD significantly outperforms the best-competing method SSKD across all network pairs with an average accuracy gain of 2.56% on CIFAR-100 and a top-1 gain of 0.77% on ImageNet. Compared with other SOTA methods, the superiority of HSAKD can be attributed to hierarchical self-supervised augmented knowledge distillation by the assistance with well-designed auxiliary classifiers.

Transferability of Learned Representations. Beyond the accuracy on the object dataset, we also expect the student network can produce the generalized feature representations that transfer well to other unseen semantic recognition datasets. To this end, we freeze the feature extractor pre-trained on the upstream CIFAR-100, and then train two linear classifiers based on frozen pooled features for downstream STL-10 and TinyImageNet respectively, following the common linear classification protocal [Tian et al., 2020]. As shown in Table 4, we can observe that both SSKD and HSAKD achieve better accuracy than other comparative methods, demonstrating that using self-supervised auxiliary tasks for distillation is conducive to generating better feature representations. Moreover, HSAKD can significantly outperform the best-competing SSKD by 3.63% on STL-10 and 3.50% on TinyImageNet. The results verify that encoding the selfsupervised auxiliary task as an augmented distribution in our HASKD has better supervision quality than the contrastive relationship in SSKD for learning good features.

Transferability for Object Detection. We further evaluate the student network ResNet-18 pre-trained with the teacher ResNet-34 on ImageNet as a backbone to carry out downstream object detection on Pascal VOC. We use FasterRCNN [Ren et al., 2016] framework and follow the standard data preprocessing and finetuning strategy. The comparison on detection performance is shown in Table 5. Our method outperforms the original baseline by 2.27% mAP and the best-competing SSKD by 0.85% mAP. These results verify that our method can guide a network to learn better feature representations for semantic recognition tasks.

Efficacy under Few-shot Scenario. We compare our method with conventional KD and SOTA CRD and SSKD under few-shot scenarios by retaining 25%, 50% and 75% training samples. For a fair comparison, we use the same data split strategy for each few-shot setting, while maintaining the original test set. As shown in Table 6, our method can consistently surpass others by large margins under various few-shot settings. Moreover, it is noteworthy that by using only 25% training samples, our method can achieve comparable accuracy with the baseline trained on the complete set. This is because our method can effectively learn general feature representations from limited data. In contrast, previous methods often focus on mimicking the inductive bias from intermediate feature maps or cross-sample relationships, which may overfit the limited set and generalize worse to the test set.

5 Conclusion

We propose a self-supervised augmented task for KD and further transfer such rich knowledge derived from hierarchical feature maps leveraging well-designed auxiliary classifiers. Our method achieves SOTA performance on the standard image classification benchmarks in the field of KD. It can guide the network to learn well-general feature representations for semantic recognition tasks. Moreover, it has no hyper-parameters to be tuned and is easy to implement.