PaperlayerView exact revision on arXiv

Source-linked AI summary

Knowledge Distillation with the Reused Teacher Classifier

Defang Chen, Jian-Ping Mei, Hailin Zhang, Can Wang, Yan Feng, Chun Chen

arXiv:2203.14001v1cs.CV

TL;DR

Knowledge distillation seeks to compress a powerful teacher into a smaller student while limiting performance loss, but many methods rely on elaborate representations. SimKD reuses the teacher classifier and trains the student with single-loss feature alignment, achieving consistently strong results across architectures while adding projector cost and remaining limited to supervised settings.

  • Problem

    Vanilla distillation leaves a significant teacher–student performance gap, while feature-distillation methods often require elaborate representations and carefully tuned losses.

  • Method

    SimKD reuses the pre-trained teacher classifier for student inference and trains the student encoder with a single ℓ2 feature-alignment loss, using a projector for dimension mismatches.

  • Results

    SimKD consistently outperforms compared state-of-the-art approaches across varied teacher–student architectures and generalizes to multi-teacher and data-free distillation.

  • Takeaways & Limitations

    Feature alignment with the reused teacher classifier provides a simple, comprehensible baseline for knowledge transfer without elaborate knowledge representations.

  • Takeaways & Limitations

    SimKD requires a projector when feature dimensions mismatch and is currently applicable only to supervised knowledge distillation.

Abstract

from arXiv · show

Knowledge distillation aims to compress a powerful yet cumbersome teacher model into a lightweight student model without much sacrifice of performance. For this purpose, various approaches have been proposed over the past few years, generally with elaborately designed knowledge representations, which in turn increase the difficulty of model development and interpretation. In contrast, we empirically show that a simple knowledge distillation technique is enough to significantly narrow down the teacher-student performance gap. We directly reuse the discriminative classifier from the pre-trained teacher model for student inference and train a student encoder through feature alignment with a single $\ell_2$ loss. In this way, the student model is able to achieve exactly the same performance as the teacher model provided that their extracted features are perfectly aligned. An additional projector is developed to help the student encoder match with the teacher classifier, which renders our technique applicable to various teacher and student architectures. Extensive experiments demonstrate that our technique achieves state-of-the-art results at the modest cost of compression ratio due to the added projector.

1. Introduction

The paper introduces SimKD, a simple distillation method that reuses the teacher’s classifier and aligns student features, avoiding elaborate knowledge representations. A projector handles feature-dimension mismatches, while experiments report strong performance across architectures.

  • Motivation: The method targets the significant performance gap remaining in vanilla knowledge distillation and avoids elaborate intermediate-feature representations.Prior approaches use attention maps, pairwise similarities, mutual information, or other designed representations.
  • SimKD: A single ℓ2 loss is used for feature alignment, with only the student encoder and dimension projector updated during training.
  • SimKD: SimKD reuses the pre-trained teacher classifier for student inference and trains the student through feature alignment before that classifier.The authors argue that perfect feature alignment would eliminate the teacher–student performance gap.
  • SimKD: A projector resolves mismatched teacher and student feature dimensions, making SimKD applicable to arbitrary model architectures.The projector adds less than 3% cost to the pruning ratio in teacher-to-student compression.
  • Results: SimKD consistently outperforms compared state-of-the-art approaches across teacher–student architecture combinations and extends to multi-teacher and data-free distillation.

2. Related Work

Related work improves student performance by transferring teacher outputs or intermediate representations, but often requires specialized representations and careful loss balancing. SimKD instead emphasizes the overlooked discriminative information in the teacher classifier.

  • Knowledge Distillation: Knowledge distillation transfers a powerful teacher’s knowledge into a smaller student model, commonly using teacher predictions or soft targets.
  • Feature Distillation: Feature distillation methods transmit information from intermediate teacher–student layers through feature maps, transformations, relations, contrastive learning, or cross-layer associations.
  • Feature Distillation: These feature-distillation approaches require carefully designed representations and hyper-parameters to balance losses, making development labor-intensive and interpretation difficult.
  • Hypothesis Transfer Learning: Hypothesis transfer learning reuses a source-domain classifier when only limited labeled target data and no source data are available.

3. Method

