Knowledge Distillation: Theory and Practice

Deep learning models have achieved remarkable performance across a wide range of perception and reasoning tasks, but state-of-the-art accuracy often requires extremely large and computationally expensive neural networks with hundreds of millions of parameters. As AI systems are increasingly deployed on resource-constrained edge devices like smartphones, wearables, and IoT sensors, there is a critical need for techniques to compress models while preserving performance.

Knowledge distillation has emerged as a powerful paradigm for training compact, efficient neural networks by transferring knowledge from a large "teacher" model to a smaller "student" model. Rather than training the student from scratch on the original data, we leverage the outputs of the teacher as an auxiliary training signal to guide the optimization of the student. In this post, we‘ll dive deep into the theory and practice of knowledge distillation, with a complete end-to-end code walkthrough on a real image classification problem.

Preliminaries: Model Compression

Modern neural networks often have massive computational and memory footprints. As a representative example, the famous GPT-3 language model has 175 billion parameters and training it is estimated to cost over $10 million! While such large-scale models are pushing the boundaries of AI capabilities, they are impractical to deploy for most real-world use cases.

Numerous techniques have been proposed to compress neural networks, including:

• Parameter pruning: Removing redundant/non-essential weights and neurons
• Low-rank factorization: Decomposing weight matrices into low-rank approximations
• Weight quantization: Clustering weights into discrete bins to store them with fewer bits
• Knowledge distillation: Training a student model on teacher outputs rather than raw data

Compared to the other model compression approaches, knowledge distillation has a key advantage – it is architecture-agnostic and can be applied to any teacher and student model. Pruning, factorization, and quantization all require directly manipulating the weights and activations which makes them tricky to apply to complex architectures. Knowledge distillation only requires access to the inputs, outputs, and optionally intermediate activations of the teacher, without needing to modify its internal structure.

Knowledge Distillation 101

The core idea behind knowledge distillation is straightforward: we want to train a student model to solve the task while encouraging its outputs to be similar to those of the teacher model. We assume the teacher has already been trained to have high accuracy, so mimicking its predictions should lead to good student performance.

Mathematically, the standard supervised learning loss for classification is the cross-entropy between the model‘s predicted class probabilities and the true target labels:

$$ \mathcal{L}{CE} = – \sum{i=1}^N \sum{c=1}^C y{i,c} \log p_{i,c} $$

$N$ is the number of examples, $C$ is the number of classes, $y{i,c} \in {0,1}$ indicates if example $i$ belongs to class $c$, and $p{i,c}$ is the predicted probability that example $i$ is in class $c$.

For knowledge distillation, in addition to the standard supervised loss, we introduce a "distillation loss" to encourage the student‘s predictions $p^S$ to match the teacher‘s predictions $p^T$. The most common choice is to use the Kullback-Leibler divergence:

$$ \mathcal{L}{KD} = \sum{i=1}^N\sum{c=1}^C p^T{i,c} \log \frac{p^T{i,c}}{p^S{i,c}} $$

The total loss for the student is a weighted sum of the standard supervised loss and the distillation loss:

$$ \mathcal{L} = (1-\alpha)\mathcal{L}{CE} + \alpha T^2 \mathcal{L}{KD} $$

$\alpha \in [0,1]$ controls the relative importance of the two loss terms, and $T$ is a "temperature" parameter that softens the output distributions. When $T=1$, the outputs are the standard softmax probabilities, while higher values produce more uniform distributions.

Intuitively, the distillation loss acts as a regularizer, constraining the student to learn representations that are predictive of the same features as the teacher. The temperature provides a mechanism to emphasize the relative probabilities of the incorrect classes. By softening the distributions, more information is preserved about which classes the teacher finds similar to the correct class.

In the Wild: Distilling Models for Chest X-Ray Diagnosis

To make things concrete, let‘s see how knowledge distillation can be applied to a real-world medical image classification problem. The task is to predict whether a chest X-ray image shows signs of pneumonia, an infection of the lungs. Distinguishing between normal and diseased tissue can be quite subtle, so having an accurate yet efficient model for screening is highly valuable.

We‘ll use the ChestX-ray14 dataset which contains 112,120 frontal-view chest X-rays from 30,805 unique patients. The images are labeled as either normal (no pneumonia) or one of 14 different types of lung pathologies including pneumonia. We‘ll focus on the binary pneumonia detection task of normal vs. pneumonia.

Teacher Model

