Distilling Knowledge Into Smaller Models

Summary & Methodology Analysis

Ensembles of models or very large neural networks are computationally expensive and cumbersome to deploy to a large number of users, while transferring knowledge from large models or ensembles to smaller deployment-ready models requires overcoming the limitations of standard hard-target training. Standard mechanics involve computing class probabilities using a softmax output layer that converts logits into probabilities based on a temperature parameter T, and training models by minimizing cross-entropy between predictions and hard target labels or using an ensemble by averaging predictions from multiple models. The novel contribution introduces distillation, which trains a smaller distilled model using a high temperature T in the softmax to produce and match soft target distributions, thereby transferring generalization ability from cumbersome models.

To implement this, the method combines objective functions using a weighted average of cross-entropy with soft targets at high temperature T and cross-entropy with correct hard labels at temperature 1, scaling gradients by T squared. Additionally, the authors introduce specialist models trained on data enriched in confusable subsets of classes by clustering the covariance matrix of generalist model predictions, combining non-specialist classes into a single dustbin class, and initializing specialist weights from the generalist model. These methods are tested across datasets such as MNIST, Android voice search, JFT, and ImageNet. For example, on MNIST, a small net regularized solely by matching soft targets from a large net at temperature 20 achieved 74 test errors compared to 146 errors for an unregularized small net and 67 errors for the large net. On the speech recognition development set, the baseline achieved a frame accuracy of 58.9 percent and a WER of 10.9 percent, an ensemble of 10 models achieved 61.1 percent frame accuracy and 10.7 percent WER, and the distilled single model achieved 60.8 percent frame accuracy and 10.7 percent WER. On the JFT development set, combining the baseline with 61 specialist models yielded a 4.4 percent relative improvement in overall test accuracy, moving test accuracy from 25.0 percent to 26.1 percent.

Despite these results, the approach has limitations. The ensemble gives a smaller improvement on the ultimate objective of Word Error Rate due to a mismatch in the objective function. Furthermore, the authors have not yet shown that they can distill the knowledge in the specialists back into the single large net. The paper does not specify hardware requirements, exact parameter counts, or training costs beyond the reported dataset metrics and error rates.