concept_number: 18 title: Knowledge Distillation stage: Advanced difficulty: โญโญโญโญโ time_to_learn: 4-5 hours status: not-started date_created: 2026-03-07 date_started: date_completed: progress: 0% prerequisites:
- "Large Language Models"
- "Transformer"
- "Small Language Models" next_concepts:
- "Quantization"
- "Small Language Models" related_concepts:
- "Fine-Tuning"
- "Reinforcement Learning" tags:
- ai
- distillation
- compression
- teacher-student
- model-optimization
๐บ๏ธ Concept Mind Map
KNOWLEDGE DISTILLATION
|
+---------------+---------------+
| | |
PROCESS TECHNIQUES BENEFITS
| | |
| | |
Teacher Model Soft Labels Size Reduction
- Large (175B) - Probability - 350GB โ 14GB
- Accurate distributions - 25x smaller
- Expensive - Richer info - Fits on GPU
- Slow - Not just 0/1 - Portable
| | |
| | |
Student Model Temperature Speed Boost
- Small (7B) Scaling - 2000ms โ 200ms
- Fast - Soften probs - 10x faster
- Cheap - T=1 (hard) - Real-time
- 90% quality - T=5 (soft) - Responsive
| | |
| | |
Training Loss Function Cost Savings
- Learn from - KL Divergence - $0.03 โ $0.003
teacher - Match teacher - 10x cheaper
- Mimic outputs - Scale easily
behavior - Soft targets - Affordable
| | |
| | |
Evaluation Progressive Quality
- Compare to Distillation - 90-95% of
teacher - Layer-by-layer teacher
- Measure gap - Gradual - Good enough
- Optimize transfer - Production-ready
Teacher-Student Process:
STEP 1: Teacher Generates Soft Labels
Input: "The movie was"
Teacher (GPT-4):
- "great" (0.6)
- "good" (0.25)
- "amazing" (0.1)
- "terrible" (0.05)
STEP 2: Student Learns from Distribution
Student (7B) learns:
- Not just "great" (hard label)
- But the full distribution (soft labels)
- Captures teacher's uncertainty
- Learns nuanced patterns
STEP 3: Student Mimics Teacher
Student output:
- "great" (0.58)
- "good" (0.27)
- "amazing" (0.09)
- "terrible" (0.06)
Very close to teacher!
Distillation vs Training from Scratch:
APPROACH TIME COST QUALITY DATA NEEDED
From Scratch Weeks $1M 100% Billions of tokens
Distillation Days $50K 90-95% Millions of tokens
Fine-Tuning Hours $5K 85% Thousands of examples
Temperature Effect:
T=1 (Hard):
[0.9, 0.05, 0.03, 0.02] โ Sharp, less info
T=5 (Soft):
[0.4, 0.25, 0.2, 0.15] โ Smooth, more info
Higher T = More knowledge transfer!
Real-World Examples:
- DistilBERT: BERT (110M) โ DistilBERT (66M), 97% quality, 60% faster
- TinyBERT: BERT โ TinyBERT (14M), 96% quality, 9x smaller
- DistilGPT-2: GPT-2 โ DistilGPT-2, 2x faster, 90% quality
Key Insight: Distillation transfers "dark knowledge" - the subtle patterns and uncertainties that make the teacher smart!
๐ Start With a Real-World Scenario
You have GPT-4 (175B parameters):
Problems:
- Size: 350GB (won't fit on most servers!)
- Speed: 2 seconds per response (slow!)
- Cost: $0.03 per 1K tokens (expensive!)
- Memory: Requires 8ร A100 GPUs ($80K!)
You need:
- Size: <10GB (fits on one GPU)
- Speed: <200ms (10x faster)
- Cost: <$0.003 per 1K tokens (10x cheaper)
- Memory: Runs on consumer GPU ($1K)
Solution: Knowledge Distillation! ๐
The Teacher-Student Analogy:
Traditional Training:
Student learns from textbooks (raw data)
- Reads millions of pages
- Takes years
- Expensive
Distillation:
Student learns from expert teacher (GPT-4)
- Teacher explains concepts
- Student learns faster
- Much cheaper!
Real Example:
Teacher: GPT-4 (175B params)
Question: "Explain photosynthesis"
GPT-4: "Photosynthesis is the process by which plants convert
light energy into chemical energy. It occurs in chloroplasts
and involves two stages: light-dependent reactions and the
Calvin cycle..."
[Detailed, nuanced, perfect answer]
Student: Distilled Model (7B params)
Question: "Explain photosynthesis"
Student: "Photosynthesis is how plants make food from sunlight.
It happens in chloroplasts and has two main steps: capturing
light energy and making sugar..."
[Simpler but still accurate!]
Student is 25x smaller but 90% as good! ๐ฏ
The Process:
Step 1: Teacher generates "soft labels"
Input: "The movie was"
Teacher probabilities:
- "great" (60%)
- "good" (25%)
- "amazing" (10%)
- "okay" (5%)
These probabilities contain MORE information than just "great"!
Step 2: Student learns from soft labels
Student learns:
- Not just the answer ("great")
- But the uncertainty (60% vs 25% vs 10%)
- The relationships between words
- The teacher's "reasoning"
Step 3: Result
Student model:
- 25x smaller
- 10x faster
- 90-95% of teacher's quality
- 10x cheaper to run
Real-World Success Stories:
1. DistilBERT (Hugging Face)
- Teacher: BERT (110M params)
- Student: DistilBERT (66M params)
- Size: 40% smaller
- Speed: 60% faster
- Quality: 97% of BERT
- Used by millions!
2. TinyBERT
- Teacher: BERT (110M params)
- Student: TinyBERT (14M params)
- Size: 87% smaller!
- Speed: 9x faster!
- Quality: 96% of BERT
3. MobileBERT (Google)
- Runs on phones
- 4x smaller than BERT
- Powers Google Assistant
- Billions of users!
The Economics:
Running GPT-4:
- Hardware: 8ร A100 GPUs = $80,000
- Power: $2,000/month
- Inference: $0.03 per 1K tokens
- Total: $100K+ initial + $2K/month
Running Distilled Model:
- Hardware: 1ร Consumer GPU = $1,000
- Power: $50/month
- Inference: $0.003 per 1K tokens
- Total: $1K initial + $50/month
Savings: 99% cost reduction! ๐ฐ
By the end of this guide, you'll understand:
- How distillation transfers knowledge
- Soft labels vs. hard labels
- Different distillation techniques
- How to distill your own models
- When distillation works (and when it doesn't)
- Real-world deployment strategies
The breakthrough: Get 90% of GPT-4's performance at 1% of the cost! ๐
Let's dive in! ๐
๐ Knowledge Distillation
Transfer knowledge from large "teacher" models to smaller "student" models while maintaining performance
zations](#interactive-tools--visualizations)
- Code Examples
- Hands-On Exercises
- Assignments & Practice
- Common Pitfalls & Debugging
- Fun Facts & Hacks
- Related Concepts
- Resources
What is Knowledge Distillation?
Knowledge Distillation is a model compression technique where a smaller "student" model learns to mimic the behavior of a larger "teacher" model. Instead of learning from raw data labels, the student learns from the teacher's soft predictions (probability distributions), which contain richer information about relationships between classes.
Key Idea
Large Teacher Model (100B params)
โ (transfers knowledge)
Small Student Model (1B params)
โ
Similar performance, 100x smaller!
Real-World Analogy
Think of it like learning from an expert teacher:
- Teacher: Experienced professor with deep knowledge
- Student: Learns not just the answers, but the reasoning
- Soft targets: Teacher's confidence levels, not just right/wrong
- Result: Student gains intuition, not just memorization
Why Distillation Matters
1. Deployment Efficiency ๐
- Run models on mobile devices
- Reduce inference costs by 10-100x
- Lower latency for real-time applications
2. Maintain Performance ๐ฏ
- Keep 95-99% of teacher's accuracy
- Preserve reasoning capabilities
- Transfer specialized knowledge
3. Practical Applications ๐ผ
- Edge AI deployment
- Cost-effective serving
- Privacy-preserving inference
4. Environmental Impact ๐ฑ
- Reduce energy consumption
- Lower carbon footprint
- Sustainable AI deployment
Core Concepts
1. Teacher-Student Framework
# Teacher: Large, accurate model
teacher = LargeModel(params=100_000_000_000)
# Student: Small, efficient model
student = SmallModel(params=1_000_000_000)
# Distillation: Transfer knowledge
distill(teacher, student, data)
2. Soft Targets vs Hard Labels
Hard Labels (Traditional):
Input: "The cat sat on the mat"
Label: [0, 0, 1, 0] # One-hot: "positive"
Soft Targets (Distillation):
Input: "The cat sat on the mat"
Teacher Output: [0.05, 0.10, 0.75, 0.10]
โ โ โ โ
negative neutral positive very_positive
3. Temperature Scaling
# Temperature controls softness of predictions
logits = model(input)
# T=1: Normal predictions (sharp)
probs_t1 = softmax(logits / 1.0)
# T=5: Softer predictions (more information)
probs_t5 = softmax(logits / 5.0)
4. Distillation Loss
# Combined loss function
loss = ฮฑ * distillation_loss(student, teacher) +
(1-ฮฑ) * task_loss(student, labels)
How It Works
Step-by-Step Process
graph TD
A[Training Data] --> B[Teacher Model]
B --> C[Soft Predictions]
A --> D[Student Model]
C --> E[Distillation Loss]
D --> E
A --> F[Hard Labels]
F --> G[Task Loss]
D --> G
E --> H[Combined Loss]
G --> H
H --> I[Update Student]
I --> D
Mathematical Foundation
1. Soft Target Generation:
q_i = exp(z_i / T) / ฮฃ_j exp(z_j / T)
Where:
z_i: Teacher's logitsT: Temperature (higher = softer)q_i: Soft target probability
2. Distillation Loss:
L_distill = KL(q_teacher || q_student)
3. Combined Loss:
L_total = ฮฑ * L_distill + (1-ฮฑ) * L_task
Types of Distillation
1. Response-Based Distillation
Transfer final layer outputs:
# Teacher's final predictions
teacher_logits = teacher(input)
# Student learns to match
student_logits = student(input)
loss = KL_divergence(teacher_logits, student_logits)
2. Feature-Based Distillation
Transfer intermediate representations:
# Teacher's hidden states
teacher_features = teacher.get_hidden_states(input)
# Student matches intermediate layers
student_features = student.get_hidden_states(input)
loss = MSE(teacher_features, student_features)
3. Relation-Based Distillation
Transfer relationships between samples:
# Teacher's similarity matrix
teacher_relations = compute_relations(teacher_features)
# Student learns relationships
student_relations = compute_relations(student_features)
loss = MSE(teacher_relations, student_relations)
4. Self-Distillation
Model learns from itself:
# Use model's own predictions as soft targets
model_v1 = train_model(data)
model_v2 = distill(teacher=model_v1, student=model_v1, data)
Local Setup
Installation (15-20 minutes)
# Create environment
conda create -n distillation python=3.10
conda activate distillation
# Install core libraries
pip install torch transformers datasets
pip install accelerate bitsandbytes
pip install wandb tensorboard
# Install distillation frameworks
pip install textbrewer # For NLP distillation
pip install neural-compressor # Intel's toolkit
Quick Start
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from textbrewer import DistillationConfig, TrainingConfig, BasicTrainer
import torch
# Load teacher (large model)
teacher = AutoModelForSequenceClassification.from_pretrained(
"bert-large-uncased",
num_labels=2
)
# Load student (small model)
student = AutoModelForSequenceClassification.from_pretrained(
"bert-base-uncased",
num_labels=2
)
# Configure distillation
distill_config = DistillationConfig(
temperature=4.0,
hard_label_weight=0.5,
kd_loss_weight=0.5
)
# Train
trainer = BasicTrainer(
train_config=train_config,
distill_config=distill_config,
model_T=teacher,
model_S=student
)
Code Examples
Example 1: Basic Distillation (Beginner)
import torch
import torch.nn as nn
import torch.nn.functional as F
class DistillationLoss(nn.Module):
"""Basic distillation loss with temperature scaling"""
def __init__(self, temperature=3.0, alpha=0.5):
super().__init__()
self.temperature = temperature
self.alpha = alpha
self.ce_loss = nn.CrossEntropyLoss()
def forward(self, student_logits, teacher_logits, labels):
# Soft targets from teacher
soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)
soft_student = F.log_softmax(student_logits / self.temperature, dim=1)
# Distillation loss (KL divergence)
distill_loss = F.kl_div(
soft_student,
soft_targets,
reduction='batchmean'
) * (self.temperature ** 2)
# Task loss (cross entropy)
task_loss = self.ce_loss(student_logits, labels)
# Combined loss
return self.alpha * distill_loss + (1 - self.alpha) * task_loss
# Usage
teacher = TeacherModel()
student = StudentModel()
criterion = DistillationLoss(temperature=4.0, alpha=0.7)
# Training loop
for batch in dataloader:
inputs, labels = batch
# Get predictions
with torch.no_grad():
teacher_logits = teacher(inputs)
student_logits = student(inputs)
# Compute loss
loss = criterion(student_logits, teacher_logits, labels)
# Backprop
loss.backward()
optimizer.step()
Example 2: BERT Distillation (Intermediate)
from transformers import (
BertForSequenceClassification,
BertTokenizer,
Trainer,
TrainingArguments
)
import torch.nn as nn
class DistillationTrainer(Trainer):
"""Custom trainer for BERT distillation"""
def __init__(self, teacher_model, temperature=2.0, alpha=0.5, *args, **kwargs):
super().__init__(*args, **kwargs)
self.teacher = teacher_model
self.teacher.eval()
self.temperature = temperature
self.alpha = alpha
def compute_loss(self, model, inputs, return_outputs=False):
# Student forward pass
outputs_student = model(**inputs)
student_logits = outputs_student.logits
# Teacher forward pass (no grad)
with torch.no_grad():
outputs_teacher = self.teacher(**inputs)
teacher_logits = outputs_teacher.logits
# Distillation loss
loss_fct = nn.KLDivLoss(reduction='batchmean')
loss_distill = loss_fct(
F.log_softmax(student_logits / self.temperature, dim=-1),
F.softmax(teacher_logits / self.temperature, dim=-1)
) * (self.temperature ** 2)
# Task loss
loss_task = outputs_student.loss
# Combined
loss = self.alpha * loss_distill + (1 - self.alpha) * loss_task
return (loss, outputs_student) if return_outputs else loss
# Setup
teacher = BertForSequenceClassification.from_pretrained(
'bert-large-uncased',
num_labels=2
)
student = BertForSequenceClassification.from_pretrained(
'bert-base-uncased',
num_labels=2
)
# Training arguments
training_args = TrainingArguments(
output_dir='./distilled_bert',
num_train_epochs=3,
per_device_train_batch_size=16,
learning_rate=2e-5,
warmup_steps=500,
logging_steps=100
)
# Train
trainer = DistillationTrainer(
teacher_model=teacher,
model=student,
args=training_args,
train_dataset=train_dataset,
temperature=4.0,
alpha=0.7
)
trainer.train()
Example 3: Feature Distillation (Advanced)
import torch
import torch.nn as nn
class FeatureDistillation(nn.Module):
"""Distill intermediate layer representations"""
def __init__(self, teacher, student, layer_mapping):
super().__init__()
self.teacher = teacher
self.student = student
self.layer_mapping = layer_mapping # {student_layer: teacher_layer}
# Projection layers to match dimensions
self.projections = nn.ModuleDict({
str(s_layer): nn.Linear(
student.config.hidden_size,
teacher.config.hidden_size
)
for s_layer in layer_mapping.keys()
})
def forward(self, inputs, labels):
# Get all hidden states
student_outputs = self.student(
**inputs,
output_hidden_states=True,
return_dict=True
)
with torch.no_grad():
teacher_outputs = self.teacher(
**inputs,
output_hidden_states=True,
return_dict=True
)
# Feature distillation loss
feature_loss = 0
for s_layer, t_layer in self.layer_mapping.items():
# Get hidden states
student_hidden = student_outputs.hidden_states[s_layer]
teacher_hidden = teacher_outputs.hidden_states[t_layer]
# Project student to teacher dimension
student_projected = self.projections[str(s_layer)](student_hidden)
# MSE loss
feature_loss += F.mse_loss(student_projected, teacher_hidden)
# Task loss
task_loss = student_outputs.loss
# Combined
total_loss = task_loss + 0.5 * feature_loss
return total_loss, student_outputs
# Usage
layer_mapping = {
3: 6, # Student layer 3 โ Teacher layer 6
6: 12, # Student layer 6 โ Teacher layer 12
9: 18, # Student layer 9 โ Teacher layer 18
12: 24 # Student layer 12 โ Teacher layer 24
}
model = FeatureDistillation(teacher, student, layer_mapping)
Example 4: DistilBERT from Scratch (Production)
from transformers import (
BertConfig,
BertForMaskedLM,
BertTokenizer,
DataCollatorForLanguageModeling,
Trainer,
TrainingArguments
)
class DistilBERTTrainer:
"""Complete DistilBERT training pipeline"""
def __init__(self, teacher_name='bert-base-uncased'):
self.teacher_name = teacher_name
self.tokenizer = BertTokenizer.from_pretrained(teacher_name)
def create_student_config(self):
"""Create smaller student configuration"""
teacher_config = BertConfig.from_pretrained(self.teacher_name)
student_config = BertConfig(
vocab_size=teacher_config.vocab_size,
hidden_size=teacher_config.hidden_size,
num_hidden_layers=6, # Half of teacher's 12 layers
num_attention_heads=teacher_config.num_attention_heads,
intermediate_size=teacher_config.intermediate_size,
hidden_dropout_prob=0.1,
attention_probs_dropout_prob=0.1
)
return student_config
def initialize_student(self, teacher):
"""Initialize student from teacher weights"""
student_config = self.create_student_config()
student = BertForMaskedLM(student_config)
# Copy embeddings
student.bert.embeddings.load_state_dict(
teacher.bert.embeddings.state_dict()
)
# Copy every other layer
for i in range(6):
student.bert.encoder.layer[i].load_state_dict(
teacher.bert.encoder.layer[i * 2].state_dict()
)
return student
def distill(self, train_dataset, output_dir='./distilbert'):
"""Run distillation"""
# Load teacher
teacher = BertForMaskedLM.from_pretrained(self.teacher_name)
teacher.eval()
# Create student
student = self.initialize_student(teacher)
# Data collator
data_collator = DataCollatorForLanguageModeling(
tokenizer=self.tokenizer,
mlm=True,
mlm_probability=0.15
)
# Training arguments
training_args = TrainingArguments(
output_dir=output_dir,
num_train_epochs=3,
per_device_train_batch_size=32,
learning_rate=5e-5,
warmup_steps=10000,
weight_decay=0.01,
logging_steps=500,
save_steps=10000,
fp16=True
)
# Custom trainer with distillation
trainer = DistillationTrainer(
teacher_model=teacher,
model=student,
args=training_args,
train_dataset=train_dataset,
data_collator=data_collator,
temperature=2.0,
alpha=0.5
)
# Train
trainer.train()
# Save
student.save_pretrained(f'{output_dir}/final')
self.tokenizer.save_pretrained(f'{output_dir}/final')
return student
# Usage
distiller = DistilBERTTrainer('bert-base-uncased')
student_model = distiller.distill(train_dataset)
Example 5: Multi-Task Distillation (Expert)
class MultiTaskDistillation(nn.Module):
"""Distill multiple tasks simultaneously"""
def __init__(self, teacher_models, student_model, task_weights):
super().__init__()
self.teachers = nn.ModuleDict(teacher_models)
self.student = student_model
self.task_weights = task_weights
# Task-specific heads for student
self.task_heads = nn.ModuleDict({
task: nn.Linear(student_model.config.hidden_size, num_labels)
for task, num_labels in task_weights.items()
})
def forward(self, inputs, task_name, labels=None):
# Student encoding
student_outputs = self.student(**inputs, output_hidden_states=True)
student_hidden = student_outputs.last_hidden_state[:, 0, :] # [CLS]
# Task-specific prediction
student_logits = self.task_heads[task_name](student_hidden)
# Teacher prediction
with torch.no_grad():
teacher = self.teachers[task_name]
teacher_outputs = teacher(**inputs)
teacher_logits = teacher_outputs.logits
# Distillation loss
distill_loss = F.kl_div(
F.log_softmax(student_logits / 2.0, dim=-1),
F.softmax(teacher_logits / 2.0, dim=-1),
reduction='batchmean'
) * 4.0
# Task loss
if labels is not None:
task_loss = F.cross_entropy(student_logits, labels)
total_loss = 0.5 * distill_loss + 0.5 * task_loss
else:
total_loss = distill_loss
return total_loss, student_logits
# Usage with multiple tasks
teacher_models = {
'sentiment': BertForSequenceClassification.from_pretrained('bert-sentiment'),
'nli': BertForSequenceClassification.from_pretrained('bert-nli'),
'qa': BertForQuestionAnswering.from_pretrained('bert-qa')
}
student = BertModel.from_pretrained('bert-base-uncased')
model = MultiTaskDistillation(
teacher_models=teacher_models,
student_model=student,
task_weights={'sentiment': 2, 'nli': 3, 'qa': 2}
)
# Train on multiple tasks
for task, dataloader in task_dataloaders.items():
for batch in dataloader:
loss, logits = model(batch['inputs'], task, batch['labels'])
loss.backward()
optimizer.step()
TODO: Implement temperature comparison
def compare_temperatures(teacher_logits, temperatures=[1, 2, 5, 10]): """ Compare soft targets at different temperatures
Args:
teacher_logits: Raw logits from teacher
temperatures: List of temperatures to try
Returns:
Dictionary of soft targets for each temperature
"""
# Your code here
pass
Test with sample logits
logits = torch.tensor([[2.0, 1.0, 0.5, 0.1]]) results = compare_temperatures(logits)
### Exercise 2: Loss Function Design
**Goal**: Implement custom distillation loss
```python
# TODO: Create adaptive distillation loss
class AdaptiveDistillationLoss(nn.Module):
"""
Distillation loss that adapts alpha based on training progress
"""
def __init__(self, initial_alpha=0.5, final_alpha=0.1):
super().__init__()
# Your code here
pass
def forward(self, student_logits, teacher_logits, labels, progress):
"""
Args:
progress: Training progress (0.0 to 1.0)
"""
# Your code here
pass
Exercise 3: Layer Selection
Goal: Find optimal layers for feature distillation
# TODO: Implement layer importance analysis
def analyze_layer_importance(teacher, student, val_dataset):
"""
Determine which teacher layers are most important
for student learning
Returns:
Dictionary mapping student layers to teacher layers
"""
# Your code here
pass
Exercise 4: Compression Ratio Analysis
Goal: Measure compression vs performance tradeoff
# TODO: Benchmark different student sizes
def compression_analysis(teacher, student_configs, test_dataset):
"""
Train students of different sizes and compare:
- Model size
- Inference speed
- Accuracy
- Memory usage
Returns:
DataFrame with comparison metrics
"""
# Your code here
pass
Exercise 5: Self-Distillation
Goal: Implement iterative self-distillation
# TODO: Implement self-distillation loop
def self_distill(model, dataset, num_iterations=3):
"""
Iteratively distill model from itself
Args:
model: Initial model
dataset: Training data
num_iterations: Number of self-distillation rounds
Returns:
List of models from each iteration
"""
# Your code here
pass
Assignments & Practice
Assignment 1: Distill GPT-2
Difficulty: โญโญโญโโ
Create a distilled version of GPT-2:
- Load GPT-2 medium as teacher
- Create GPT-2 small as student
- Implement distillation training
- Compare perplexity and speed
- Deploy to Hugging Face
Deliverables:
- Training script
- Evaluation metrics
- Model card
- Inference demo
Assignment 2: Mobile BERT
Difficulty: โญโญโญโญโ
Build a BERT model for mobile deployment:
- Distill BERT-base to 6 layers
- Apply quantization
- Convert to ONNX
- Optimize for mobile
- Benchmark on device
Deliverables:
- Distilled model
- ONNX export
- Mobile app demo
- Performance report
Assignment 3: Multi-Teacher Ensemble
Difficulty: โญโญโญโญโญ
Distill knowledge from multiple teachers:
- Train 3 different teacher models
- Create ensemble distillation
- Compare with single teacher
- Analyze knowledge transfer
- Write research report
Deliverables:
- Ensemble distillation code
- Ablation studies
- Visualization of learned features
- Technical report
Student accuracy much lower than teacher
teacher_acc = 0.95 student_acc = 0.75 # Too low!
**Solutions**:
```python
# 1. Increase temperature
temperature = 5.0 # Try higher values
# 2. Adjust alpha (more weight on distillation)
alpha = 0.8 # Increase from 0.5
# 3. Use feature distillation
add_feature_matching_loss()
# 4. Better initialization
initialize_student_from_teacher()
Issue 2: Training Instability
Symptoms:
# Loss oscillates or diverges
epoch_losses = [2.5, 1.8, 3.2, 1.5, 4.1, ...]
Solutions:
# 1. Lower learning rate
optimizer = AdamW(student.parameters(), lr=1e-5)
# 2. Gradient clipping
torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
# 3. Warmup schedule
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=1000,
num_training_steps=total_steps
)
# 4. Smaller batch size
batch_size = 16 # Reduce from 32
Issue 3: Memory Issues
Symptoms:
# CUDA out of memory during distillation
RuntimeError: CUDA out of memory
Solutions:
# 1. Gradient accumulation
accumulation_steps = 4
for i, batch in enumerate(dataloader):
loss = compute_loss(batch)
loss = loss / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
# 2. Mixed precision training
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
loss = compute_loss(batch)
# 3. Freeze teacher
teacher.eval()
for param in teacher.parameters():
param.requires_grad = False
Issue 4: Slow Convergence
Symptoms:
# Student takes too long to train
epochs_needed = 50 # Should be ~10
Solutions:
# 1. Better initialization
copy_teacher_weights_to_student()
# 2. Layer-wise learning rates
optimizer = AdamW([
{'params': student.embeddings.parameters(), 'lr': 1e-5},
{'params': student.encoder.parameters(), 'lr': 2e-5},
{'params': student.classifier.parameters(), 'lr': 5e-5}
])
# 3. Progressive distillation
# Start with high temperature, gradually decrease
temperature = max(1.0, initial_temp * (1 - progress))
Hack 1: Data-free distillation
Generate synthetic data from teacher
synthetic_data = generate_from_teacher(teacher, num_samples=10000) distill(teacher, student, synthetic_data)
Hack 2: Progressive layer dropping
Gradually remove teacher layers during training
def progressive_distillation(teacher, student, epochs): for epoch in range(epochs): # Use fewer teacher layers as training progresses num_layers = max(6, 12 - epoch) teacher_subset = teacher.layers[:num_layers] train_epoch(teacher_subset, student)
Hack 3: Attention transfer
Distill attention patterns, not just outputs
def attention_distillation(teacher_attn, student_attn): return F.mse_loss( student_attn / student_attn.sum(dim=-1, keepdim=True), teacher_attn / teacher_attn.sum(dim=-1, keepdim=True) )
Hack 4: Task-specific distillation
Different temperatures for different tasks
task_temperatures = { 'classification': 3.0, 'regression': 1.0, 'generation': 5.0 }
Hack 5: Curriculum distillation
Start with easy examples, progress to hard
sorted_data = sort_by_difficulty(dataset) for difficulty_level in ['easy', 'medium', 'hard']: subset = sorted_data[difficulty_level] distill(teacher, student, subset)
### ๐ง Debugging Tricks
```python
# Visualize soft vs hard targets
import matplotlib.pyplot as plt
def visualize_targets(teacher_logits, labels, temperature=3.0):
soft = F.softmax(teacher_logits / temperature, dim=-1)
hard = F.one_hot(labels, num_classes=len(teacher_logits[0]))
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.bar(range(len(soft[0])), soft[0].detach().numpy())
plt.title('Soft Targets')
plt.subplot(1, 2, 2)
plt.bar(range(len(hard[0])), hard[0].numpy())
plt.title('Hard Labels')
plt.show()
# Monitor distillation quality
def distillation_metrics(teacher, student, val_loader):
metrics = {
'agreement': 0, # How often they agree
'kl_divergence': 0, # Distribution difference
'confidence_gap': 0 # Confidence difference
}
for batch in val_loader:
teacher_pred = teacher(batch).argmax(dim=-1)
student_pred = student(batch).argmax(dim=-1)
metrics['agreement'] += (teacher_pred == student_pred).float().mean()
# ... compute other metrics
return metrics
Related Concepts
Prerequisites
- Large Language Models - Understanding model architecture
- Transformer - Base architecture
- Fine-Tuning - Training techniques
Related Topics
- Reinforcement Learning - Alternative training methods
- Model Pruning - Removing unnecessary weights
- Neural Architecture Search - Finding optimal architectures
Resources
๐ Essential Papers
"Distilling the Knowledge in a Neural Network" (Hinton et al., 2015)
- Original distillation paper
- https://arxiv.org/abs/1503.02531
"DistilBERT" (Sanh et al., 2019)
- Practical BERT distillation
- https://arxiv.org/abs/1910.01108
"TinyBERT" (Jiao et al., 2020)
- Advanced distillation techniques
- https://arxiv.org/abs/1909.10351
"Patient Knowledge Distillation" (Sun et al., 2019)
- Layer-wise distillation
- https://arxiv.org/abs/1908.09355
๐ฅ Video Tutorials
"Knowledge Distillation Explained" - Yannic Kilcher
"Distilling BERT" - Hugging Face
"Model Compression Techniques" - Stanford CS224N
๐ ๏ธ Tools & Libraries
TextBrewer - NLP distillation framework
Neural Compressor - Intel's compression toolkit
Hugging Face Optimum - Model optimization
PyTorch Distillation - Official examples
๐ Tutorials & Guides
Hugging Face Distillation Guide
Distillation Best Practices
Model Compression Survey
Progress Tracking
- Understand teacher-student framework
- Implement basic distillation loss
- Experiment with temperature scaling
- Try feature-based distillation
- Complete BERT distillation project
- Deploy distilled model
- Benchmark performance improvements
- Explore advanced techniques
๐ Master Knowledge Distillation!
Transfer knowledge efficiently and deploy AI everywhere
