Similarity Transfer for Knowledge Distillation

Recently, CNNs become predominant in many computer vision tasks. Network architectures become deeper and wider which achieve better performances than the shallow ones. However, these complicated networks also bring in high computation and memory costs, which obstruct us to deploy these models into real-time applications. Thus, model compression and acceleration aim to reduce the model size and computation cost meanwhile maintaining high performance. One kind of method directly designs small network structure by modifications on convolution since the original deep model has many redundant parameters. For example, MobileNet [33] introduces depth-wise separable convolution to build block instead of standard convolution. ShuffleNet [34] proposes point-wise group convolution and channel shuffle to maintain a decent performance without adding the burden of computation. These approaches could retain a decent performance without adding too much computing burden at inference phase.

Besides, there has been another kind of method which attempts to remove the redundant information of teacher network. For example, network pruning aims to boost the speed of inference by getting rid of redundancy of the large CNN model. Han et al. [35] propose to boost the speed of inference by deleting unimportant connections. However, network pruning methods require many iterations to converge and the pruning threshold needs to be set manually. Network decomposition use matrix decomposition technique to decompose the original convolution kernel in CNNs model. For example, Novikov et al. [36] propose to reduce the number of parameters while preserve the original performance of network by converting the dense weight matrices of the fully-connected layers to the Tensor Train format. Network quantization methods aim to represent the float weights with fewer bits. Yang et al. [37] propose to simulate the quantization process with a differentiable function. Note that, above mentioned works can be combined with KD methods for further improvement.

Different from above methods, knowledge distillation enrich and get the student model by extracting kinds of knowledge from the fixed teacher model. To address the challenge of deploying CNNs in resource-constrained edge devices, Bucilua et al. [38] first propose to transfer the knowledge of an ensemble of models to a small model. Then Caruana et al. [39] propose to train student model by mimicking the teacher model’s logits. Later, Hinton et al. [12] popularize the idea of knowledge distillation, which efficiently transfers knowledge from large teacher network to compact student network by mimicking the class probabilities outputs. Note that, the outputs of the teacher network are defined as the dark knowledge in KD, which provides similarity information as extra supervisions compared with one-hot labels. In other words, there are two sources of supervision used to train the student in KD, one from the ground-truth label, and the other from the soft targets of teacher network.

Afterward, some recent works [14] [15] extend KD by distilling knowledge from intermediate feature representations instead of soft labels. For example, FitNets [14] propose to train student network by mimicking the intermediate feature maps of teacher network, which are defined as hints. Inspired by this, Zagoruyko et al.  [15] propose to match the attention maps between the teacher and the student, which are defined from the original feature maps as knowledge. Wang et al. [40] propose to improve the performance of student network by matching the distributions of spatial neuron activations between the teacher and the student. Recently, Heo et al. [41] introduce the activation boundary of the hidden neuron as knowledge for distilling the compact student network.

However, the aforementioned knowledge distillation methods only utilize the knowledge contained in the output of specific layers of the teacher network. More richer knowledge between different layers is explored and utilized for knowledge distillation. For example, Yim et al. [16] propose to use Gram matrix between different feature layers as distilled knowledge, which named flow of solution process (FSP) that reflects the relations of different features maps. Lee et al. [42] propose to extract valuable correlation information as knowledge between different feature maps using singular value decomposition (SVD). However, these methods only focus on the knowledge contained in the individual data samples and ignore the similarities between categories of multiple samples, which also valuable for knowledge distillation.

Thus, some very recent works [17][18] aim to explore the relationship between data samples. Park et al. [17] propose a relation knowledge distillation (RKD) method, which extracts mutual relations of data samples by the proposed distance-wise and angle-wise distillation losses. Tung et al. propose a similarity preserving (SP) knowledge distillation method which transfers the similar activations of input pairs from the teacher network to the student network. Peng et al. [43] propose to capture the high order correlation between samples using the kernel method. However, the similarities information between instances is hard to learn for the student network due to the high dimension in the embedding space. Different from existing works, the proposed approach presents the similarity correlation between instances through the probability distribution outputs of virtual samples, which are created by the mixup technique. In contrast to the vanilla KD which softens the softmax outputs using a temperature hyperparameter, our approach directly outputs a softened probability distribution using the virtual samples.

