PaperlayerView exact revision on arXiv

Source-linked AI summary

Knowledge Diffusion for Distillation

Tao Huang, Yuan Zhang, Mingkai Zheng, Shan You, Fei Wang, Chen Qian, Chang Xu

arXiv:2305.15712v2cs.CVcs.AI

TL;DR

Teacher–student feature gaps make knowledge distillation challenging: direct matching may disturb student optimization, and existing approaches can rely on complicated or task-specific methods. DiffKD denoises student features with a teacher-trained diffusion model, using lightweight components and adaptive noise matching. It reports consistent state-of-the-art results across image classification, object detection, and semantic segmentation.

  • Problem

    Teacher–student feature discrepancies can be large, and directly aligning mismatched features may disturb student optimization; existing methods can require complicated schemes or task-specific priors.

    Source: §1 · p. 1

  • Method

    DiffKD uses a diffusion model trained with teacher features to denoise student features, with a lightweight architecture, linear autoencoder, and adaptive noise matching.

    Source: §1 · p. 2

  • Results

    DiffKD consistently surpasses current state-of-the-art methods on image classification, object detection, and semantic segmentation; it reaches 73.62% accuracy on ImageNet with MobileNetV1 and ResNet-50, surpassing DKD by 1.57%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report that knowledge diffusion applies across intermediate features, classification outputs, and regression outputs, and generalizes across three task types.

    Source: §1 · p. 2

    DiffKD costs more computation than simple logits-distillation methods such as KD and DIST, though its cost is comparable to existing feature-distillation methods and does not affect inference computation.

    Source: §D.1 · p. 16

Abstract

from arXiv · show

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.

1 Introduction

Knowledge distillation can be hindered by the gap between teacher and student features: direct alignment may disturb student optimization, while existing approaches can require complex schemes or task-specific priors. DiffKD treats student features as noisy teacher features, denoises them with a teacher-trained diffusion model, and reports strong results across several tasks.

1 Introduction

  • Directly aligning mismatched student and teacher features can disturb student optimization, while existing approaches may require complex training schemes or task-specific priors.The introduction describes assistant teachers, student-friendly teacher training, correlation-based losses, and identifying useful receptive regions as examples.Source: §1 · p. 1
  • DiffKD treats student features as noisy versions of teacher features and uses a diffusion model trained on teacher features to denoise them for more consistent distillation.The paper reports that denoised student features are very similar to their corresponding teacher features.Source: §1 · p. 2, Figure 1 · p. 2
  • DiffKD uses two ResNet bottleneck blocks and a linear autoencoder to reduce diffusion computation, and adaptive noise matching to estimate each student feature’s noise level and initialize denoising accordingly.The student feature’s corresponding noise level is not directly known, which can weaken denoising when initialization is inexact.Source: §1 · p. 2
  • DiffKD consistently surpasses current state-of-the-art methods on image classification, object detection, and semantic segmentation; on ImageNet it reaches 73.62% accuracy with MobileNetV1 and ResNet-50, surpassing DKD by 1.57%.On Cityscapes, it outperforms MasKD by 1% with PSPNet-R18; with Swin-T and Swin-L, it reaches 82.5% ImageNet accuracy, improving the KD baseline by 1%.Source: §1 · p. 2

2 Preliminaries

Knowledge distillation minimizes discrepancies between teacher and student outputs, using a distance function suited to the output type. Diffusion models progressively add noise to data and learn to reverse that process; DiffKD applies this approach to student features.

2.1 Knowledge distillation

  • Knowledge distillation minimizes a distance between student and teacher outputs, which can be predictions or intermediate features.The distance function can be KL divergence for probabilistic outputs or MSE for intermediate features and regression outputs.Source: §2.1 · p. 3

2.2 Diffusion models

  • A diffusion model progressively adds Gaussian noise to sample data, then trains a network to predict that noise and reconstruct the sample through iterative denoising.The forward process can directly sample a noisy state at an arbitrary timestep; DDIM accelerates inference by sampling with fewer score-function evaluations.Source: §2.2 · p. 3
  • DiffKD uses a diffusion model to eliminate noise in student features for knowledge distillation.Source: §2.2 · p. 3

3 Method

DiffKD turns teacher–student feature alignment into denoising: it trains a diffusion model with teacher features and denoises student features before distillation. A linear autoencoder reduces diffusion cost, while adaptive noise matching handles variable student–teacher gaps.

3.1 Knowledge diffusion for distillation

  • DiffKD treats student features as noisier than teacher features, trains a diffusion model on teacher features, then distills the denoised student features against the original teacher features.The method reframes feature alignment as denoising to match student and teacher outputs.Source: §3 · p. 3, §3.1 · p. 4

3.2 Efficient diffusion model with linear autoencoder

  • The diffusion process can be costly for high-dimensional features, so DiffKD uses a two-bottleneck diffusion model and a linear autoencoder to compress teacher channels.The autoencoder is trained with mean-squared reconstruction loss, and its latent teacher feature is used to train diffusion and supervise the student.Source: §3.2 · pp. 4–5

