Model Distillation

Overview

Model Distillation (Knowledge Distillation) trains a small "student" model to mimic a large "teacher" model, transferring knowledge through softer probability distributions. Result: Smaller models with comparable quality to large models, but 10-100x faster inference.

• Foundational Paper: "Distilling the Knowledge in a Neural Network" (Hinton et al., 2015)
• Key Insight: Soft targets (teacher outputs) are better for training than hard targets
• Trade-off: 5-10% quality loss for 10-100x faster inference
• Adoption: DistilBERT (40% smaller, 60% faster), MobileBERT, student LLMs
• Modern Use: Common for deployment-critical scenarios

-

The Problem: Large Models Are Slow

Model Size vs Speed Trade-off

Model Parameters Memory FP32 Speed Int8 Speed Use Case
──────────────────────────────────────────────────────────────────────
GPT-3 175B 350GB Very slow Slow Research
Llama 70B 70B 140GB Very slow Slow Server
Llama 7B 7B 14GB Slow OK Desktop
DistilBERT 66M 130MB Fast Very fast Mobile
TinyBERT 14M 28MB Very fast Fastest Edge

Inference speed (tokens/sec):

Model GPU Mobile (4C) Edge Device
───────────────────────────────────────────────────────
Llama 70B 50 — —
Llama 7B 200 2 t/s —
DistilBERT 2000 100 t/s 10 t/s
TinyBERT 5000 500 t/s 50 t/s

Problem:
 - Users want quality of 70B but speed of 7B
 - Simple scaling down loses quality dramatically
 - Solution: Distillation preserves quality while compressing!

-

How Model Distillation Works

Core Concept

Standard supervised learning:
- ┌─────────────────────────┐
 - Large training data │
 - (hard labels: 0 or 1) │
 - ┬────────────┘
 ↓
- ┌────────────────┐
 - Train model │
 - ┘
 ↓
- ┌────────────────┐
 - Learned model │
 - ┘

Problem:
- Hard labels contain little information
- Model must infer patterns from data alone

Knowledge Distillation:
- ┌─────────────────────────┐
 - Large training data │
 - ┬────────────┘
 ↓
- ┌────────────────────────┐
 - Teacher model (large) │ ← Pre-trained, high quality
 - Generate soft targets │
 - ┬───────────┘
 ↓
- ┌────────────────────┐
 - Student soft targets│ (probability distributions)
 - from teacher output │
 - ┬────────┘
 ↓
- ┌────────────────────────────────────┐
 - Train student model to match │
 - teacher's probability distributions │
 - ┘
 ↓
- ┌────────────────────┐
 - Compressed student │
 - with teacher's │
 - knowledge │
 - ┘

Benefit:
- Teacher outputs encode richer patterns
- Student learns faster and better

Temperature-Scaled Softmax

Problem: Teacher outputs confident probabilities
 Student tries to match these directly
 Problem: Extreme values hard to match

Solution: Temperature scaling

Standard softmax:
p_i = exp(z_i) / Σ exp(z_j)

Temperature-scaled:
p_i = exp(z_i / T) / Σ exp(z_j / T)

Effect of temperature T:

T = 1 (standard):
 - Output: [0.95, 0.04, 0.01]
 - Distribution: Sharp, concentrated
 - Problem: Extreme probabilities hard to match

T = 5 (soft):
 - Output: [0.68, 0.21, 0.11]
 - Distribution: Smooth, spread out
 - Benefit: Student can match easier!

Higher T → Softer targets
 - "Softer" = more information about wrong answers
 - "Why wrong" is as important as "why right"

Example:
Student tries to predict: Is this an apple?

Hard target: [1.0, 0.0] (yes/no)
 - Student learns: "If I see red round fruit, say yes"
 - Doesn't learn other properties

Soft target (from teacher): [0.95, 0.05]
 - Information: "95% sure yes, 5% could be tomato or cherry"
 - Student learns: "Red roundness → mostly apple, but could be similar"
 - Transfers richer knowledge!

Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationTrainer:
 def __init__(self, teacher_model, student_model, temperature=5.0):
 self.teacher = teacher_model
 self.student = student_model
 self.temperature = temperature

 # Teacher should be frozen (not training)
 self.teacher.eval()
 for param in self.teacher.parameters():
 param.requires_grad = False

 def compute_distillation_loss(self, logits_student, logits_teacher):
 """
 KL divergence between student and teacher distributions
 (at temperature T)
 """
 # Soften probabilities using temperature
 q_teacher = F.softmax(logits_teacher / self.temperature, dim=-1)
 log_p_student = F.log_softmax(logits_student / self.temperature, dim=-1)

 # KL divergence
 kl_loss = F.kl_div(log_p_student, q_teacher, reduction='batchmean')

 # Scale by T^2 (for proper gradient scaling)
 distillation_loss = kl_loss * (self.temperature ** 2)
 return distillation_loss

 def train_step(self, batch, alpha=0.9):
 """
 alpha: Weight between distillation loss and task loss
 alpha=0.9: 90% distillation, 10% task loss
 """
 x, y_hard = batch

 # Teacher forward (no gradient)
 with torch.no_grad():
 logits_teacher = self.teacher(x)

 # Student forward
 logits_student = self.student(x)

 # Compute losses
 loss_distillation = self.compute_distillation_loss(
 logits_student, logits_teacher
)
 loss_task = F.cross_entropy(logits_student, y_hard)

 # Combine losses
 loss = alpha * loss_distillation + (1 - alpha) * loss_task

 return loss

# Usage:
teacher = LargeModel()
student = SmallModel()

trainer = DistillationTrainer(teacher, student, temperature=5.0)

optimizer = torch.optim.Adam(student.parameters())

