SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tao Huang $^{1,2}$  Yuan Zhang $^{3}$  Mingkai Zheng $^{1}$  Shan You $^{2*}$
Fei Wang<sup>4</sup> Chen Qian<sup>2</sup> Chang Xu<sup>1</sup>
<sup>1</sup>School of Computer Science, Faculty of Engineering, The University of Sydney
$^{2}$ SenseTime Research  $^{3}$ Peking University
The representation gap between teacher and student is an emerging topic in knowledge distillation (KD). To reduce the gap and improve the performance, current methods often resort to complicated training schemes, loss functions, and feature alignments, which are task-specific and feature-specific. In this paper, we state that the essence of these methods is to discard the noisy information and distill the valuable information in the feature, and propose a novel KD method dubbed DiffKD, to explicitly denoise and match features using diffusion models. Our approach is based on the observation that student features typically contain more noises than teacher features due to the smaller capacity of student model. To address this, we propose to denoise student features using a diffusion model trained by teacher features. This allows us to perform better distillation between the refined clean feature and teacher feature. Additionally, we introduce a light-weight diffusion model with a linear autoencoder to reduce the computation cost and an adaptive noise matching module to improve the denoising performance. Extensive experiments demonstrate that DiffKD is effective across various types of features and achieves state-of-the-art performance consistently on image classification, object detection, and semantic segmentation tasks. Code is available at https://github.com/hunto/DiffKD.
The success of deep neural networks is generally accomplished with the requirements of large computation and memory, which restricts their applications on resource-limited devices. One widely-used solution is knowledge distillation (KD) [14], which aims to boost the performance of efficient model (student) by transferring the knowledge of a larger model (teacher).
The key to knowledge distillation lies in how to transfer the knowledge from teacher to student by matching the output features (e.g., representations and logits). Recently, some studies [17, 29] have shown that the discrepancy between student feature and teacher feature can be significantly large due to the capacity gap between the two models. Directly aligning those mismatched features would even disturb the optimization of student and weaken the performance. As a result, the essence of most state-of-the-art KD methods is to shrink this discrepancy and only select the valuable information for distillation. For example, TAKD [29] introduces multiple middle-sized teach assistant models to bridge the gap; SFTN [30] learns a student-friendly teacher by regularizing the teacher training with student; DIST [17] relaxes the exact matching of teacher and student features of Kullback-Leibler (KL) divergence loss by proposing a correlation-based loss; MasKD [18] distills the valuable information in the features and ignores the noisy regions by learning to identify receptive regions that contribute to the task precision. However, these methods need to resort to either complicated training schemes or task-specific priors, making them challenging to apply to various tasks and feature types.
*Correspondence to: Shan You <youshan@sensetime.com>.
Figure 1: Diffusion model in DiffKD. The diffusion model is trained with teacher feature in diffusion process  $q$  (red-dashed arrow), while we feed the student feature to the reverse denoising process  $p_{\theta}$  (blue arrows) to obtain a denoised feature for distillation. We find that due to capacity limitation, student feature contains more noises and its semantic information is not as salient as the teacher's. Therefore, we treat student feature as a noisy version of teacher feature, and propose to denoise student feature using a diffusion model trained with teacher feature.
In this paper, we proceed from a different perspective and argue that the devil of knowledge distillation is in the noise within the distillation features. Intuitively, we regard the student as a noisy version of the teacher due to its limited capacity or training recipe to learn truly valuable and decent features. However, distilling knowledge with this noise can be detrimental for the student, and may even lead to undesired degradation. Therefore, we propose to eliminate the noisy information within student and distill only the valuable information accordingly. Concretely, inspired by the success of generative tasks, we leverage diffusion models [15, 40], a class of probabilistic generative models that can gradually remove the noise from an image or a feature, to perform the denoising module. An overview of our DiffKD is illustrated in Fig. 1. We empirically show that this simple denoising process can generate a denoised student feature that is very similar to the corresponding teacher feature, ensuring that our distillation can be performed in a more consistent manner.
Nevertheless, directly leveraging diffusion models in knowledge distillation has two major issues. (1) Expensive computation cost. The conventional diffusion models use a UNet-based architecture to predict the noise, and take a large amount of computations to generate high-quality images<sup>2</sup>. In DiffKD, a lighter diffusion model would suffice since we only need to denoise the student feature. We therefore propose a light-weight diffusion model consisting of two bottleneck blocks in ResNet [11]. Besides, inspired by Latent Diffusion [35], we also adopt a linear autoencoder to compress the teacher feature, which further reduces the computation cost. (2) Inexact noisy level of student feature. The reverse denoising process in diffusion requires to start from a certain initial timestep, but in DiffKD, the student feature is used as the initial noisy feature and we cannot directly get its corresponding noisy level (timestep); therefore, the inexact noisy level would weaken the denoising performance. To solve this problem, we propose an adaptive noise matching module, which measures the noisy level of each student feature adaptively and specifies a corresponding Gaussian noise to the feature to match the correct noisy level in initialization. With these two improvements, our resulting method DiffKD is efficient and effective, and can be easily implemented on various tasks.
It is worth noting that one of the merits of our method DiffKD is feature-agnostic, and the knowledge diffusion can be applied to different types of features including intermediate feature, classification output, and regression output. Extensive experimental results show our DiffKD surpasses current state-of-the-art methods consistently on standard model settings of image classification, object detection, and semantic segmentation tasks. For instance, DiffKD obtains  $73.62\%$  accuracy with MobileNetV1 student and ResNet-50 teacher on ImageNet, surpassing DKD [53] by  $1.57\%$ ; while on semantic segmentation, DiffKD outperforms MasKD [18] by  $1\%$  with PSPNet-R18 student on Cityscapes test set. Moreover, to demonstrate our efficacy in eliminating discrepancy between teacher and student features, we also implement DiffKD on stronger teacher settings that have much more advanced teacher models, and our method significantly outperforms existing methods. For example, with Swin-T student and Swin-L teacher, our DiffKD achieves remarkable  $82.5\%$  accuracy on ImageNet, improving KD baseline with a large margin of  $1\%$ .
For example, the diffusion model in Stable Diffusion [35] takes 104 GFLOPs on an input resolution of 256  $\times$  256 [33].
Figure 2: Visualization of student features, denoised student features and teacher features on COCO dataset. Details and more visualizations can be found in Appendix F.
2.1 Knowledge distillation
Conventional knowledge distillation methods transfer the knowledge of a pretrained and fixed teacher model to a student model by minimizing the discrepancies between teacher and student outputs. Typically, the outputs can be the predictions (e.g., logits in classification task) and intermediate features of model. Given the student outputs  $\pmb{F}^{(s)}$  and teacher outputs  $\pmb{F}^{(t)}$ , the knowledge distillation loss is defined as
$$
\mathcal {L} _ {\mathrm {k d}} := d \left(\boldsymbol {F} ^ {(s)}, \boldsymbol {F} ^ {(t)}\right), \tag {1}
$$
where  $d$  denotes distance function that measures the discrepancy of two outputs. For example, we can use Kullback-Leibler (KL) divergence for probabilistic outputs, and mean square error (MSE) for intermediate features and regression outputs.
2.2 Diffusion models
Diffusion models are a class of probabilistic generative models that progressively add noise to the sample data, and then learn to reverse this process by predicting and removing the noise. Formally, given the sample data  $\mathbf{z}_0 \in \mathbb{R}^{C \times H \times W}$ , the forward noise process iteratively adds Gaussian noise to it, i.e.,
$$
q \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z} _ {0}\right) := \mathcal {N} \left(\boldsymbol {z} _ {t} \mid \sqrt {\bar {\alpha} _ {t}} \boldsymbol {z} _ {0}, (1 - \bar {\alpha} _ {t}) \boldsymbol {I}\right), \tag {2}
$$
where  $\mathbf{z}_t$  is the transformed noisy data at timestep  $t \in \{0,1,\dots,T\}$ ,  $\bar{\alpha}_t \coloneqq \Pi_{s=0}^t \alpha_s = \Pi_{s=0}^t (1 - \beta_s)$  is a notation for directly sampling  $\mathbf{z}_t$  at arbitrary timestep with noise variance schedule  $\beta$  [15]. Therefore, we can express  $\mathbf{z}_t$  as a linear combination of  $\mathbf{z}_0$  and noise variable  $\epsilon_t$ :
$$
\boldsymbol {z} _ {t} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {z} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon} _ {t}, \tag {3}
$$
where  $\epsilon_t \in \mathcal{N}(\mathbf{0},\mathbf{I})$ . During training, a neural network  $\Phi_{\theta}(z_t,t)$  is trained to predict the noise in  $z_{t}$  w.r.t.  $z_0$  by minimizing the L2 loss between them, i.e.,
$$
\mathcal {L} _ {\text {d i f f}} := \left| \left| \boldsymbol {\epsilon} _ {t} - \Phi_ {\theta} \left(\boldsymbol {z} _ {t}, t\right) \right| \right| _ {2} ^ {2}. \tag {4}
$$
During inference, with the initial noise  $z_{t}$ , the data sample  $z_{0}$  is reconstructed with an iterative denoising process using the trained network  $\Phi_{\theta}$ :
$$
p _ {\theta} \left(\boldsymbol {z} _ {t - 1} \mid \boldsymbol {z} _ {t}\right) := \mathcal {N} \left(\boldsymbol {z} _ {t - 1}; \Phi_ {\theta} \left(\boldsymbol {z} _ {t}, t\right), \sigma_ {t} ^ {2} \boldsymbol {I}\right), \tag {5}
$$
where  $\sigma_t^2$  denotes the transition variance in DDIM [40], which accelerates denoising process by sampling with a small number of score function evaluations (NFEs), i.e.,  $z_{T} \rightarrow z_{T - \Delta} \rightarrow \dots \rightarrow z_{0}$  where  $\Delta$  is the sampling interval.
In this paper, we leverage a diffusion model to eliminate the noises in student feature in our knowledge distillation method DiffKD, which will be introduced in the next section.
In this section, we formulate our proposed knowledge distillation method DiffKD. We first convert the feature alignment task in KD to the denoising procedure in diffusion models, this enables us to use diffusion models to match student and teacher outputs for more accurate and effective distillation.
Figure 3: Architecture of DiffKD. Bottleneck denotes the Bottneck block in ResNet [11].
To further improve the computation efficiency, we introduce a feature autoencoder to reduce the dimensions of feature maps, thereby streamlining the diffusion process. Additionally, we propose an adaptive noise matching module to enhance the denoising performance of student feature. The architecture of DiffKD is illustrated in Fig. 3.
3.1 Knowledge diffusion for distillation
Generally, models with different capacities and architectures produce varying preferences on feature representations, even when trained on the same dataset. This discrepancy between teacher and student is crucial to the success of knowledge distillation. Previous studies [19, 21] have investigated the differences between teacher and student features. Kundu et al. [19] observes that the predicted probabilistic distribution of teacher is more sharp and confident than the student's. Similarly, ATS [21] discovers that the variance of wrong class probabilities of teacher is smaller than that of student, indicating that the teacher output is cleaner and more salient. In summary, both studies find that the student has larger values and variances on wrong classes than the teacher, suggesting that the predictions of student contains more noises than the teacher's.
In this paper, we demonstrate that the same trend holds for intermediate features. We visualize the first feature map of FPN in RetinaNet [23] on COCO dataset [24] in Fig. 2 and find that the semantic information in teacher feature is much more salient than the student feature. Therefore, we can conclude that both predictions and intermediate features of student model contain more noises than the teacher's, and these noises are difficult to eliminate through simple imitation of the teacher model in KD due to the capacity gap [17]. As a result, a proper solution is to disregard the noises and only imitate the valuable information from both teacher and student. Inspired by the success of eliminating noises in diffusion models [15, 35, 40], we propose to treat the student feature as a noisy version of teacher feature, and train a diffusion model with teacher feature, then use it to denoise the student feature. With the denoised feature that contains only valuable information as teacher feature, we can perform noiseless distillation on them.
Formally, with teacher feature  $\pmb{F}^{(tea)}$  and student feature  $\pmb{F}^{(stu)}$  used in distillation, we use  $\pmb{F}^{(tea)}$  in the forward noise process  $q(\pmb{F}_t^{(tea)}|\pmb{F}^{(tea)})$  (Eq. (2)) to train the diffusion model with  $\mathcal{L}_{\mathrm{diff}}$  (Eq. (4)). Then the student feature is fed into the iterative denoising process of the learned diffusion model, i.e.,  $p_{\theta}(F_{t - 1}^{(stu)}|F_t^{(stu)})$  in Eq. (5), where  $\pmb{F}^{(stu)}$  is the initial noisy feature of the process. After this denoising process, we obtain a denoised student feature  $\hat{\pmb{F}}^{(stu)}$ , which is used to compute the KD loss with the original teacher feature  $\pmb{F}^{(tea)}$  in Eq. (1).
3.2 Efficient diffusion model with linear autoencoder
However, we find that the denoising process in DiffKD can be computationally expensive due to the large dimensions of the teacher feature. During training, DiffKD needs to forward the noise prediction network  $\Phi_{\theta}$  for  $T$  times (we use  $T = 5$  in our method) for denoising the student feature and 1 time for training the noise prediction network with teacher feature. This  $T + 1$  times of forwarding can result in a high computation cost when the dimension of teacher feature is large. To reduce the computation cost of diffusion model, we propose a light diffusion model which is stacked with two bottleneck block in ResNet [11]. Then we follow Latent Diffusion Models [35] and propose to compress the number of channels using a linear autoencoder module. The compressed latent feature is used as the input of diffusion model. As shown in Fig. 3, our linear autoencoder module is composed of two convolutions only, one is the encoder for reducing the number of channels, another one is the decoder
Table 1: Training strategies on image classification tasks. BS: batch size; LR: learning rate; WD: weight decay; LS: label smoothing; EMA: model exponential moving average; RA: RandAugment [7]; RE: random erasing; CJ: color jitter.
Strategy	Dataset	Epochs	Total BS	Initial LR	Optimizer	WD	LS	EMA	LR scheduler	Data augmentation
A1	CIFAR-100	240	64	0.05	SGD	5 × 10-4	-	-	×0.1 at 150,180,210 epochs	crop + flip
B1	ImageNet	100	256	0.1	SGD	1 × 10-4	-	-	×0.1 every 30 epochs	crop + flip
B2	ImageNet	450	768	0.048	RMSProp	1 × 10-5	0.1	0.9999	×0.97 every 2.4 epochs	{B1} + RA + RE
B3	ImageNet	300	1024	5e-4	AdamW	5 × 10-2	0.1	-	cosine	{B2} + CJ + Mixup + CutMix
for reconstruct the teacher features. The output feature of encoder is used to train the diffusion models (Eq. (2)) and supervise the student.
The autoencoder is trained with a reconstruction loss only, which is the mean square error between the original teacher teacher  $\pmb{F}^{(tea)}$  and reconstructed teacher feature  $\tilde{\pmb{F}}_{ae}^{(tea)}$ , i.e.,
$$
\mathcal {L} _ {\mathrm {a e}} := \left\| \tilde {\boldsymbol {F}} ^ {(t e a)} - \boldsymbol {F} ^ {(t e a)} \right\| _ {2} ^ {2}. \tag {6}
$$
Note that the latent teacher feature  $Z^{(tea)}$  used to train diffusion model is detached and has no gradient backward from the diffusion model.
We also use a convolution layer to project the student feature to the same dimension as teacher latent feature  $Z^{(tea)}$ , denoted as  $Z^{(stu)}$ . It is then passed to the diffusion model to perform a reverse denoising process (Eq. (5)) and generate the denoised student feature  $\hat{Z}^{(stu)}$ . Now we have the latent representation  $Z^{(tea)}$  of teacher and the denoised representation of student  $\hat{Z}^{(stu)}$ , they are then used to compute the KD loss and supervise the student, i.e.,
$$
\mathcal {L} _ {\text {d i f f k d}} := d \left(\hat {\boldsymbol {Z}} ^ {(s t u)}, \boldsymbol {Z} ^ {(t e a)}\right). \tag {7}
$$
Note that our DiffKD is generic and applies to various tasks (e.g., classification, object detection, and semantic segmentation) and feature types (e.g., intermediate feature, classification output, and regression output). We use simple MSE loss and KL divergence loss as the distance function  $d$  to compute the discrepancy of denoised student feature and teacher feature as a baseline implementation of DiffKD, while it is possible to achieve better performance with more advanced distillation losses.
3.3 Adaptive noise matching
As previously discussed, we treat student feature as a noisy version of teacher feature. However, the noisy level, which represents the gap between the teacher and student features, is unknown and may vary depending on different training samples. Therefore, we cannot directly determine which initial timestep we should start the diffusion process. To address this issue, we introduce an adaptive noise matching module to match the noise level of student feature to a pre-defined noise level.
As shown in Fig. 3, we construct a simple convolutional module to learn a weight  $\gamma$  that fuses student output and Gaussian noise, which helps us match the student output to the same noisy level of noisy feature at initial time step  $T$ . Therefore, the initial noisy feature in the denoising process becomes
$$
\boldsymbol {Z} _ {T} ^ {(s t u)} = \gamma \boldsymbol {Z} ^ {(s t u)} + (1 - \gamma) \boldsymbol {\epsilon} _ {T}. \tag {8}
$$
This noise adaptation can be naturally optimized with the KD loss  $\mathcal{L}_{\mathrm{kd}}$ , since the optimal denoised student feature that has minimal discrepancy to the teacher feature is obtained when the student feature matches the appropriate noise level in the denoising process.
Overall loss function. The overall loss function of DiffKD is composed of the original task loss, a diffusion loss that optimize the diffusion model, a reconstruction loss to learn the autoencoder, and a KD loss for distillation on teacher features and denoised student features, i.e.,
$$
\mathcal {L} _ {\text {t r a i n}} = \mathcal {L} _ {\text {t a s k}} + \lambda_ {1} \mathcal {L} _ {\text {d i f f}} + \lambda_ {2} \mathcal {L} _ {\text {a e}} + \lambda_ {3} \mathcal {L} _ {\text {d i f f k d}}, \tag {9}
$$
where  $\lambda_1, \lambda_2$ , and  $\lambda_3$  are loss weights to balance the losses. We simply set  $\lambda_1 = \lambda_2 = 1$  in all experiments.
Table 2: Evaluation results of baseline settings on ImageNet. We use ResNet-34 and ResNet-50 released by Torchvision [28] as our teacher networks, and follow the standard training strategy (B1). MSE: we implement our baseline for comparisons.  $\dagger$ : we replace KL divergence loss with more advanced DIST loss in DiffKD.
Student (teacher)	Tea.	Stu.	KD [14]	Review [6]	DKD [53]	DIST [17]	MSE	DiffKD	DiffKD†
R18 (R34)	Top-1	73.31	69.76	70.66	71.61	71.70	72.07	70.58	72.22	72.49
Top-5	91.42	89.08	89.88	90.51	90.41	90.42	89.95	90.64	90.71
MBV1 (R50)	Top-1	76.16	70.13	70.68	72.56	72.05	73.24	72.39	73.62	73.78
Top-5	92.86	89.49	90.30	91.00	91.05	91.12	90.74	91.34	91.48
Table 3: Performance of students trained with strong strategies on ImageNet. The Swin-T is trained with strategy B3 in Table 1, others are trained with B2. The ResNet-50 is trained by TIMM [44], and Swin-L is pretrained on ImageNet-22K.
Teacher	Student	Top-1 ACC (%)
Tea.	Stu.	KD [14]	RKD [31]	SRRL [47]	DIST [17]	DiffKD
ResNet-50	ResNet-34	80.1	76.8	77.2	76.6	76.7	77.8	78.1
MobileNetV2	73.6	71.7	73.1	69.2	74.4	74.9
EfficientNet-B0	78.0	77.4	77.5	77.3	78.6	78.8
Swin-L	ResNet-50	86.3	78.5	80.0	78.9	78.6	80.2	80.5
Swin-T	81.3	81.5	81.2	81.5	82.3	82.5
In this paper, to sufficiently validate the generalization of our DiffKD, we conduct extensive experiments on image classification, object detection, and semantic segmentation tasks.
4.1 ImageNet classification
Settings. Following DIST [17], we conduct experiments on baseline settings and stronger teacher settings. The training strategies for CIFAR-100 and ImageNet datasets are summarized in Tab. 1. On baseline settings, we use ResNet-18 [11] and MobileNet V1 [16] as student models, and ResNet-34 and ResNet-50 as teachers, respectively; and the training strategy (B1) is the most common one in previous methods [6, 17, 53]. While for the stronger teacher settings, we train students with much stronger teacher models (ResNet-50 and Swin-L [26]) and strategies (B2 and B3).
We implement our DiffKD on the output feature of backbone before average pooling, and the output logits of classification head, and the distance functions are MSE and KL divergence (with a temperature factor of 1), respectively. We set  $\lambda_1 = \lambda_2 = \lambda_3 = 1$ .
Results on baseline settings. We summarized the results on baseline settings in Tab. 2. Our methods outperforms existing KD methods with a large margin, especially on the MobileNet and ResNet-50 setting, DiffKD significantly surpasses previous state-of-the-art DIST [17] by  $0.38\%$  on top-1 accuracy. For comparisons with our baseline one feature distillation, we also report the MSE results with the same distillation location as DiffKD. We can see that, by only adding our diffusion model for feature alignment, DiffKD obviously improves the MSE results by  $1.74\%$  on ResNet-18 and  $1.23\%$  on MobileNet V1. Moreover, we replace the KL divergence loss in the output logits of DiffKD with the advanced loss function DIST, which achieves further improvements. For instance, DiffKD with DIST loss improves DIST by  $0.42\%$  on ResNet-18. This indicates that our feature alignment in DiffKD is generic to different KD losses and can be further improved by changing the losses.
Results on stronger teacher settings. To fully investigate the efficacy of DiffKD on reducing the distillation gap between teacher feature and student feature, we further conduct experiments on much stronger teachers and training strategies following DIST. From the results summarized in Tab. 3, we can see that DiffKD outperforms DIST on all model settings, especially for the most light-weight MobileNetV2 in the table, DiffKD surpasses DIST by  $0.5\%$ . It is worth to remind that our DiffKD only uses the simple KL divergence loss and MSE loss on the stronger settings, the performance could be further improved if we use more advanced loss functions such as DKD [53] and DIST [17].