[CVPR 2024 Highlight] Logit Standardization in Knowledge Distillation

[Project Page] [arXiv] [Supplementary Materials] [Zhihu (in Chinese)]

Vanilla KD	KD w/ our logit standardization

Abstract

Knowledge distillation involves transferring soft labels from a teacher to a student using a shared temperature-based softmax function. However, the assumption of a shared temperature between teacher and student implies a mandatory exact match between their logits in terms of logit range and variance. This side-effect limits the performance of student, considering the capacity discrepancy between them and the finding that the innate logit relations of teacher are sufficient for student to learn. To address this issue, we propose setting the temperature as the weighted standard deviation of logit and performing a plug-and-play Z-score pre-process of logit standardization before applying softmax and Kullback-Leibler divergence. Our pre-process enables student to focus on essential logit relations from teacher rather than requiring a magnitude match, and can improve the performance of existing logit-based distillation methods. We also show a typical case where the conventional setting of sharing temperature between teacher and student cannot reliably yield the authentic distillation evaluation; nonetheless, this challenge is successfully alleviated by our Z-score. We extensively evaluate our method for various student and teacher models on CIFAR-100 and ImageNet, showing its significant superiority. The vanilla knowledge distillation powered by our pre-process can achieve favorable performance against state-of-the-art methods, and other distillation variants can obtain considerable gain with the assistance of our pre-process.

:tada: News

2024.4: Selected as Highlight in CVPR 2024

2024.3: Release the code and arXiv

2024.2: Accepted by CVPR 2024

2023.7: Rejected by ICCV 2023

Usage

The code is built on mdistiller, Multi-Level-Logit-Distillation, CTKD and tiny-transformers.

Installation

Environments:

• Python 3.8
• PyTorch 1.7.0

Install the package:

sudo pip3 install -r requirements.txt
sudo python setup.py develop

Distilling CNNs

CIFAR-100

• Download the cifar_teachers.tar and untar it to ./download_ckpts via tar xvf cifar_teachers.tar.

1. For KD

   # KD
   python tools/train.py --cfg configs/cifar100/kd/resnet32x4_resnet8x4.yaml
   # KD+Ours
   python tools/train.py --cfg configs/cifar100/kd/resnet32x4_resnet8x4.yaml --logit-stand --base-temp 2 --kd-weight 9 

2. For DKD

   # DKD
   python tools/train.py --cfg configs/cifar100/dkd/resnet32x4_resnet8x4.yaml 
   # DKD+Ours
   python tools/train.py --cfg configs/cifar100/dkd/resnet32x4_resnet8x4.yaml --logit-stand --base-temp 2 --kd-weight 9 

3. For MLKD

   # MLKD
   python tools/train.py --cfg configs/cifar100/mlkd/resnet32x4_resnet8x4.yaml
   # MLKD+Ours
   python tools/train.py --cfg configs/cifar100/mlkd/resnet32x4_resnet8x4.yaml --logit-stand --base-temp 2 --kd-weight 9 

4. For CTKD

Please refer to CTKD.

Results and Logs

We put the training logs in ./logs and hyper-linked below. The name of each log file is formated with KD_TYPE,TEACHER,STUDENT,BASE_TEMPERATURE,KD_WEIGHT.txt. The possible third value for DKD is the value of BETA. Due to average operation and randomness, there may be slight differences between the reported results and the logged results.

1. Teacher and student have identical structures:

2. Teacher and student have distinct structures:

Training on ImageNet

• Download the dataset at https://image-net.org/ and put it to ./data/imagenet

  # KD
  python tools/train.py --cfg configs/imagenet/r34_r18/kd.yaml
  # KD+Ours
  python tools/train.py --cfg configs/imagenet/r34_r18/kd.yaml --logit-stand --base-temp 2 --kd-weight 9 

Distilling ViTs

Please refer to tiny-transformers.

Results

Reproduce Figure 3, 4, 5

Please refer to visualizations.

Acknowledgement

Sincere gratitude to the contributors of mdistiller, CTKD, Multi-Level-Logit-Distillation and tiny-transformers for their distinguished efforts.

:mailbox_with_mail: Contact

Shangquan Sun: shangquansun@gmail.com

:mega: Citation

If you find that this project helps your research, please consider citing some of the following paper:

@inproceedings{sun2024logit,
  title={Logit standardization in knowledge distillation},
  author={Sun, Shangquan and Ren, Wenqi and Li, Jingzhi and Wang, Rui and Cao, Xiaochun},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={15731--15740},
  year={2024}
}

@article{sun2024logit,
  title={Logit Standardization in Knowledge Distillation},
  author={Sun, Shangquan and Ren, Wenqi and Li, Jingzhi and Wang, Rui and Cao, Xiaochun},
  journal={arXiv preprint arXiv:2403.01427},
  year={2024}
}