Model Distillation: Compressing Knowledge into Smaller Models

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 - Complement to distillation
• Llm Inference Optimization - Complete inference optimization stack