Distilling the Knowledge in a Neural Network
📜 Abstract
A very simple way to improve the performance of almost any machine learning algorithm is to train many different models on the same data and then to average their predictions [3]. Unfortunately, making predictions using a whole ensemble of models is cumbersome and may be too computationally expensive to allow deployment to a large number of users, especially if the individual models are large neural nets. Caruana and his collaborators [1] have shown that it is possible to compress the knowledge in an ensemble into a single model which is much easier to deploy and we develop this approach further using a different compression technique. We achieve some surprising results on MNIST and we show that we can significantly improve the acoustic model of a heavily used commercial system by distilling the knowledge in an ensemble of models into a single model. We also introduce a new type of ensemble composed of one or more full models and many specialist models which learn to distinguish fine-grained classes that the full models confuse. Unlike a mixture of experts, these specialist models can be trained rapidly and in parallel.
✨ Summary
The paper introduces knowledge distillation, a method for transferring information from a large or ensemble teacher model to a smaller student model by training the student to match the teacher’s softened output probabilities. The temperature parameter exposes relative probabilities among incorrect classes, which convey information about class similarities and the teacher’s generalization behavior. When hard labels are also available, the authors combine hard-label and soft-target cross-entropy objectives, scaling the soft-target gradients by the square of the temperature. In the high-temperature limit, this objective is approximately equivalent to matching the teacher’s logits after per-example mean-centering.
Empirically, distillation substantially improves compact models on MNIST, transfers most of the ensemble benefit to a single acoustic model in speech recognition, and acts as a strong regularizer: a speech model trained with soft targets on only 3% of the original data nearly recovers the test performance of a model trained on the full dataset. The paper also proposes specialist models trained on clusters of frequently confused classes. These specialists can be trained independently and combined with a generalist model; on Google’s JFT image dataset, 61 specialists improve top-1 accuracy from 25.0% to 26.1%.
The work became a foundational reference for subsequent knowledge-distillation research. Kim and Rush explicitly extend the teacher–student approach to neural machine translation through sequence-level distillation, reporting a student that runs 10 times faster than its teacher with limited performance loss. (arxiv.org) Later surveys characterize knowledge distillation as a major model-compression and acceleration direction for deploying neural networks on resource-constrained devices. (arxiv.org) Subsequent research has applied the framework beyond direct output matching, including symbolic distillation in which large language models generate structured training data for smaller models. (arxiv.org) The paper’s bibliographic record and original publication metadata are corroborated by arXiv, Google Research, and DBLP. (arxiv.org)