Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

How to Train the Teacher Model for Effective Knowledge Distillation

Shayan Mohajer Hamidi, Xizhen Deng, Renhao Tan, Linfeng Ye, Ahmed Hussein Salamah

TL;DR

This paper addresses KD by reframing the teacher’s role as providing an estimate of the true Bayes conditional probability distribution $p^*_{oldsymbol{x}}$. It proves that training a neural network with MSE loss is equivalent to minimizing the expected MSE between its output and $p^*_{oldsymbol{x}}$, arguing that an MSE-trained teacher yields a closer BCPD estimate for KD than a CE-trained one. Extensive experiments on CIFAR-100 and ImageNet show that replacing CE-trained teachers with MSE-trained teachers consistently improves student accuracy across a wide range of KD methods, including semi-supervised and binary classification settings, with gains up to around 2.6–2.7%. The results support adopting MSE-trained teachers as a plug-and-play enhancement to KD pipelines, without altering the KD loss functions or hyperparameters, though they observe a slight drop in teacher performance under MSE training.

Abstract

Recently, it was shown that the role of the teacher in knowledge distillation (KD) is to provide the student with an estimate of the true Bayes conditional probability density (BCPD). Notably, the new findings propose that the student's error rate can be upper-bounded by the mean squared error (MSE) between the teacher's output and BCPD. Consequently, to enhance KD efficacy, the teacher should be trained such that its output is close to BCPD in MSE sense. This paper elucidates that training the teacher model with MSE loss equates to minimizing the MSE between its output and BCPD, aligning with its core responsibility of providing the student with a BCPD estimate closely resembling it in MSE terms. In this respect, through a comprehensive set of experiments, we demonstrate that substituting the conventional teacher trained with cross-entropy loss with one trained using MSE loss in state-of-the-art KD methods consistently boosts the student's accuracy, resulting in improvements of up to 2.6\%.

How to Train the Teacher Model for Effective Knowledge Distillation

TL;DR

This paper addresses KD by reframing the teacher’s role as providing an estimate of the true Bayes conditional probability distribution . It proves that training a neural network with MSE loss is equivalent to minimizing the expected MSE between its output and , arguing that an MSE-trained teacher yields a closer BCPD estimate for KD than a CE-trained one. Extensive experiments on CIFAR-100 and ImageNet show that replacing CE-trained teachers with MSE-trained teachers consistently improves student accuracy across a wide range of KD methods, including semi-supervised and binary classification settings, with gains up to around 2.6–2.7%. The results support adopting MSE-trained teachers as a plug-and-play enhancement to KD pipelines, without altering the KD loss functions or hyperparameters, though they observe a slight drop in teacher performance under MSE training.

Abstract

Recently, it was shown that the role of the teacher in knowledge distillation (KD) is to provide the student with an estimate of the true Bayes conditional probability density (BCPD). Notably, the new findings propose that the student's error rate can be upper-bounded by the mean squared error (MSE) between the teacher's output and BCPD. Consequently, to enhance KD efficacy, the teacher should be trained such that its output is close to BCPD in MSE sense. This paper elucidates that training the teacher model with MSE loss equates to minimizing the MSE between its output and BCPD, aligning with its core responsibility of providing the student with a BCPD estimate closely resembling it in MSE terms. In this respect, through a comprehensive set of experiments, we demonstrate that substituting the conventional teacher trained with cross-entropy loss with one trained using MSE loss in state-of-the-art KD methods consistently boosts the student's accuracy, resulting in improvements of up to 2.6\%.
Paper Structure (22 sections, 1 theorem, 6 equations, 2 figures, 4 tables)

This paper contains 22 sections, 1 theorem, 6 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related works
  3. Knowledge distillation
  4. Training a customized teacher for KD
  5. Training DNNs using MSE
  6. Notation and Preliminaries
  7. Notation
  8. True risk Vs. empirical risk
  9. Estimating BCPD by the teacher in KD
  10. Student's generalization error and accuracy
  11. Methodology
  12. MSE loss Vs. CE loss
  13. MSE proximity Vs. CE proximity
  14. Set 1: the student's accuracy as a function of MSE($\boldsymbol{p}^*_{\boldsymbol{x}}$, $\boldsymbol{p}_{\boldsymbol{x}}$) and CE($\boldsymbol{p}^*_{\boldsymbol{x}}$, $\boldsymbol{p}_{\boldsymbol{x}}$).
  15. Set 2: empirically validating Theorem 1.
  16. ...and 7 more sections

Key Result

theorem thmcountertheorem

For $\ell= \{ \text{CE}, \text{MSE}\}$

Figures (2)

  • Figure 1: Student accuracy as a function of (left) the MSE between $\boldsymbol{p}_{\boldsymbol{x}}$ and $\Tilde{\boldsymbol{p}}_{\boldsymbol{x}}$; and (right) CE between $\boldsymbol{p}_{\boldsymbol{x}}$ and $\Tilde{\boldsymbol{p}}_{\boldsymbol{x}}$. The gray dots are noisy versions of the true $\boldsymbol{p}_{\boldsymbol{x}}$. Also, the red and blue dots represent the points corresponding to the estimates provided by MSE and CE teachers, respectively.
  • Figure 2: The student's accuracy in semi supervised distillation for CE and MSE teachers.

Theorems & Definitions (2)

  • theorem thmcountertheorem
  • proof