SimKD simplifies distillation by reusing the teacher classifier and aligning student features with a single ℓ2 loss, avoiding elaborate knowledge representations and multiple feature-distillation losses. A projector resolves feature-dimension mismatches, while the reused classifier supports student inference and can substantially alleviate compression performance degradation.

  • Vanilla Knowledge Distillation: Vanilla KD combines cross-entropy with prediction alignment through Kullback–Leibler divergence, using temperature T to soften teacher and student distributions.The alignment term conveys information about incorrect classes, motivating a raised temperature.
  • Feature Distillation: Feature distillation extracts additional intermediate-layer gradient information but depends on designed knowledge representations and carefully tuned hyper-parameters.These requirements make development labor-intensive and complicate interpretation of each representation’s role.
  • Simple Knowledge Distillation: SimKD directly reuses the pre-trained teacher classifier for student inference, making feature alignment the only gradient-generating loss.The student is furnished with the teacher classifier rather than training a new classifier, eliminating the cross-entropy requirement for student predictions.
  • Simple Knowledge Distillation: Teacher and SimKD student features are nearly indistinguishable in t-SNE visualizations of ten sampled CIFAR-100 classes, while reused layers can trade higher accuracy for greater inference burden.The single-loss formulation is presented as interpretable, and reusing more than the final classifier is permitted.
  • Simple Knowledge Distillation: The SimKD objective minimizes the squared ℓ2 distance between teacher features and projected student features.The projector P(·) matches feature dimensions before alignment.
  • Simple Knowledge Distillation: A projector addresses mismatched teacher and student feature dimensions at less than 3% cost to the pruning ratio while supporting arbitrary architectures.In some cases, the projector plus reused classifier uses fewer parameters than the original student classifier.

4. Experiments

Experiments on CIFAR-100, ImageNet, multi-teacher, and data-free settings evaluate SimKD’s accuracy, classifier reuse, projector design, and compression cost. SimKD consistently performs strongly, while classifier reuse and projector capacity materially affect outcomes.

  • Comparison of Test Accuracy: SimKD consistently outperforms competing distillation approaches on CIFAR-100, including teacher-student pairs with similar or different architectures.For ResNet-8x4 & ResNet-32x4, it improves absolute accuracy by 3.66%, versus 1.81% for the best competitor over vanilla KD.
  • Comparison of Test Accuracy: SimKD can outperform a larger student baseline using a ResNet-110 student with a projector containing 0.05M additional parameters.The comparison uses the same ResNet-110x2 teacher and contrasts ResNet-110 with ResNet-116, which has about 0.10M more parameters.
  • Classifier-Reusing Operation Analysis: Replacing classifier reuse with joint training substantially reduces student accuracy, whether inference uses the student’s classifier or the reused teacher classifier.The two-loss joint-training strategy also requires case-by-case hyper-parameter tuning to surpass vanilla KD.
  • Classifier-Reusing Operation Analysis: Sequentially training a new student classifier after feature alignment often causes a precipitous accuracy drop, making classifier training difficult even with aligned features.Directly reusing the pretrained teacher classifier performs better than this sequential alternative.
  • Classifier-Reusing Operation Analysis: Reusing deeper teacher layers further improves student performance, although the associated model complexity increases.The results support the benefit of transferring deeper teacher information for student performance.
  • Projector Analysis: Increasing the projector bottleneck ratio raises pruning ratio but causes performance to drop because the projector’s representation ability is restricted.The paper therefore uses r = 2 as the default balance between performance and pruning ratio.
  • Projector Analysis: The added projector costs less than 1% pruning ratio in most evaluated cases, while SimKD can exceed vanilla KD’s pruning ratio in some settings.The less-than-1% cost occurs in 10 of 14 teacher-student combinations; r = 2 is selected by default.

5. Conclusion

The paper presents SimKD as a simple distillation technique that reuses the teacher classifier and trains the student with a single ℓ2 feature-alignment loss.

  • SimKD reuses the pre-trained teacher classifier for student inference and trains the student model with a single ℓ2 feature-alignment loss.

6. Limitation and Future Work

The method has projector-related complexity and supervised-learning scope limitations that motivate future work on simpler and unsupervised variants.

  • A projector is required when teacher and student feature dimensions mismatch, increasing model complexity and motivating projector-free alternatives.
  • SimKD is limited to supervised knowledge distillation, leaving successful unsupervised-learning variants for future investigation.

A.1. Datasets and Training Details

