[Paper Review] Understanding and Improving Knowledge Distillation

Jiaxi Tang, Rakesh Shivanna+5 more|arXiv (Cornell University)|Feb 10, 2020

Machine Learning and Data Classification39 references89 citations

TL;DR

The paper decomposes knowledge distillation (KD) into three hierarchical effects—universal label smoothing, domain knowledge of class relationships, and instance-specific gradient rescaling—and introduces partial KD methods to isolate and validate each effect, with extensive empirical validation on synthetic and real datasets.

ABSTRACT

Knowledge Distillation (KD) is a model-agnostic technique to improve model quality while having a fixed capacity budget. It is a commonly used technique for model compression, where a larger capacity teacher model with better quality is used to train a more compact student model with better inference efficiency. Through distillation, one hopes to benefit from student's compactness, without sacrificing too much on model quality. Despite the large success of knowledge distillation, better understanding of how it benefits student model's training dynamics remains under-explored. In this paper, we categorize teacher's knowledge into three hierarchical levels and study its effects on knowledge distillation: (1) knowledge of the `universe', where KD brings a regularization effect through label smoothing; (2) domain knowledge, where teacher injects class relationships prior to student's logit layer geometry; and (3) instance specific knowledge, where teacher rescales student model's per-instance gradients based on its measurement on the event difficulty. Using systematic analyses and extensive empirical studies on both synthetic and real-world datasets, we confirm that the aforementioned three factors play a major role in knowledge distillation. Furthermore, based on our findings, we diagnose some of the failure cases of applying KD from recent studies.

Motivation & Objective

• Decompose KD into three knowledge sources: universal label smoothing, domain class-relationship priors, and instance-specific gradient rescaling.
• Provide theoretical analyses of how these effects improve student training and generalization.
• Propose partial KD techniques to isolate and validate each effect.
• Empirically validate effects on synthetic and real-world datasets and diagnose KD failure cases.

Proposed method

• Theoretical analysis linking KD to label smoothing and gradient rescaling with explicit equations for KD gradients.
• Introduce partial KD methods (KD-pt and KD-sim) to isolate gradient rescaling and class-relationship priors.
• Develop synthetic teacher distributions to emulate partial effects and test their impact on learning dynamics.
• Empirical evaluation on CIFAR-100, ImageNet, and PTB to compare LS, KD, and partial KD variants.
• Analyze class correlations and the effect of top-k probability preservation (KD-topk) on distillation performance.

Experimental results

Research questions

• RQ1 What are the distinct mechanisms by which KD improves student learning beyond standard label smoothing?
• RQ2 How do universal, domain, and instance-specific knowledge contribute to KD's effectiveness?
• RQ3 Can we isolate and validate each KD component with partial KD methods?
• RQ4 Under what data conditions do KD components provide the most benefit or cause failures?

Key findings

• KD’s benefits come from three levels of knowledge: a regularization effect (label smoothing), domain knowledge (class relationships shaping logit geometry), and instance-specific gradient rescaling based on teacher confidence.
• Partial KD methods (KD-pt and KD-sim) can isolate these effects and, when combined (KD-pt+sim), can approach or surpass standard KD performance on some datasets.
• On synthetic data, KD-pt dominates when classes are uncorrelated, while KD-sim adds value as class correlations increase; combining them yields further gains.
• On CIFAR-100 and ImageNet, KD generally improves student performance beyond LS, with KD-topk often providing additional gains by preserving only the most informative class relationships and reducing noise.
• The study identifies scenarios where LS can hinder KD due to loss of class-relationship information and skewed ground-truth distributions, offering diagnostic guidance for KD failures.