Compress large language models using knowledge distillation from teacher to student models. Use when deploying smaller models with retained performance, transferring GPT-4 capabilities to open-source models, or reducing inference costs. Covers temperature scaling, soft targets, reverse KLD, logit distillation, and MiniLLM training strategies.
Compress large language models using knowledge distillation from teacher to student models. Use when deploying smaller models with retained performance, transferring GPT-4 capabilities to open-source models, or reducing inference costs. Covers temperature scaling, soft targets, reverse KLD, logit distillation, and MiniLLM training strategies.
# Standard transformers
pip install transformers datasets accelerate
# For training
pip install torch deepspeed wandb
# Optional: MiniLLM implementation
git clone https://github.com/microsoft/LMOps
cd LMOps/minillm
pip install -e .
Quick Start
Basic Knowledge Distillation
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
# 1. Load teacher (large) and student (small) models
teacher = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-70b-hf", # Large teacher
torch_dtype=torch.float16,
device_map="auto"
)
student = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-7b-hf", # Small student
torch_dtype=torch.float16,
device_map="cuda:0"
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-hf")
# 2. Define distillation lossdefdistillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
"""
Combine hard loss (cross-entropy) with soft loss (KL divergence).
Args:
temperature: Softens probability distributions (higher = softer)
alpha: Weight for distillation loss (1-alpha for hard loss)
"""# Hard loss: Standard cross-entropy with true labels
hard_loss = F.cross_entropy(student_logits.view(-1, student_logits.size(-1)), labels.view(-1))
# Soft loss: KL divergence between student and teacher
soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
soft_student = F.log_softmax(student_logits / temperature, dim=-1)
soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean') * (temperature ** 2)
# Combined lossreturn alpha * soft_loss + (1 - alpha) * hard_loss
# 3. Training loopfor batch in dataloader:
# Teacher forward (no grad)with torch.no_grad():
teacher_outputs = teacher(**batch)
teacher_logits = teacher_outputs.logits
# Student forward
student_outputs = student(**batch)
student_logits = student_outputs.logits
# Compute distillation loss
loss = distillation_loss(
student_logits,
teacher_logits,
batch['labels'],
temperature=2.0,
alpha=0.7# 70% soft, 30% hard
)
# Backward and optimize
loss.backward()
optimizer.step()
optimizer.zero_grad()
MiniLLM (Reverse KLD)
Source: arXiv 2306.08543 (2024)
Innovation: Use reverse KLD instead of forward KLD for better generative model distillation.
defreverse_kl_loss(student_logits, teacher_logits, temperature=1.0):
"""
Reverse KL divergence: KL(Teacher || Student)
Better for generative models than forward KL.
"""# Teacher distribution (target)
p_teacher = F.softmax(teacher_logits / temperature, dim=-1)
# Student distribution (model)
log_p_student = F.log_softmax(student_logits / temperature, dim=-1)
# Reverse KL: Sum over teacher, student learns to cover teacher's modes
reverse_kl = -(p_teacher * log_p_student).sum(dim=-1).mean()
return reverse_kl * (temperature ** 2)
# Training with MiniLLMfor batch in dataloader:
with torch.no_grad():
teacher_logits = teacher(**batch).logits
student_logits = student(**batch).logits
# Reverse KLD (better for generation)
loss = reverse_kl_loss(student_logits, teacher_logits, temperature=1.0)
loss.backward()
optimizer.step()
Why reverse KL?
Forward KL (standard): Student learns to match teacher's mean
Reverse KL (MiniLLM): Student learns to cover all teacher's modes
Better for diverse text generation
Response Distillation
# Generate synthetic data from teacher, train student to imitate# 1. Generate synthetic responses from teacher
prompts = ["Explain AI:", "What is ML?", "Define NLP:"]
teacher_responses = []
for prompt in prompts:
inputs = tokenizer(prompt, return_tensors='pt').to(teacher.device)
outputs = teacher.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
teacher_responses.append(response)
# 2. Train student on teacher's responses (standard fine-tuning)
train_dataset = [
{"text": f"{prompt}\n{response}"}
for prompt, response inzip(prompts, teacher_responses)
]
# 3. Fine-tune student
trainer = Trainer(
model=student,
args=TrainingArguments(output_dir="./student", num_train_epochs=3, learning_rate=2e-5),
train_dataset=train_dataset,
)
trainer.train()
Core Concepts
1. Temperature Scaling
Purpose: Soften probability distributions to expose teacher's uncertainty.
# Low temperature (T=1): Sharp distribution
logits = [3.0, 2.0, 1.0]
probs_T1 = softmax(logits / 1.0) # [0.67, 0.24, 0.09]# High temperature (T=4): Soft distribution
probs_T4 = softmax(logits / 4.0) # [0.42, 0.34, 0.24]# Higher T reveals more information about relative rankings
Rule: Use T=2-5 for distillation (2 is common default).
2. Loss Function Components
# Total loss = alpha * soft_loss + (1 - alpha) * hard_loss# Soft loss: Learn from teacher's knowledge
soft_loss = KL(student || teacher)
# Hard loss: Learn from ground truth labels
hard_loss = CrossEntropy(student_output, true_labels)
# Typical values:
alpha = 0.5# Balanced
alpha = 0.7# More emphasis on teacher
alpha = 0.3# More emphasis on labels
3. Forward vs Reverse KLD
# Forward KL: KL(Student || Teacher)# - Student matches teacher's average behavior# - Mode-seeking: Student focuses on teacher's highest probability modes# - Good for classification# Reverse KL: KL(Teacher || Student)# - Student covers all of teacher's behaviors# - Mode-covering: Student learns diverse behaviors# - Good for generation (MiniLLM)
Training Strategies
Strategy 1: Logit Distillation
# Train student to match teacher's logits directlydeflogit_distillation_trainer(student, teacher, dataloader, temperature=2.0):
optimizer = torch.optim.AdamW(student.parameters(), lr=2e-5)
for epoch inrange(3):
for batch in dataloader:
# Get logitswith torch.no_grad():
teacher_logits = teacher(**batch).logits
student_logits = student(**batch).logits
# MSE on logits (alternative to KLD)
loss = F.mse_loss(student_logits, teacher_logits)
# Or use KLD# loss = F.kl_div(# F.log_softmax(student_logits/temperature, dim=-1),# F.softmax(teacher_logits/temperature, dim=-1),# reduction='batchmean'# ) * (temperature ** 2)
loss.backward()
optimizer.step()
optimizer.zero_grad()
return student
Strategy 2: Two-Stage Distillation
# Stage 1: Distill from teacher
student = distill(teacher, student, epochs=5)
# Stage 2: Fine-tune on task-specific data
student = fine_tune(student, task_data, epochs=3)
# Results in better task performance than single-stage
Strategy 3: Multi-Teacher Distillation
# Learn from multiple expert teachersdefmulti_teacher_distillation(student, teachers, batch):
"""Distill from ensemble of teachers."""
teacher_logits_list = []
# Get logits from all teacherswith torch.no_grad():
for teacher in teachers:
logits = teacher(**batch).logits
teacher_logits_list.append(logits)
# Average teacher predictions
avg_teacher_logits = torch.stack(teacher_logits_list).mean(dim=0)
# Student learns from ensemble
student_logits = student(**batch).logits
loss = F.kl_div(
F.log_softmax(student_logits, dim=-1),
F.softmax(avg_teacher_logits, dim=-1),
reduction='batchmean'
)
return loss
# Temperature
T = 1.0# Sharp (less knowledge transfer)
T = 2.0# Standard (good balance)
T = 5.0# Soft (more knowledge transfer)# Alpha (weight)
alpha = 0.5# Balanced
alpha = 0.7# Emphasize teacher knowledge
alpha = 0.9# Strong distillation# Rule: Higher T + higher alpha = stronger distillation
2. Model Size Ratio
# Good ratios (teacher/student)
70B / 7B = 10× # Excellent
13B / 1B = 13× # Good
7B / 1B = 7× # Acceptable# Avoid too large gap
70B / 1B = 70× # Too large, ineffective
3. Data Quality
# Best: Use teacher-generated data + real data
train_data = {
"teacher_generated": 70%, # Diverse, high-quality"real_data": 30% # Ground truth
}
# Avoid: Only real data (doesn't utilize teacher fully)
Evaluation
from transformers import pipeline
# Compare student vs teacher
teacher_pipe = pipeline("text-generation", model=teacher)
student_pipe = pipeline("text-generation", model=student)
prompts = ["Explain quantum computing:", "What is AI?"]
for prompt in prompts:
teacher_out = teacher_pipe(prompt, max_new_tokens=100)
student_out = student_pipe(prompt, max_new_tokens=100)
print(f"Prompt: {prompt}")
print(f"Teacher: {teacher_out[0]['generated_text']}")
print(f"Student: {student_out[0]['generated_text']}")
print(f"Match quality: {calculate_similarity(teacher_out, student_out):.2f}")