Experiments use CIFAR-100 and ImageNet with standard normalization and augmentation, while data-free distillation comparisons use reproduced baselines and specified teacher accuracy.

  • Experiments use CIFAR-100 and ImageNet, with channel-wise normalization and horizontal-flip augmentation.CIFAR-100 has 50,000 training and 10,000 test images across 100 classes; ImageNet has about 1.3 million training and 50,000 validation images across 1,000 classes.
  • Data-free distillation comparisons reproduce ZSKT, DAFL, and CMI using default model hyper-parameters.The teacher WRN-40-2 reaches 76.31% top-1 test accuracy, and a student trained with the original dataset is included for comparison.
  • All experiments use PyTorch on NVIDIA RTX 2080Ti or A40 GPU servers with CUDA 11.2 or 11.4.

A.2. Network Architectures

The experiments span diverse teacher-student architectures and use a projector for feature-dimension alignment, with added-parameter analysis and supplementary convergence results.

  • Performance evaluation covers VGG, ResNet, WRN, MobileNetV2, ShuffleNetV1, and ShuffleNetV2 teacher-student combinations.Architecture notation includes depth for VGG and ResNet, depth and width factor for WRN, and filter-scaling ratios such as ResNet-32x4.
  • The projector uses convolutional layers with batch normalization and ReLU, while r denotes the reduction ratio.When feature-map spatial dimensions differ, average pooling is applied before alignment; otherwise H and W are assumed equal.
  • The added projector parameter count is modeled as F(r) = C_t(C_s + C_t + 4)/r + 9C_t^2/r^2 + 2C_t.
  • Under mild conditions, the projector parameter count satisfies 2F(2r) < F(r) < 4F(2r).The proof relies on positive teacher and student channel dimensions and practical reduction-ratio settings.
  • SimKD achieves faster convergence: at epoch 30 it matches the baseline student at epoch 60, and at epoch 60 it surpasses the baseline at epoch 120.

B.2. Joint Training Results

Joint training results show that using only one classifier yields meaningful predictions, while the other classifier degrades to random guessing at α = 0 or α = 1.

  • B.2. Joint Training Results: At α = 0 or α = 1, only one classifier produces meaningful results, while the other degrades into random guessing.Random guesses are denoted by “–”.

B.3. Sequential Training Results

Sequential training uses a regular SGD schedule, and training a student classifier from scratch is difficult across tested learning rates. SimKD instead reaches 78.08 ± 0.15 test accuracy by reusing the pre-trained teacher classifier and aligning features with an ℓ2 loss.

  • B.3. Sequential Training Results: Sequential training uses SGD with 0.9 Nesterov momentum, 5 × 10−4 weight decay, 240 epochs, and scheduled learning-rate reductions.The initial learning rate is 0.01 for MobileNet/ShuffleNet architectures and 0.05 for other architectures, with batch size 64.
  • B.3. Sequential Training Results: Student accuracy stays near 50% for initial learning rates from 0.01 to 0.5 when training a new classifier from scratch.The result indicates difficulty in obtaining a satisfactory student classifier under these settings.
  • B.3. Sequential Training Results: 78.08 ± 0.15 test accuracy is achieved by SimKD without classifier retraining, using the pre-trained teacher classifier instead.The default feature-alignment loss is applied in the preceding layer of the teacher classifier with an ℓ2 loss; an alternative applies it in the succeeding layer.

B.5. Comparison of Pruning Ratio

SimKD generally combines strong test accuracy with a modest pruning-ratio cost from its projector. Reusing additional teacher layers can improve performance but sharply reduces the pruning ratio, while feature alignment makes teacher and student representations nearly indistinguishable in visualization.

  • B.5. Comparison of Pruning Ratio: The added projector requires less than or about 3% pruning-ratio cost while SimKD achieves state-of-the-art performance.Tables report top-1 test accuracy, pruning ratio, and the projector-to-student parameter ratio across dimension-reduction factors.
  • B.5. Comparison of Pruning Ratio: For r = 8, SimKD has higher pruning ratios than vanilla KD in the MobileNetV2x2–ResNet-32x4 and ShuffleNetV1–ResNet-32x4 settings.SimKD also attains second-best accuracy in three reported settings, with 77.49%, 78.96%, and 77.35%, respectively.
  • B.5. Comparison of Pruning Ratio: SimKD+ and SimKD++ improve performance over SimKD but sharply reduce the pruning ratio, making final-classifier reuse a balance between accuracy and parameter complexity.The comparison concerns reusing progressively more deep teacher layers.
  • B.5. Comparison of Pruning Ratio: With a simple ℓ2 loss, teacher and student features become almost indistinguishable in CIFAR-100 t-SNE visualizations.The visualization covers all 100 CIFAR-100 classes and uses magenta for teacher features and cyan for student features.
Loading 2203.14001v1…