For the teacher model, we‘ll use a standard ResNet-50 architecture pretrained on ImageNet and fine-tuned on the ChestX-ray14 data. The model is trained with a batch size of 64 for 50 epochs using Adam with an initial learning rate of 0.001. Data augmentation including random horizontal flips, rotations, and contrast and brightness shifts is applied to reduce overfitting.

On the held-out test set of 22,424 images (20% of the total data), the fine-tuned ResNet-50 teacher achieves an AUC of 0.887, accuracy of 82.1%, sensitivity (recall) of 80.4%, and specificity of 84.2%. Not bad for an off-the-shelf model with minimal task-specific tuning! However, with 25.6 million parameters and 8.2 billion FLOPs per inference, this model is still far too large and expensive to deploy in many real-world clinical settings.

Student Model

For the student, we use a lightweight residual CNN with only 570K parameters – a 45x reduction from the teacher. The architecture consists of:

• 3×3 conv, 16 filters
• Residual block: 3×3 conv, 16 filters → ReLU → 3×3 conv, 16 filters
• Max pool, /2
• Residual block: 3×3 conv, 32 filters → ReLU → 3×3 conv, 32 filters
• Max pool, /2
• Residual block: 3×3 conv, 64 filters → ReLU → 3×3 conv, 64 filters
• Global average pool
• FC, 2 outputs

We train this student model with knowledge distillation using a temperature of T=10 and loss weighting α=0.7. Prior work has shown that relatively high temperatures and distillation loss weights tend to work best for knowledge distillation.

After 50 epochs of distillation, the student model achieves an AUC of 0.861, accuracy of 80.5%, sensitivity of 78.1%, and specificity of 83.3% on the test set. This is impressively close to the teacher‘s performance, despite the drastically reduced model size! For comparison, training the same student architecture from scratch to convergence (without distillation) only reaches 76.2% accuracy.

To further validate the effectiveness of knowledge distillation, we perform a more thorough hyperparameter search over temperatures from 1 to 50 and α values from 0.1 to 0.9. The best result of 81.2% test accuracy is achieved with T=20 and α=0.5, demonstrating the value of tuning these parameters for the task at hand.

We also visualize the learned representations of the student and teacher models using t-SNE plots of their penultimate layer activations. Interestingly, while the decision boundaries are fairly similar, the student‘s representations are somewhat more spread out, which could indicate that distillation has encouraged greater class separation compared to training from scratch.

Reflections and Looking Ahead

Knowledge distillation has become an indispensable technique in the deep learning toolbox, empowering practitioners to reap the benefits of extremely high-capacity models without the burdensome computational costs. Through our theoretical exposition and practical case study on medical imaging, we‘ve seen how distillation can massively compress models while preserving performance. A few key takeaways:

1. Knowledge distillation is a general-purpose method that can be flexibly applied to a wide range of teacher and student model architectures, with minimal manual engineering required beyond tuning hyperparameters.

2. Distillation is particularly powerful for deploying deep learning in real-world applications with compute, memory, power, or latency constraints. Edge devices and cost-sensitive domains like healthcare are prime candidates.

3. Temperature and loss weighting hyperparameters have a significant effect on knowledge distillation performance and should be carefully tuned for best results. Higher temperatures and balanced loss terms tend to work well.

4. Distillation can produce student models that are more robust and generalize better than ones trained from scratch, by leveraging the inductive biases and dark knowledge learned by the teacher.

Looking ahead, there are many exciting opportunities for future research on knowledge distillation:

• Theoretical analysis on the optimization and generalization properties of distillation, e.g. analyzing the dynamics of the teacher and student networks over the course of training.

• Experiments with more sophisticated and dynamic distillation strategies like curriculum learning, where the temperature and loss weights change over time, or meta-learning the optimal hyperparameters.

• Distillation with more than two models, e.g. teacher-assistant-student hierarchies or co-distillation between multiple students.

• Applying knowledge distillation to other domains like natural language processing, speech recognition, and reinforcement learning.

• Combining knowledge distillation with other techniques like semi-supervised learning, few-shot learning, federated learning, and differential privacy for more data-efficient and secure model training.

As AI systems become increasingly ubiquitous, knowledge distillation will be a key enabler for democratizing access to state-of-the-art performance beyond the data centers of large tech companies. We‘re excited to see what the coming years will bring as researchers and practitioners continue to push the boundaries of efficient, lightweight deep learning!