Paper: "Distilling the Knowledge in a Neural Network" (Hinton et al., 2015) ArXiv: https://arxiv.org/abs/1503.02531 Key Insight: Don't train a small model directly. Use a large trained model (teacher) to guide a small model (student) via soft targets. Student learns faster and better than training alone.
Soft Targets vs Hard Targets
Hard Target (standard):
Teacher output: [0.7, 0.2, 0.1]
Student should predict: class 0
Loss = -log(student_prob[0])
Soft Target (distillation):
Teacher output: [0.7, 0.2, 0.1]
Student should match teacher exactly
Loss = KL(teacher_probs || student_probs)
= sum(teacher_prob[i] * log(teacher_prob[i] / student_prob[i]))
Advantage: Student learns the full distribution, not just class
Temperature
Soft targets need temperature to smooth teacher outputs:
Hard output: softmax(logits / T=1) = [0.7, 0.2, 0.1]
Soft output: softmax(logits / T=4) = [0.6, 0.25, 0.15]
(smoother, more information for student)
Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
class DistillationLoss(nn.Module):
def __init__(self, temperature=4.0, alpha=0.5):
super().__init__()
self.temperature = temperature
self.alpha = alpha # weight for distillation loss
def forward(self, student_logits, teacher_logits, targets):
"""
student_logits: logits from student model
teacher_logits: logits from teacher model (detached)
targets: ground truth labels
"""
# Hard target loss (standard CE)
hard_loss = F.cross_entropy(student_logits, targets)
# Soft target loss (KL divergence)
student_soft = F.softmax(student_logits / self.temperature, dim=-1)
teacher_soft = F.softmax(teacher_logits / self.temperature, dim=-1)
soft_loss = F.kl_div(
F.log_softmax(student_logits / self.temperature, dim=-1),
teacher_soft,
reduction='batchmean'
)
# Combined loss
loss = self.alpha * hard_loss + (1 - self.alpha) * soft_loss
return loss
# Training
teacher = load_model('large_model') # 175B
student = load_model('small_model') # 7B
teacher.eval() # Freeze teacher
distill_loss = DistillationLoss(temperature=4.0, alpha=0.3)
optimizer = torch.optim.Adam(student.parameters(), lr=1e-4)
for epoch in range(10):
for images, labels in dataloader:
# Teacher forward (no grad)
with torch.no_grad():
teacher_logits = teacher(images)
# Student forward
student_logits = student(images)
# Distillation loss
loss = distill_loss(student_logits, teacher_logits, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Benchmarks
ImageNet Classification (ResNet-50 as teacher)
Training method | Model | Accuracy
Standard training | 18M | 68.8%
Standard training | 50M | 76.1%
Distillation from ResNet-50| 18M | 73.8% (5% improvement!)
Distillation from ResNet-50| 50M | 77.5% (1.4% improvement!)
Distillation helps both small and large models!
Why It Works
Teacher has learned a smooth decision boundary:
- Correctly classified hard examples
- Also knows which wrong classes are "near" correct
- This information is in soft probabilities
Student learns from both:
1. Hard targets: what is correct
2. Soft targets: what is almost correct
Result: Student generalizes better
Extensions
1. Feature Distillation
- Match hidden layers, not just outputs
- More detailed knowledge transfer
2. Attention Transfer
- Match attention maps
- Student learns what to focus on
3. Relation-based Distillation
- Match relationships between examples
- Higher-order knowledge transfer
Our Analysis: Why Distillation Matters
Knowledge distillation is practical magic: take a expensive large model, compress it with minimal accuracy loss. This is how models get deployed on edge devices. Modern transformers often use distillation:
- BERT → DistilBERT (40% smaller, 60% faster)
- GPT-3 → smaller models for inference
References
- Paper: Distilling the Knowledge in a Neural Network (Hinton et al., 2015)
- Extensions: FitNet, Attention Transfer, RKD
Conclusion
Knowledge distillation enables model compression without sacrificing quality. By learning from soft targets, students gain access to the teacher's knowledge beyond classification accuracy. This principle drives production deployment of large models. Next: we'll explore model merging and ensemble techniques.