3.3 Adaptive noise matching

  • Because the student–teacher gap varies across samples, adaptive noise matching learns a weight to fuse student output with Gaussian noise at a predefined starting noise level.The module’s noise adaptation is optimized through the KD loss.Source: §3.3 · p. 5
  • The overall objective combines task, diffusion, reconstruction, and KD losses; λ1 and λ2 are set to 1 in all experiments.Source: §3.3 · p. 5

4 Experiments

DiffKD is evaluated across image classification, object detection, and semantic segmentation, with ablations examining feature types, autoencoder compression, adaptive noise matching, and denoising steps. Results report improvements over prior methods across the tasks, while the ablations identify design choices balancing performance and efficiency.

4 Experiments

  • Experiments use baseline and stronger-teacher settings, and the evaluated tasks include varied detector architectures and teacher strengths.The ImageNet settings follow DIST; detection settings span two-stage, one-stage, and anchor-free detectors.Source: §4.1 · p. 6, §4.2 · p. 7

4.1 ImageNet classification

  • DiffKD surpasses DIST by 0.38% top-1 accuracy on baseline ImageNet settings with MobileNet and ResNet-50, and also outperforms DIST across stronger-teacher settings.On the stronger settings, it surpasses DIST by 0.5% with MobileNetV2; Table 3 results are described in prose but that table is not supplied.Source: §4.1 · p. 6

4.2 Object detection

  • DiffKD improves FCOS-R50 by 4.0 AP and exceeds FGD by 0.4 AP on baseline object detection settings.The attention-based MSE variant adds approximately 0.1 AP over vanilla DiffKD.Source: Table 4 · p. 7, §4.2 · p. 7

4.3 Semantic segmentation

  • DiffKD outperforms MasKD across all reported semantic-segmentation settings, including gains of 0.78% on validation and 0.65 on test for DeepLabV3-R18.Table 6 reports semantic-segmentation results on Cityscapes.Source: §4.3 · p. 8, Table 6 · p. 8

4.4 Ablation study

  • Combining feature-level and logits-level DiffKD yields the best reported ImageNet result, 73.62% top-1 accuracy for MobileNetV1 with ResNet-50.Each distillation type also improves performance without changing the loss functions.Source: Table 7 · p. 8, §4.4 · p. 8
  • A 1024-dimensional linear autoencoder offers a performance-efficiency trade-off, while compression to 128 or 256 dimensions severely weakens distillation performance.The experiments compare top-1 accuracy and diffusion-model FLOPs for MobileNetV1 and ResNet-50 on ImageNet.Source: §4.4 · pp. 8–9, Table 8 · p. 9
  • Removing adaptive noise matching lowers MobileNetV1 top-1 accuracy to 73.34%, a 0.28-point decrease from DiffKD with the module.The authors attribute the improvement to better denoised student features.Source: §4.4 · p. 9
  • Five score-function evaluations suffice for promising performance, whereas one-step denoising improves on the 72.39% MSE baseline but performs worse than larger NFE settings.The authors use five NFEs throughout for an efficiency-accuracy trade-off.Source: §4.4 · p. 9

5 Conclusion

DiffKD addresses teacher–student representation discrepancy by denoising student features with diffusion, alongside a lightweight model, linear autoencoder, and adaptive noise matching. The paper reports generalization across feature types and tasks.

5 Conclusion

  • DiffKD explicitly removes noise from student features with a diffusion model to reduce teacher–student discrepancy.The conclusion also identifies a lightweight diffusion model, linear autoencoder, and adaptive noise matching as components of the approach.Source: §5 · p. 9

A.1 Representation gap in knowledge distillation

  • Prior approaches to the representation gap include assistant teachers, student-friendly teachers, and correlation-based losses, with cited limitations in complexity or task applicability.TAKD [29], DAKD [39], SFTN [30], and DIST [17] are discussed; the text notes DIST is classification-output focused.Source: §A.1 · p. 12

5 Conclusion

  • DiffKD is presented as applicable to multiple feature types and tasks, including intermediate features, classification outputs, and regression outputs.The paper reports experiments across image classification, object detection, and semantic segmentation.Source: §5 · p. 9, §A.1 · p. 12

A.2 KD in dense prediction tasks

Dense prediction requires object-level or pixel-level predictions, making knowledge transfer more challenging than classification. Existing approaches use task-specific feature and relation strategies, motivating interest in methods that generalize across tasks.

A.2 KD in dense prediction tasks

  • Object detection and semantic segmentation require predictions for objects or pixels, making effective teacher-knowledge distillation more challenging than classification.Source: §A.2 · p. 12
  • Detection distillation methods target logits, regressions, or FPN features, while foreground-background imbalance has prompted methods that select valuable features.The cited examples include Chen et al., Mimicking, and methods [10] [18] [38] [48].Source: §A.2 · pp. 12–13
  • Segmentation methods transfer structural or relational information through latent-space features, pixel relationships, score maps, intra-class variation, channel-wise distillation, or image relations.Examples include He et al. [12], SKD, IFVD, CWD [38], and CIRKD [46].Source: §A.2 · p. 13
  • Dense-prediction KD methods are often task-specific and costly to adapt, motivating research into generalizing and unifying methods across tasks.Source: §A.2 · p. 13

