[Paper Review] Knowledge Distillation: Bad Models Can Be Good Role Models

Gal Kaplun, Eran Malach+2 more|arXiv (Cornell University)|Mar 28, 2022

Machine Learning and Algorithms6 citations

TL;DR

This paper proposes that overparameterized neural networks, despite poor generalization, can serve as effective knowledge distillation teachers by acting as conditional samplers of label noise. Theoretically, distillation from such 'bad' models guarantees the student network approximates the Bayes optimal classifier, even when the teacher is a poor predictor.

ABSTRACT

Large neural networks trained in the overparameterized regime are able to fit noise to zero train error. Recent work \citep{nakkiran2020distributional} has empirically observed that such networks behave as "conditional samplers" from the noisy distribution. That is, they replicate the noise in the train data to unseen examples. We give a theoretical framework for studying this conditional sampling behavior in the context of learning theory. We relate the notion of such samplers to knowledge distillation, where a student network imitates the outputs of a teacher on unlabeled data. We show that samplers, while being bad classifiers, can be good teachers. Concretely, we prove that distillation from samplers is guaranteed to produce a student which approximates the Bayes optimal classifier. Finally, we show that some common learning algorithms (e.g., Nearest-Neighbours and Kernel Machines) can generate samplers when applied in the overparameterized regime.

Motivation & Objective

• To understand why overparameterized models generalize well despite fitting training noise perfectly.
• To investigate the theoretical properties of models that act as conditional samplers of label noise.
• To explore whether such noisy models can serve as effective teachers in knowledge distillation.
• To establish a theoretical guarantee that distillation from these samplers leads to Bayes-optimal student models.

Proposed method

• Theoretical analysis of overparameterized models as conditional samplers of label noise in the training distribution.
• Formalization of knowledge distillation from such samplers using soft labels on unlabeled data.
• Proof that distillation from conditional samplers converges to the Bayes optimal classifier under mild assumptions.
• Identification of common learning algorithms (e.g., Nearest-Neighbours, Kernel Machines) that generate such samplers in the overparameterized regime.
• Use of distributional generalization theory to analyze the behavior of these models on unseen data.
• Derivation of generalization bounds showing the student network's convergence to the Bayes risk when trained via distillation from a sampler.

Experimental results

Research questions

• RQ1 Can overparameterized models that perfectly fit noisy training data still serve as useful teachers in knowledge distillation?
• RQ2 What theoretical properties do these noisy models exhibit in terms of generalization and sampling behavior?
• RQ3 Does knowledge distillation from such models lead to improved generalization in the student network?
• RQ4 Which standard learning algorithms naturally produce conditional samplers in the overparameterized regime?
• RQ5 Is there a theoretical guarantee that distillation from a sampler yields a student close to the Bayes optimal classifier?

Key findings

• Overparameterized models trained on noisy data behave as conditional samplers, replicating label noise on unseen inputs.
• Despite poor predictive accuracy, such models are effective teachers in knowledge distillation.
• Distillation from these samplers guarantees that the student network approximates the Bayes optimal classifier.
• Common algorithms like Nearest-Neighbours and Kernel Machines generate conditional samplers when applied in the overparameterized regime.
• Theoretical analysis confirms that distillation from samplers leads to generalization performance close to the Bayes risk.