Online Knowledge Distillation via Multi-branch Diversity Enhancement

Knowledge distillation[1] has been widely used in many scenarios involving deep learning algorithms, such as virtual experiments in VR, autonomous driving and so on. It provides an useful method that allows the complex teacher model to be compressed to a more lightweight student model by aligning the student model with the teacher model. When training the target model, this method takes advantage of the extra supervisory signal provided by the soft output of the teacher model. there are also many works[13, 14, 15, 27] made explorations based on this idea. In FitNets[13], the student model attempts to mimic the intermediate representation directly from the teacher network. Attention Transfer[14] transfers an attention map of a teacher model into a student and [28] proposes a similar method using mean weights. In flow-based knowledge distillation[15], the student is encouraged to mimic the teacher’s flow matrices, which are derived from the inner product between feature maps in two layers. [29] saves the computation by using singular value decomposition to compress feature maps.

There are also innovative works exploring alternatives to the usual student-teacher training paradigm. Generative Adversarial Learning[30] is proposed to generate realistic-looking images from random noise using neural networks. The ideas in the adversarial network are applied to knowledge distillation[31, 32, 33]. In MEAL[31], the generators were employed to synthesize student features and the discriminator was used to discriminate between teacher and student outputs for the same image. In [33], this work adopts adversarial method to discover adversarial samples supporting decision boundary. With the supervision of discriminator, student can better mimic the behavior of teacher model. In addtion, many works[34, 35, 36, 37] have also explored the relationship between the samples. [34] propose that similar input pairs in the teacher network tends to produce similar activations in the student network. A few recent papers[37, 38, 39] have shown that models of the same architecture can also be distilled. Snapshot distillation[39] uses the cyclic learning rate policy, in which the last snapshot of each cycle is used as the teacher for all iterations in the next cycle, and the teacher is used to provide supervision signal.

Traditional knowledge distillation methods have two stages that require a pretrained teacher model to provide soft output for distillation. Different from above complex training methods, several works adopts collaboratively training strategy. Simultaneously training a group of student models based on each other’s predictions is an effective single-stage distillation method, which can be a good substitute for pretrained teacher models. Some methods[16, 18] solve this problem. The online knowledge distillation was completed through mutual instruction between two peers[16]. However, the lack of a high-capacity teacher model will decrease the distillation efficiency. In [17, 40], each student model learns from the average of the predictions generated by a group of students and obtains a better teacher model effect. ONE found that simply averaging the results would reduce the diversity among students, affecting the training of branch-based models. ONE generates the importance score corresponding to each student through the gate module. By assigning different importance score to each branch, a high-capacity teacher model is constructed, which can leverage knowledge from training data more effectively. OKDDip[19] proposed the concept of two-level distillation. The ensemble results of auxiliary peer networks were distilled into the group leader. The diversified peer network plays a key role in improving distillation performance.

In knowledge distillation, the student uses the output of the teacher as an additional supervisory signal for network training. Given a dataset of N training samples $D=\{(x_{i},y_{i})\}^{N}_{i}$, where $y_{i}\in\{1,2,...,C\}$. Here, $x_{i}$ is the $i^{th}$ training sample, $y_{i}$ is the corresponding ground truth label and C is the total number of image classes. Take the training sample as the input of the teacher network, we will get the output logits $t_{i}=(t^{1}_{i},...,t^{c}_{i})$. The logits vector after softmax will get the $i^{th}$ probability value $q^{j}_{i}$,

where T is the temperature parameter. An increase in the parameter T will make the probability distribution smoother. When training teachers, T is set to 1. When distilling knowledge from the teacher model to the student model, T is usually set to 3.

In order to train a multi-class image classification model, our goal is to minimize the cross entropy between the predicted class probabilities $q_{i}$ and the corresponding ground truth label distribution $y_{i}$,

where $H(p,q)=-\sum_{i}p_{i}logq_{i}$.

Knowledge transfer is achieved by aligning the probability distribution $q$ generated by the student with the target distribution $t$. The temperature parameter T should be the same for teacher and student networks. Specifically, we use KL(Kullback-Leibler) Divergence as the loss function:

An overview of the Feature Fusion Module is described in Fig. 2. Features from a single layer contain less information than features from multiple layers. Many approaches[41, 42, 43, 44] try to take advantage of more diversed features to get better model performance. We take the features of the last block from multiple branches as the input of the Feature Fusion Module. Since deeper layers in the network lead to richer semantic information, this approach can enrich features with high-level semantic information. Our experiment proves that the weights generated from this method can achieve better results.

where $f(\cdot)$ denotes the function of center block in the FFM. This function will output the corresponding importance score for each branch and also satisfy $\sum_{i=1}^{m-1}f_{i}(s_{1},s_{2},...,s_{m-1})=1$. $s_{i}$ denotes the feature map of the last block from the $i^{th}$ branch. $t_{i}$ denotes the logits from the $i^{th}$ branch. $t_{e}$ denotes the weighted ensemble target.

The diversity has an important influence on the accuracy of the final ensemble results. For better results, we expect peer classifiers to classify samples based on different viewpoints. So we restrict the weight of classifiers, force them to be diversed. We use

where $W_{i}$ is the fully connected layers’ weights of peer classifiers. If the weights of fully connected layers between peers get similar, it means there are more homogenization among them. This loss function acts as a regularization term to prevent homogenization. This will force each classifier to learn different features under this limit. Experiments show that this loss function improves the diversity of peer classifiers and improves the distillation efficiency. We will explain in detail in the ablation study.

To get a better understanding of our method, we describe the process in Algorithm 1. Our distillation method is a two-level procedure. For the first level distillation, each auxiliary branch learns the knowledge distilled from the soft targets $t_{e}$ generated by FFM. The distillation loss of all auxiliary branches is

In the second-level distillation, the knowledge learned by the group will be distilled to the group leader. Same as OKDDip, we average the predictions of all branches to get $t_{avg}$. The distillation of the group leader is

To sum up, the loss function of the whole neural network is:

where $\alpha$, $\beta$ and $\gamma$ are the balance parameter to balance the loss term. The first term is the sum of all branches’ cross entropy loss.

Datasets. There are three datasets in our experiments. CIFAR-10 and CIFAR-100[25] both contains 50,000 training images and 10,000 test images, which come from 10/100 classes. CINIC-10 consists of images from both CIFAR and ImageNet[48]. It has 270,000 images and 10 classes. The size in CINIC-10 is the same as in CIFAR. It contains 90,000 training images and 90,000 test images, all at a resolution of 32 x 32. The top-1 classification error rate are reported.

Settings. We implement all the networks and training procedures in Pytorch[49]. We conduct all experiments on an NVIDIA GeFore RTX 2080Ti GPU. For all datasets, we follow the experimental setting of [19]. For data augmentation, we apply standard random crop and horizontal flip to all images. We use SGD[50] as the optimizer with Nesterov momentum 0.9 and weight decay $5e-4$ during training. We set mini-batch size to 128. We use the standard learning schedule. The learning rate starts from 0.1 and divided by 10 at 150 and 225 iterations, for a total of 300 iterations. We set m=4, means that there are three auxiliary branches and a group leader. We separate the last two blocks of each backbone network for CIFAR-10/100 and CINIC-10. We empirically set $T$=3 to generate soft predictions. We set $\alpha$=$1$, $\beta$=$2$ and $\gamma$=$5e-8$ to balance the loss term in Equation 6.

We compare our method with several online knowledge distillation methods. In OKDDip, it has two network settings: branch-based and network-based. The branch-based approach refers to student models sharing multiple convolutional layers, separated from each other after a specified layer. The network-based method means that all student models do not share any convolutional layers, and each student is an independent model. The principles of these two approaches are close, so the branch-based method can well validate the effectiveness of our method. In all the experiments, we use branch-based setting for comparison. Baseline means the original model trained on the dataset without any modification.

Table 1 and Table 2 compares the top-1 classification error rate on CIFAR-10 and CIFAR-100 based on five different backbone networks. The result generated by ONE is the averaged accuracy of all branches. The results of OKDDip and ours are the accuracy of the group leader. From these two tables, it clearly shows that our method achieves a lower error rate on the same backbone network. Specifically, our method improves the accuracy of various baseline network by 3% to 4% on CIFAR-100. The network with higher capacity generally benefits more from our method. Our methods improves the state-of-the-art methods by 0.61%, 0.49% and 0.35% with ResNet-32, ResNet-110 and ResNext-50, respectively. These results showing that our method is more effective than existing methods. When the baseline model has lower capacity, our method can also slightly improve the accuracy compared with other methods.

