Memory-Replay Knowledge Distillation

Abstract

Knowledge Distillation (KD), which transfers the knowledge from a teacher to a student network by penalizing their Kullback-Leibler (KL) divergence, is a widely used tool for Deep Neural Network (DNN) compression in intelligent sensor systems. Traditional KD uses pre-trained teacher, while self-KD distills its own knowledge to achieve better performance. The role of the teacher in self-KD is usually played by multi-branch peers or the identical sample with different augmentation. However, the mentioned self-KD methods above have their limitation for widespread use. The former needs to redesign the DNN for different tasks, and the latter relies on the effectiveness of the augmentation method. To avoid the limitation above, we propose a new self-KD method, Memory-replay Knowledge Distillation (MrKD), that uses the historical models as teachers. Firstly, we propose a novel self-KD training method that penalizes the KD loss between the current model's output distributions and its backup outputs on the training trajectory. This strategy can regularize the model with its historical output distribution space to stabilize the learning. Secondly, a simple Fully Connected Network (FCN) is applied to ensemble the historical teacher's output for a better guidance. Finally, to ensure the teacher outputs offer the right class as ground truth, we correct the teacher logit output by the Knowledge Adjustment (KA) method. Experiments on the image (dataset CIFAR-100, CIFAR-10, and CINIC-10) and audio (dataset DCASE) classification tasks show that MrKD improves single model training and working efficiently across different datasets. In contrast to the existing fancy self-KD methods with various external knowledge, the effectiveness of MrKD sheds light on the usually abandoned historical models during the training trajectory.

Keywords: Deep Neural Network; Fully Connected Network; Knowledge Adjustment; audio classification; image classification; self-knowledge distillation; training trajectory.

Figure 1

The framework of our proposed memory replay Knowledge Method with Fully Connected Network and Knowledge Adjustment.

Figure 2

Simplified graphical illustration for different self-knowledge distillation methods. (a) Class-wise self-Knowledge Distillation [19]; (b) self-Knowledge Distillation [21]; (c) memory replay Knowledge Distillation.

Figure 3

The framework of MrKD-plain without Fully Connected Network (FCN) ensemble and Knowledge Adjustment (KA).

Figure 4

The framework of the Fully Connected Network.

Figure 5

Knowledge Adjustment of a wrong probability offered by an imperfect teacher. The distribution is from a sample of the CIFAR-10 training dataset, whose ground truth label is ‘ship’, but the teacher’s prediction is ‘car’. Their values are exchanged.

Figure 6

Samples from the CIRFAR-100 dataset. Upper: original image. Lower: augmented image. The labels from left to right: leopard, train, fox, truck, snail, wolf, castle, and cockroach.

Figure 7

Performance on CIFAR-100 with different update frequency $\kappa$.

Figure 8

Performance on CIFAR-100 with different copy amount n.

Figure 9

Performance of MrKD on CIFAR-100 with different depth FCN.

Figure 10

Samples from the CINIC-10 dataset. Upper: inherit from the CIFAR-10 dataset. Lower: extended from ImageNet dataset. The labels from left to right: plane, ship, bird, dog, car, cat, horse, and deer.

Figure 11

The 10-s audio clips from DCASE’18 ASC [30]. Left: the raw waveform data. Right: the corresponding log Mel spectrogram. The acoustic scenes from top to down: airport, park, shopping mall, bus.

Figure 12

The 10-s audio clips from DCASE’20 Low Complexity ASC [31]. Left: the raw waveform data. Right: the corresponding log Mel spectrogram. The acoustic scenes from top to down: transportation (traveling by tram), indoor (metro station), outdoor (park).