for epoch in range(num_epochs):
 for batch in dataloader:
 loss = trainer.train_step(batch, alpha=0.9)
 loss.backward()
 optimizer.step()
 optimizer.zero_grad()

---

Distillation Strategies

1. Response-Based Distillation (Basic)

Match final output layer:

Teacher:
- ┌──────────────┐
 - Layer 96 │ ← Output layer
 - ┤
 - [0.92, 0.07, 0.01]
 - ┘
 ↓ (soft target)
Student:
- ┌──────────────┐
 - Layer 12 │ ← Much smaller
 - ┤
 - Train to match teacher's distribution
 - ┘

Pros:
Simple: Only match outputs
Works reasonably well
No access to teacher internals needed

Cons:
Doesn't leverage teacher's internal knowledge
Less efficient than deeper distillation

2. Feature-Based Distillation

Match intermediate layers:

Teacher:
- ┌─────────────────────────────┐
 - Layer 96 │ ← Match this
 - ┤
 - Feature size: (batch, seq, 2048)
 - ┘

Adapter:
- ┌─────────────────────────────┐
 - Project student features │ ← Adapt to match teacher size
 - (batch, seq, 256) → (batch, seq, 2048)
 - ┘

Student:
- ┌─────────────────────────────┐
 - Layer 6 │ ← Train to match
 - ┤
 - Feature size: (batch, seq, 256)
 - ┘

Pros:
Transfers deeper knowledge
Better quality improvement
Helpful for layer-to-layer alignment

Cons:
More complex
Requires careful layer matching
Teacher architecture knowledge needed

3. Relation-Based Distillation

Match the relationships between data points:

Teacher attention:
- ┌───────────────────────────────────┐
 - Attention matrix │
 - How important is each token? │
 - ┘

Student training:
 - Match not just outputs
 - But also attention patterns
 - Student learns: "teacher focuses on X and Y together"
 - Captures reasoning structure

Pro:
Teaches reasoning patterns
Better generalization

Con:
Complex to implement

-

Real-World Examples

DistilBERT

Teacher: BERT-base
 - 12 layers
 - 110M parameters
 - Baseline

Student: DistilBERT
 - 6 layers (50% of teacher)
 - 66M parameters (40% of teacher)
 - Training: Distillation + task loss (90/10)

Results:
 - Inference speed: 60% faster
 - Size: 40% smaller
 - Accuracy: 97% of original (3% loss)
 - Quality-speed trade: Excellent!

Deployment impact:
 - Can run on mobile / edge devices
 - Inference cost: 1/3 of BERT
 - Suitable for real-time applications
 - Industry widely adopted

Student LLMs

Concept: Distill large LLM to smaller one

Teacher: LLaMA 70B
 - Trained on 1.4T tokens
 - High quality but slow inference

Student: LLaMA 7B (improved)
 - Standard LLaMA 7B: 1T tokens
 - With distillation: Mix hard + soft targets from 70B
 - Result: Better quality than standard 7B

Quality comparison:

Model MMLU Benchmark Speed
──────────────────────────────────────────
LLaMA 7B (standard) 45.3% Baseline Baseline
LLaMA 7B (distilled) 46.8% +1.5% Baseline
Distilled maintains same speed but with better quality!

Trade-off vs full training:
 - Training distilled 7B: 1/10 compute of 70B teacher
 - Quality improvement: +1-2%
 - Better than training 7B from scratch!

-

Advantages and Limitations

Advantages

Smaller models (10-50x parameters reduction)
Faster inference (10-100x speedup possible)
Less memory needed (mobile/edge deployment)
Cheaper inference (lower compute cost)
Effective knowledge transfer
Can combine with quantization for 1000x speedup

Limitations

Quality loss (5-15% typical)
Requires pre-trained teacher model
Training complexity (tuning temperature, alpha)
Task-specific (model for one task, must distill again for another)
Not suitable for novel tasks (teacher knowledge doesn't apply)

---

Advanced Techniques

Progressive Distillation

Staged approach:

Step 1: 70B → 30B (distillation)
 Quality loss: 3%

Step 2: 30B → 7B (distillation)
 Quality loss: 2%

Step 3: 7B → 1B (distillation)
 Quality loss: 2%

Total: 70B → 1B with only 7% quality loss
 - Better than single-step distillation (would be 20%+ loss)

Intuition:
- Small steps allow better adaptation
- Each stage preserves more knowledge
- Better than trying to compress too much at once

Multi-Teacher Distillation

Instead of one teacher, use multiple:

Teachers:
 - Teacher A: Optimized for accuracy
 - Teacher B: Optimized for speed
 - Teacher C: Different architecture
 - Each captures different aspect of knowledge

Student:
 - Learns from all teachers
 - Captures diverse knowledge
 - Better generalization
 - More robust student

Ensemble benefit:
 - Student better than any single teacher could produce!

-

When to Use Distillation

Use Distillation When

Need fast inference on mobile/edge
Quality loss of 5-10% acceptable
Have pre-trained teacher model
Want to deploy existing model at scale
Inference cost is main concern

Avoid Distillation When

Quality must be pristine (loss unacceptable)
No suitable teacher model exists
Need novel capabilities (teacher doesn't have)
Training cost is main concern (not inference)

---

Key Takeaways

Transfer knowledge: Larger model teaches smaller model 10-100x inference speedup possible Temperature scaling creates soft, informative targets 5-10% quality loss typical, acceptable for many use cases DistilBERT, MobileBERT: Proven in production

-

Related Notes

• Quantization - Combine with distillation for maximum compression
• Lora - Alternative for model adaptation
• [Pruning & Sparsity](/01-modeling/02-training/03-compression/(pruning-sparsity/) - Complement to distillation
• Llm Inference Optimization - Complete inference optimization stack