There are many training strategies to further improve the performance of CNNs, such as data augmentation and regularization techniques. In particular, mixed sample data augmentation has gain increasing attention due to its state-of-the-art performance. It trains the models using the augmented data set by combining data samples according to some policy.

For example, Zhang et al. [32] first introduce a data augmentation approach named mixup, which could improve the generalization of CNNs by linearly interpolating a random pair of original training data to create a new training data and corresponding labels. Inspired by this, some mixup variants [44][45] [46] have been proposed. Later, Cutmix [47] proposes to replace the removed regions with a patch from another image. FMix [48] proposes to use binary masks obtained by applying a threshold to low frequency images sampled from Fourier space. Following them, Pairing Samples [49] propose to take an average of two samples for each pixel. Recently, Cubuk et al. [50] employ automatically searching to improve data augmentation policies that called AutoAugment. Although these methods have achieved remarkable performance to the training of CNNs, they suffer from the prohibitive training time.

In our approach, we aim to explore a more effective knowledge, which contains in the similarity correlation between instances, to guide the training of student model. Moreover, we employ mixup technique to represent these knowledge by introducing virtual samples.

For simplicity, we define the large, complex Teacher deep neural network as T with learned parameters $P_{T}$, the less complex Student network as S with learned parameters $P_{S}$. The original training data consist of tuples of input data and target $(x,y)\in D$. In general, the conventional knowledge distillation methods train the Student by minimizing the following objective function with respect to the parameters $P_{S}$ over the training samples $(x,y)\in D$:

For example, Hinton et al. [12] use pre-softmax outputs for $T(x,P_{T})$ and $S(x,P_{S})$, and Kullback-Leibler divergence for $\mathcal{L}_{KD}$:

Likewise, other works [14][16] involved an objective function that can also be formulated as a form of Eq. 1. However, these conventional KD methods only focus on the similarity correlation between categories of individual samples while ignore the similarity correlation between different samples. Moreover, the temperature hyperparameter $t$ need to be set manually which is used to soften the probability distribution.

In this section, we introduce similarity transfer for knowledge distillation in detail. Usually, traditional knowledge distillation methods [12] utilize the temperature hyperparameter to soften the one-hot labels to soft labels. Different from existing works, STKD presents the similarity correlation between different samples through the probability distribution outputs of virtual samples, which are created by the mixup technique.

To be specific, we first employ the recently proposed mixup [32] to create the virtual samples. We randomly sample two samples $(x_{i},y_{i})$ and $(x_{j},y_{j})$ from $D$. Then it generates the new synthetic sample $(x,y)$ by a weighted linear interpolation of these two samples:

As can be seen from Figure. 1, the virtual samples are created by this mixup technique. In other words, each sample from the virtual set is formed by two arbitrary original images by a weighted linear interpolation.

By Eq. 4, we turn the one-hot labels to mixed labels, which also contribute to the similarity information between different samples. Thus, the mixed labels of the virtual samples are employed in our framework. Therefore, the mix loss can be formulated as:

Then we distill the similarity knowledge between multiple samples from teacher network to student network using the created virtual samples. Here we define the total loss function for similarity transfer knowledge distillation as follows:

