06. Knowledge Distillation

Chapter 6 of 18 · 15 min

KEY INSIGHT

Knowledge distillation transfers capabilities from a large model to a smaller one by training the compact model to match both the training labels and the large model's soft probability distributions. The core insight behind knowledge distillation: large models capture richer information than their predictions alone indicate. When a language model predicts "cat" with 0.7 probability and "dog" with 0.2 probability, those relative probabilities encode information about semantic similarity between categories. A compact model trained only on hard labels lacks access to this dark knowledge. The distillation training procedure uses two loss components. The first term is the standard cross-entropy against training labels, ensuring the student learns correct categories. The second term matches the student's soft probability distributions against the teacher's distributions, teaching the student the teacher's generalization behavior. ```python import torch import torch.nn as nn import torch.nn.functional as F class DistillationLoss(nn.Module): """ Combines label loss with soft target loss from teacher model. """ def __init__(self, temperature=4.0, alpha=0.5, target='soft'): super().__init__() self.temperature = temperature self.alpha = alpha # Weight for soft targets vs hard labels self.target_type = target def forward(self, student_logits, teacher_logits=None, labels=None): """ Args: student_logits: Raw logits from student model teacher_logits: Raw logits from teacher model (optional) labels: Ground truth labels for hard targets """ total_loss = 0.0 # Hard label loss (standard cross-entropy) if labels is not None: hard_loss = F.cross_entropy(student_logits, labels) total_loss += (1 - self.alpha) * hard_loss # Soft target loss from teacher if teacher_logits is not None: soft_student = F.log_softmax(student_logits / self.temperature, dim=-1) soft_teacher = F.softmax(teacher_logits / self.temperature, dim=-1) soft_loss = F.kl_div( soft_student, soft_teacher, reduction='batchmean' ) * (self.temperature ** 2) total_loss += self.alpha * soft_loss return total_loss ``` The temperature parameter controls the softness of probability distributions. Higher temperatures spread probability across more classes, magnifying differences in teacher confidence. Lower temperatures sharpen distributions toward dominant classes. Temperature values between 2 and 10 typically work best, with the specific value depending on how many classes the model distinguishes. A failure mode involves temperature-sensitive loss scaling. The KL divergence loss term must be scaled by temperature squared to compensate for the softened distributions. Forgetting this scaling results in the soft loss dominating early training and the hard loss dominating late training, destabilizing learning. Distillation does not always improve performance. If the student architecture lacks sufficient capacity to represent the teacher's knowledge, distillation cannot create information that was never present. The student must be large enough to capture the essential patterns, even if smaller than the teacher.

EXERCISE

Train a small CNN as a student model using a larger pre-trained CNN as a teacher. Compare student accuracy with and without soft target distillation. Measure whether distillation improves generalization on a held-out test set.