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.
Instrucciones de origen · Vista previa de solo lectura
name
knowledge-distillation
description
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.
"""
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
1
1
1
# Soft loss: KL divergence between student and teacher
1
1
'batchmean'
2
# Combined loss
return
1
# 3. Training loop
for
in
# Teacher forward (no grad)
with
# Student forward
# Compute distillation loss
'labels'
2.0
0.7
# 70% soft, 30% hard
# Backward and optimize
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}")