In Table 3, we compare our method with another two-level distillation method OKDDip on three backbone networks. The results of compared methods are the averaged ensemble results of three branches on three backbone networks in the second-level distillation. Since the ensemble results act as a teacher to teach the group leader, a more accurate result can train a better group leader. It is also seen that our method improves the OKDDip method by 0.59%, 0.57% and 0.34% with ResNet-32, ResNet-110 and ResNext-50. Generally, our method successfully enhanced the diversity among different branches and brings improvement to distillation performance.

Diversity Measurement. We use the interrater agreement in [21] as the metric to measure the branch diversity. This method is defined as:

where $T$ is the total number of classifiers, $\rho({x_{k}})$ is the number of classifiers that classify $x$ correctly, $\bar{p}$ is the average accuracy of individual classifiers and $m$ is the total number of test samples. OKDDip and our method obtained 0.633 and 0.549 respectively (CIFAR-100 & ResNet-32). The smaller the $s$ measurement, the larger the diversity. From this results, we can see that our method actually increase the branch diversity.

CINIC-10 dataset is larger and more challenging than CIFAR-10 but not as difficult as ImageNet. We adopt the same data preprocessing as those of CIFAR-10/100 experiments.

Table 4 compares the top-1 classification error rates based on three backbone networks trained by different methods. From this table, we observed that our method outperforms baseline by 1.68%, 2.13% and 1.63% on ResNet-32, ResNet-110 and ResNext-50 respectively. Our method also improves the state-of-the-art method by 0.13%, 0.35% and 0.18% on three backbone networks. We can find that the improvement in generalization performance is very limited on this dataset. High-capacity networks tend to perform better. But the accuracy of ResNext-50 is slightly lower than ResNet-110 although its baseline performance is better.

In Table 5, we compare our method with OKDDip. We can find that our method outperforms OKDDip by 0.11%, 0.37% and 0.23% on ResNet-32, ResNet-110 and ResNext-50. While it can be observed that all the methods seem not to increase as much as that in CIFAR-100 experiments. We guess it is because the homogenization problem becomes serious when we conduct experiments on easier datasets. We still need to explore solutions to solve the homogenization problem in the future.

In this section, we conduct various ablation studies to validate the effectiveness of our proposed FFM and CD loss. We use ResNet-32 on the CIFAR-100 dataset to show the benefit of our components. We also compare our FFM with other knowledge distillation methods, including gate module in ONE and self-attention(SA) mechanism in OKDDip.

In Table 6, we report the top-1 and top-5 error rates of different methods. The remaining experimental settings are consistent with previous experiment. We carefully conducted six experiments on the network components. We compared the performance of three attention modules in the same experimental settings. When FFM is used only, the performance of our method has slightly exceeded other methods. This shows that FFM makes the student network learns more knowledge during the distillation. Compared with gate module in ONE, our method improves the top-1 error rates by 0.36% and top-5 error rates by 0.2%. This result proves that our method effectively utilizes the rich semantic information of multiple branches. When we combine different attention mechanism with classifier diversification loss, our results clearly show that our method surpasses other methods. The combination of FFM and CD loss has more obvious improvement. Compared with the independent FFM, the combination improves the top-1 error rates by 0.56% and the top-5 error rates by 0.08%. Our method clearly enhances the diversity among branches and improves the generalization ability of the student model. From this table, we observe that CD loss really plays the most important role in the overall improvements.

Fig. 3 demonstrates how the performance of our method is affected by the choice of hyperparameter $\gamma$ of the CD loss. We plot the top-1 accuracy on the CIFAR-100 for ResNet-32 group leader trained with $\gamma$ ranging from $1e-10$ to $1e-4$. In this figure, the dash line indicates the mean accuracy of other methods. We can find that our method still has robust performance against varying $\gamma$ values. The green dot indicates the parameter we are using. We should note that the choice of parameters will affect the optimization process. If the parameter is too large, this will lead to too much diversity among the branches, and eventually will not converge. If the parameter is too small, the CD loss function will be difficult to play the role of regularization. In that case, the value of this loss function will be very small, making the loss function ineffective. This figure shows that CD loss has a significant effect on distillation performance within a proper range.