B Implementation Details

Implementation uses DDIM for reverse denoising, with task-specific feature targets, losses, and autoencoder choices. The reported details distinguish classification, detection, and segmentation setups.

B.1 Diffusion model

  • All experiments use DDIM with 1,000 training timesteps and an initial denoising timestep of 500.Source: §B.1 · p. 13

B.2 Image classification

  • ImageNet classification distills backbone features and logits, using a 1,024-channel autoencoder for ResNet teachers and 768 channels for Swin-L features.The classification logits use an MLP instead of a bottleneck, and CIFAR-100 removes the autoencoder.Source: §B.2 · p. 13

B.3 Object detection

  • Object detection distills predicted feature maps alongside regression KD and task losses, without an autoencoder because FPN has 256 channels.The reported setup uses ImageNet-pretrained backbones and official 2× MMDetection strategies.Source: §B.3 · p. 13

B.4 Semantic segmentation

  • Semantic segmentation distills teacher segmentation predictions using DIST as the distance function and omits the autoencoder because Cityscapes predictions have 19 channels.Training uses standard augmentations and a polynomial learning-rate schedule.Source: §B.4 · pp. 13–14

C.1 Results on CIFAR-100 dataset

The supplementary results report DiffKD’s CIFAR-100 performance, efficient diffusion-model comparisons, adaptive noise matching, learned noise weights, and object-detection feature-type experiments.

C.1 Results on CIFAR-100 dataset

  • DiffKD surpasses previous methods in most CIFAR-100 cases, with larger gains in heterogeneous than homogeneous architecture settings.The authors say this indicates DiffKD can better handle discrepancies between teacher and student models.Source: Table 9 · p. 14, §C.1 · p. 14

C.2 Effect of efficient diffusion model

  • On COCO, the efficient diffusion model achieves similar performance to the original UNet while using fewer parameters and GFLOPs and taking one-third the training time.The authors attribute this to feature generation being easier than image generation, so a small model suffices.Source: §C.2 · p. 14

C.3 Effect of adaptive noise matching (ANM)

  • ANM improves MobileNetV2 by 1.2%, compared with 0.3% for MobileNetV1 and 0.5% for ResNet-18.The authors suggest stronger augmentation and teachers create more varied noisy gaps, where ANM is more effective at matching noise levels.Source: §C.3 · p. 14, Table 11 · p. 15

C.4 Statistics of the learned noise weight γ

  • Many learned γ values exceed 0.9, while average γ decreases as training converges.The authors interpret high γ as indicating that student features already contain substantial noise and need little added noise; cleaner samples need more. As training progresses, student features become cleaner, so γ decreases to match noise levels.Source: Figure 5 · p. 15, §C.4 · p. 15

C.5 Effects of different feature types of distillations on object detection

  • On COCO, denoising improves distillation across FPN, classification, and regression features; FPN-feature distillation performs relatively well, and combining all three is similar to FPN alone.The authors infer that FPN-feature distillation is sufficient for good performance.Source: §C.5 · p. 15

C.6 Efficiency analysis

DiffKD adds computation for denoising intermediate features and output logits, with the reported overhead positioned against other feature-based distillation methods.

C.6 Efficiency analysis

  • 800M additional FLOPs for ResNet-18 with a ResNet-34 teacher, versus 260M for CRD and 1900M for ReviewKD.The paper characterizes DiffKD’s cost as acceptable relative to other feature-based methods.Source: §C.6 · p. 16

D.1 Limitation

The authors identify implementation and training-cost limits, discuss societal impacts, and provide visualizations of predicted scores and feature maps.

D.1 Limitation

  • DiffKD uses simple convolutions and DDIM, while newer transformer-based diffusion models and more efficient inference algorithms remain available.Its computational cost exceeds simple logits distillation, is comparable to feature-distillation methods, and does not affect inference cost.Source: §D.1 · p. 16

D.2 Societal impacts

  • Evaluating DiffKD consumes substantial computing resources and may increase carbon emissions, while deploying compact models in place of heavy models could save energy.The authors say adequate validation is necessary.Source: §D.2 · p. 16

E Visualization of Predicted Classification Scores

F.1 Visualization details

  • For feature visualization, channels are averaged and softmax is applied across spatial positions, with temperature τ set to 0.5.This produces a spatial attention map for each feature map.Source: §F.1 · p. 17

F.2 More visualizations

The visualizations compare student, denoised student, and teacher features from the first stride-4 FPN output on COCO.

F.2 More visualizations

  • Denoised student features look similar to teacher features, whereas the original student features are noisier and less salient.The authors infer that distilling the denoised features can remove noise that disturbs optimization.Source: §F.2 · p. 17, Figure 7 · p. 18

Related papers

Loading 2305.15712v2…