Knowledge Distillation

Knowledge distillation article.

Motivation

Large models (ensembles, billion-parameter transformers) achieve state-of-the-art results but are expensive to deploy. Distillation lets you train a compact model that retains much of the large models accuracy, making it practical for edge devices, mobile, and low-latency serving.

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    """Hinton distillation: weighted sum of hard and soft targets."""
    soft_targets = F.softmax(teacher_logits / T, dim=-1)
    student_soft = F.log_softmax(student_logits / T, dim=-1)

    soft_loss = F.kl_div(student_soft, soft_targets, reduction="batchmean") * (T * T)
    hard_loss = F.cross_entropy(student_logits, labels)

    return alpha * soft_loss + (1 - alpha) * hard_loss

Temperature Scaling

The temperature T controls how soft the teachers output distribution is. At T=1, its the normal softmax. As T \to \infty, the distribution approaches uniform, revealing the relative similarities the teacher has learned between classes. Higher T means the student learns finer-grained inter-class structure.