Different from the vanilla KD method, we get rid of the temperature hyperparameter in our method. Thus, the probability distribution output of $T$ and $S$ is defined as $T(x,P_{T})=softmax(T(x_{i},P_{T})$ and $S(x,P_{S})=softmax(S(x_{i},P_{S})$ respectively. Note that the classification loss has the hyper-parameter $\lambda$ to balance the weight of the mix labels.

Unlike previous methods, we define the similarity correlation between instances as knowledge and present it through probability distribution of the virtual images. To this end, we employ the mixup technique to create virtual samples which is effective for our method. It generates mixup images that contain multiple images’ similarity correlation on the one hand, on the other hand also regularizes the student model through knowledge distillation to favor simple linear behavior in-between training samples. Moreover, we could generate different mixup images by varying the coefficient $\lambda$. We study this coefficient in the ablation study.

In this way, teacher network transfers more valuable information to the student network than existing methods. That means the teacher’s dark knowledge could be better transferred to the student.

Algorithm 1 describes the details of the whole training paradigm. We first obtain a pre-trained teacher model $T$ with parameters $P_{T}$ which is trained using standard back-propagation on original dataset $(x,y)\in D$.

Then we iterate the following stages until the accuracy of the student $S$ converges. 1) We create the virtual set instead of original dataset using the mixup technique. Specifically, we use a single data loader to obtain the mini-batch and apply the mixup to the same mini-batch after random shuffling. 2) We feed these virtual samples to our knowledge distillation framework, which distills the similarity correlation between instances from the pre-trained teacher to the student network using our STKD loss in Eq. 6.

Fig. 1 illustrates the overall framework of STKD. Different from the previous methods, the proposed STKD regularizes the student network to favor the simple linear behavior under the teacher’s guidance. In other words, we employ the mixup technique to provide more similarity information among classes and make the high dimensional representation of teacher more easily transferable to the student. We will demonstrate its effectiveness in the following section.

On CIFAR dataset, we use three groups of teacher/student model pairs including ResNet-110/ResNet-20, WRN-40-2/WRN-40-1 and VGG-13/VGG-8. We apply a standard horizontal flip and random crop data augmentation scheme which is widely adopted for these datasets. We set the weight decay to $10^{-4}$., batch size to 128, and use stochastic gradient descent (SGD) with momentum 0.9. The initial learning rate is set to $\gamma$ = 0.1 and divided by 10 at 80th, 100th, 150th epochs, totally 200 epochs. We use the same setting for training the teacher network and distilling the student network.

We conduct baseline comparisons with respect to the vanilla KD [12], AT [15], SP [18] and RKD [17]. For vanilla KD, we set temperature hyperparameter $t=4$ following the experiments in [15]. We set the hyperparameter $\beta=1000$ for AT, $\gamma=0.1$ for SP respectively.

Table I summarizes the classification accuracy of several teacher/student model pairs on CIFAR-10/100 datasets. For WRN network architecture, the teacher and student networks have the same depth but different width (WRN-40-2 teacher with WRN-40-1 student).

Our method gets 94.44%, 93.52% and 93.68% of top-1 accuracy for WRN-40-1, ResNet-20 and VGG-8 on CIFAR-10, respectively. It substantially surpasses the vanilla KD by 1.6%, 1,9% and 1.8%. Compared to other baselines, the student network from STKD still significantly surpass the three related SOTA methods. For CIFAR-100, our method also has 1.6%, 1.7% and 1.8% top-1 accuracy promoting over vanilla KD. Moreover, our approach still surpasses the three SOTA distillation related methods, which verifies the effectiveness of our approach. These results validate our intuition that the similarity correlation between instances encodes valuable knowledge from the teacher network and provides an important supervision for knowledge distillation.

Figure 2 plots the accuracy change curves for WRN-40-2/WRN-40-1 experiments on CIFAR-10. As can be seen from Figure 2 (a), it shows the testing accuracy curves of teacher network, student network which is trained individually and student network from our method. Note that, STKD gets a significant improvement compared to the student trained individually and its final accuracy value is close to that of the teacher. Figure 2 (b) shows the validation accuracy change curves of WRN-40-1 over time among different knowledge distillation approaches on CIFAR-10 dataset. We can observe that our method performs a significant improvement on the final accuracy and outperforms the other SOTA approaches.

From the perspective of the teacher model, our method also shows the potential to compress large networks into more compact ones with minimal accuracy loss. For example, we distill the knowledge from the pre-trained WRN-40-2 teacher network, which contains $2.20M$ parameters, to a much smaller WRN-40-1 student network, which contains $0.56M$ parameters. The student network gets a 4$\times$ compression rate with only $0.5\%$ loss in classification accuracy using off-the-shelf Pytorch.

