[Paper Review] Knowledge Distillation as Semiparametric Inference

Tri Dao, Govinda M. Kamath+2 more|arXiv (Cornell University)|Apr 20, 2021

Adversarial Robustness in Machine Learning58 references4 citations

TL;DR

This paper frames knowledge distillation as a semiparametric inference problem, where the student model learns from teacher-provided class probabilities treated as plug-in estimates of unknown Bayes class probabilities. By applying semiparametric theory, the authors derive new error guarantees and propose two enhancements—cross-fitting and loss correction—that reduce the negative impact of teacher model overfitting and underfitting, leading to consistent accuracy gains on tabular and image data.

ABSTRACT

A popular approach to model compression is to train an inexpensive student model to mimic the class probabilities of a highly accurate but cumbersome teacher model. Surprisingly, this two-step knowledge distillation process often leads to higher accuracy than training the student directly on labeled data. To explain and enhance this phenomenon, we cast knowledge distillation as a semiparametric inference problem with the optimal student model as the target, the unknown Bayes class probabilities as nuisance, and the teacher probabilities as a plug-in nuisance estimate. By adapting modern semiparametric tools, we derive new guarantees for the prediction error of standard distillation and develop two enhancements -- cross-fitting and loss correction -- to mitigate the impact of teacher overfitting and underfitting on student performance. We validate our findings empirically on both tabular and image data and observe consistent improvements from our knowledge distillation enhancements.

Motivation & Objective

• To explain why knowledge distillation often outperforms direct student training despite using a teacher model's soft labels.
• To formalize knowledge distillation as a semiparametric inference problem with unknown Bayes class probabilities as the target and teacher probabilities as a nuisance estimate.
• To develop theoretically grounded enhancements that mitigate performance degradation due to teacher model overfitting and underfitting.
• To validate the proposed methods empirically across tabular and image classification benchmarks.

Proposed method

• The authors model knowledge distillation as a semiparametric inference problem, treating the optimal student model as the target parameter and the true Bayes class probabilities as a nonparametric nuisance component.
• Teacher model outputs are used as plug-in estimates for the nuisance component, enabling the derivation of asymptotic bias and variance bounds for student prediction error.
• The authors introduce cross-fitting to reduce bias from overfitting in the teacher model by using data splitting to estimate nuisance parameters.
• Loss correction is proposed to adjust the student's training objective to account for potential underfitting or miscalibration in the teacher model.
• Theoretical guarantees are derived using modern semiparametric tools, including influence functions and efficient estimation theory.
• The framework allows for error decomposition and provides a principled way to assess and improve student generalization beyond standard distillation.

Experimental results

Research questions

• RQ1 Why does knowledge distillation with a teacher model often yield better student performance than direct supervised training?
• RQ2 How can the impact of teacher model overfitting and underfitting be formally quantified and mitigated in knowledge distillation?
• RQ3 Can semiparametric inference theory be applied to derive new, more robust knowledge distillation methods?
• RQ4 What theoretical guarantees can be established for the prediction error of student models trained via knowledge distillation?

Key findings

• The proposed cross-fitting and loss correction methods consistently improve student model accuracy across both tabular and image datasets by reducing bias from teacher model errors.
• Theoretical analysis shows that standard knowledge distillation can inherit bias from an overfitted teacher, and the new methods mitigate this through proper estimation of nuisance parameters.
• Loss correction effectively reduces the impact of underfitting in the teacher model by adjusting the student's training objective to better align with the true Bayes decision boundary.
• Empirical results demonstrate that the enhanced distillation methods outperform standard knowledge distillation, especially when the teacher model is poorly calibrated or overfitted.