In this section, we conduct experiments on CINIC-10 dataset, whose scale is closer to that of ImageNet. We use WRN-40-2 and WRN-40-1 as the teacher network and student network respectively. For training phase, we apply CIFAR-style data augmentation with horizontal flips and random crops. The beginning learning rate is set to 0.01 and decayed by a factor of 10 after the 100th, 160th and 220th epochs. All the networks are trained for 240 epochs using SGD with Nesterov momentum, a batch size of 96.

For fair comparison, we adopt the same setting as mentioned above for all baseline methods. We set the hyperparameters for the baseline methods ($t=8$ for KD; $\beta=100$ for AT, $\gamma=0.1$ for SP). The results on CINIC-10 are shown in Table II. Our method achieves a 85.94% classification accuracy, which surpasses the student network trained individually by promoting 1.64%. Moreover, STKD substantially surpasses the vanilla KD by 0.9%. It also shows comparable performance with other SOTA approaches. These results imply that STKD method induces more meaningful dark knowledge such as the similarity information between instances than other baseline methods.

For rapid experimentation, we conduct experiments on Tiny-ImageNet dataset, which is a popular subset of the ImageNet. We adopt the VGG network architecture for the following experiments. To be specific, VGG-13 is employed for the teacher network, which achieves 62.21% classification accuracy. VGG-8 is used for the student network. Both of the networks are trained for 240 epochs.

In our Tiny ImageNet classification experiments, we apply random rotation and horizontal flipping for data augmentation. We optimize the model using stochastic gradient descent(SGD) with mini-batch 128 and momentum 0.9. The learning rate starts from 0.1 and is multiplied by 0.2 at 60, 120, 160, 200, 250 epochs. We totally train the network for 300 epochs and adopt the deep and wide WRN (WRN-40-1) for a teacher model and WRN-16-1 as a student model.

In this section, ResNet with four blocks is selected with channel size of 16, 32, 64 and 128. Random crop and horizontal flip are used for data augmentation. Table III shows the results, where baseline denotes the student network trained individually. As can be seen from Table III, our method, as well as the other SOTA methods, gets higher classification accuracy than the original student network. Moreover, our STKD method achieves better performance than the KD and shows comparable performance with the other knowledge distillation methods, and further overcomes them.

Evaluation on Different Network Architectures. To further evaluate the performance of STKD, we perform additional experiments on different network architecture families for the student/teacher pairs on CIFAR-100. Table IV shows the experimental results.

As can be seen from Table IV, STKD continuously outperforms KD and other SOAT methods, which indicates the effectiveness of our method. In particular, STKD achieves the accuracy of 69.07%, 73.81% and 68.22% for three network pairs of different network families respectively, obtaining a performance gain of 1.07%, 0.92% and 2.33% over SP. However, we observe that some methods such as AT [15] perform quite poorly when the teacher and student belong to different network architectures. Thus, the results validate that our approach is robust to teacher/student pairs. And we conclude that the knowledge contained in the similarity correlation between different instances helps to improve the robustness of student network.

Impact of Different coefficients for mixup. In this subsection, we explore the impact of coefficients for mixup. We adopt ResNet-110 as the teacher network, ResNet-20 as the student network. The results on CIFAR-100 could be found in Table V, which illustrates how the performance of STKD is affected by the choice of the coefficient $\lambda$. By setting different $\lambda$ for STKD, we observe that STKD achieves the accuracy of 72.14% ($\lambda=0.50$), which is the best performance in our settings.

Furthermore, we show the mixup images from the same sample pairs by varying the mixup coefficient $\lambda$ in Figure 3. The outputs of probability distribution are also plotted under each mixup image. As can be seen from Figure 3, when $\lambda=0.50$, the mixup image obtains a smoother probability distribution than other settings. It means that the relative probability assigned to similar categories encodes semantic similarity between similar categories. Note that, previous conventional KD methods usually achieve this goal by setting a large temperature hyperparameter $t$.

Visualization the distribution of student networks. In this subsection, we visualize the output distribution of student from STKD and vanilla training methods on CIFAR-100. Figure 4 shows the visualization results over 100 classes and 10 sampled classes using tSNE [59] respectively. Notably, the feature space of STKD is significantly more separable than it in vanilla training manner.