语言模型蒸馏简明教程
知识蒸馏是一种训练较小神经网络以执行与较大网络相同功能的技术。
梯形图转SCL | 博途AI辅助编程文档 | AI模型价格对比 | AI工具导航 | ONNX模型库 | Vibe Coding教程 | PLC在线仿真器 | Tripo 3D | Meshy AI | ElevenLabs | KlingAI | ArtSpace | Phot.AI | InVideo
知识蒸馏是一种训练较小神经网络以执行与较大网络相同功能的技术。基本思想很简单：训练一个小型"学生"模型来复制大型"教师"模型的行为。这让你可以将多年的训练和数十亿参数压缩成可以实际部署的东西。但要使其良好工作，需要理解概率匹配、中间表示和训练动态。
P.S. 缩略图是使用GPT 5.4生成的
1、蒸馏的工作原理
Hinton和他的同事在注意到一些有用的东西后创建了知识蒸馏：当训练模型进行预测时，它输出的概率分布包含比最终答案更多的信息。假设教师模型分配80%概率给正确类别，15%给相似类别，5%给不相关类别。这个分布告诉你关于类别之间的关系，而简单的正确/错误标签会丢弃这些信息。学生学习的不仅是预测什么，还有教师如何思考替代方案。
蒸馏损失函数结合了两个部分。首先，它使用KL散度衡量学生与教师软预测的匹配程度。其次，它确保学生仍然从实际标签中使用交叉熵损失学习。你用一个称为alpha的参数来控制它们之间的平衡。
import torch
import torch.nn as nn
import torch.nn.functional as F
class DistillationLoss(nn.Module):
def __init__(self, temperature=3.0, alpha=0.7):
"""
Args:
temperature: Controls softness of probability distributions
alpha: Weight balancing distillation vs hard label loss
"""
super().__init__()
self.temperature = temperature
self.alpha = alpha
self.kl_div = nn.KLDivLoss(reduction='batchmean')
self.ce_loss = nn.CrossEntropyLoss()
def forward(self, student_logits, teacher_logits, labels):
# Soften distributions with temperature scaling
student_soft = F.log_softmax(student_logits / self.temperature, dim=1)
teacher_soft = F.softmax(teacher_logits / self.temperature, dim=1)
# Distillation loss (KL divergence between soft distributions)
distillation_loss = self.kl_div(student_soft, teacher_soft) * (self.temperature ** 2)
# Standard cross-entropy with hard labels
student_loss = self.ce_loss(student_logits, labels)
# Combined objective
return self.alpha * distillation_loss + (1 - self.alpha) * student_loss
温度在这里很重要。较高的温度使概率分布"更软"，揭示更多关于类别关系的信息。在温度1时，你得到正常的softmax。随着温度升高，分布变得更加均匀。最佳温度通常在2到5之间。损失中的温度平方项保持不同温度设置下的梯度一致性。
2、处理大小差异
教师和学生大小之间的差距从根本上限制了蒸馏的效果。如果学生太小，它无法捕捉教师所知道的东西。如果太大，你无法获得太多压缩。你需要找到学生足够小以便有用但又足够大以学习重要模式的点。
不同的方法以不同的方式处理这个差距。基本蒸馏只匹配输出分布。中间蒸馏还对齐隐藏层表示。这对于transformer更有效，其中中间注意力模式编码重要的语言知识。耐心知识蒸馏更进一步，匹配层之间的关系，而不仅仅是单个层。
class IntermediateDistillation(nn.Module):
def __init__(self, student_dim, teacher_dim, num_student_layers, num_teacher_layers):
super().__init__()
self.num_student_layers = num_student_layers
self.num_teacher_layers = num_teacher_layers
# Layer mapping strategy: map student layers to teacher layers
self.layer_mapping = self._create_layer_mapping()
# Projection layers if dimensions don't match
if student_dim != teacher_dim:
self.projections = nn.ModuleList([
nn.Linear(student_dim, teacher_dim)
for _ in range(num_student_layers)
])
else:
self.projections = None
def _create_layer_mapping(self):
# Map student layers uniformly across teacher layers
# For 6 student, 12 teacher: [1, 3, 5, 7, 9, 11]
step = self.num_teacher_layers / self.num_student_layers
return [int(i * step) for i in range(self.num_student_layers)]
def forward(self, student_hidden_states, teacher_hidden_states):
"""
Args:
student_hidden_states: List of tensors [batch, seq_len, student_dim]
teacher_hidden_states: List of tensors [batch, seq_len, teacher_dim]
"""
total_loss = 0
for student_idx, teacher_idx in enumerate(self.layer_mapping):
student_hidden = student_hidden_states[student_idx]
teacher_hidden = teacher_hidden_states[teacher_idx]
# Project student to teacher dimension if needed
if self.projections is not None:
student_hidden = self.projections[student_idx](student_hidden)
# MSE loss between intermediate representations
layer_loss = F.mse_loss(student_hidden, teacher_hidden.detach())
total_loss += layer_loss
return total_loss / len(self.layer_mapping)
3、转移注意力模式
Transformer中的注意力机制捕获输入token之间的复杂依赖关系，编码结构、关系和上下文。蒸馏这些注意力模式比蒸馏输出更难，因为注意力矩阵很大且对架构差异敏感。注意力转移方法通常关注保留注意力的结构而不是精确权重。
注意力转移损失衡量教师和学生注意力分布之间的距离，通常使用均方误差或KL散度。你需要处理不同数量的注意力头。一些方法在计算损失之前跨头平均注意力。其他方法维护头特定的对齐，这保留更多细节但需要仔细映射。
class AttentionTransfer(nn.Module):
def __init__(self, student_heads, teacher_heads, use_head_mapping=True):
super().__init__()
self.student_heads = student_heads
self.teacher_heads = teacher_heads
self.use_head_mapping = use_head_mapping
if use_head_mapping and student_heads != teacher_heads:
# Learn which teacher heads to map to which student heads
self.head_mapping = nn.Parameter(
torch.randn(student_heads, teacher_heads)
)
def forward(self, student_attentions, teacher_attentions):
"""
Args:
student_attentions: [batch, num_heads, seq_len, seq_len]
teacher_attentions: [batch, num_heads, seq_len, seq_len]
"""
batch_size, _, seq_len, _ = student_attentions.shape
if self.use_head_mapping and hasattr(self, 'head_mapping'):
# Apply learned head mapping
mapping_weights = F.softmax(self.head_mapping, dim=1)
# [student_heads, teacher_heads] × [batch, teacher_heads, seq, seq]
teacher_mapped = torch.einsum(
'st,bthw->bshw',
mapping_weights,
teacher_attentions
)
else:
# Simple averaging if heads match or no mapping desired
if self.student_heads == self.teacher_heads:
teacher_mapped = teacher_attentions
else:
# Average teacher heads to match student count
teacher_mapped = teacher_attentions.reshape(
batch_size, self.student_heads, -1, seq_len, seq_len
).mean(dim=2)
# MSE between attention distributions
attention_loss = F.mse_loss(student_attentions, teacher_mapped.detach())
return attention_loss
4、渐进式蒸馏
渐进式蒸馏通过使用中间教师来解决将非常大的教师蒸馏到非常小的学生的问题。你不是直接从GPT-3规模模型到移动友好大小，而是创建一系列逐渐变小的教师，每个都从前一个蒸馏。这种分阶段方法让每个学生从更接近其自身容量的教师学习，这减少了知识差距并提高了最终性能。
蒸馏中的课程学习意味着仔细排序训练示例。早期，学生从教师更自信的较容易示例中学习。随着训练进行，你引入更多模糊案例，其中教师的软标签提供最大信息。你可以根据预测熵、损失大小或示例复杂性定义这个课程。
class ProgressiveDistillationTrainer:
def __init__(self, teachers, student, device='cuda'):
"""
Args:
teachers: List of teacher models ordered from largest to smallest
student: Student model to train
"""
self.teachers = teachers
self.student = student
self.device = device
# Move all models to device and set teachers to eval
for teacher in self.teachers:
teacher.to(device)
teacher.eval()
self.student.to(device)
def get_curriculum_weight(self, epoch, total_epochs):
# Linearly increase difficulty over training
return min(1.0, epoch / (total_epochs * 0.7))
def compute_example_difficulty(self, teacher_logits):
# Use entropy of teacher predictions as difficulty measure
probs = F.softmax(teacher_logits, dim=-1)
entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=-1)
return entropy
def progressive_distill(self, dataloader, stage, optimizer,
temperature=3.0, epochs=10):
"""
Distill from teachers[stage] into student or next teacher
"""
current_teacher = self.teachers[stage]
criterion = DistillationLoss(temperature=temperature)
for epoch in range(epochs):
curriculum_weight = self.get_curriculum_weight(epoch, epochs)
for batch in dataloader:
inputs, labels = batch
inputs = inputs.to(self.device)
labels = labels.to(self.device)
# Get teacher predictions
with torch.no_grad():
teacher_logits = current_teacher(inputs)
difficulties = self.compute_example_difficulty(teacher_logits)
# Filter or weight examples based on curriculum
difficulty_threshold = torch.quantile(
difficulties, curriculum_weight
)
example_weights = (difficulties <= difficulty_threshold).float()
# Student forward pass
student_logits = self.student(inputs)
# Compute weighted distillation loss
loss = criterion(student_logits, teacher_logits, labels)
weighted_loss = (loss * example_weights.mean())
# Optimization step
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
def train_all_stages(self, dataloader, optimizer, epochs_per_stage=10):
"""
Execute progressive distillation through all teacher stages
"""
for stage in range(len(self.teachers)):
print(f"Stage {stage}: Distilling from teacher {stage}")
self.progressive_distill(
dataloader, stage, optimizer, epochs=epochs_per_stage
)
return self.student
5、任务特定和多任务蒸馏
通用蒸馏训练学生以匹配教师在所有任务上的行为。任务特定蒸馏针对特定应用进行优化。这让你可以更积极地压缩，因为学生只需要与目标任务相关的知识。例如，将通用语言模型蒸馏到情感分类器可以实现更高的压缩，同时保持或超过任务性能。
多任务蒸馏通过同时训练学生多个相关任务来扩展这一点。教师可能是任务特定专家模型的集合，学生学习在一个架构中处理所有任务。当任务共享底层模式时，这很有效，让学生开发跨任务泛化的共享表示。
class MultiTaskDistillation(nn.Module):
def __init__(self, task_weights=None):
super().__init__()
self.task_weights = task_weights or {}
def forward(self, student_outputs, teacher_outputs, task_names, labels):
"""
Args:
student_outputs: Dict mapping task names to student logits
teacher_outputs: Dict mapping task names to teacher logits
task_names: List of tasks in current batch
labels: Dict mapping task names to ground truth labels
"""
total_loss = 0
task_losses = {}
for task in task_names:
# Task-specific distillation loss
criterion = DistillationLoss(
temperature=self.get_task_temperature(task),
alpha=self.get_task_alpha(task)
)
task_loss = criterion(
student_outputs[task],
teacher_outputs[task],
labels[task]
)
# Weight by task importance
weight = self.task_weights.get(task, 1.0)
total_loss += weight * task_loss
task_losses[task] = task_loss.item()
return total_loss, task_losses
def get_task_temperature(self, task):
# Different tasks may benefit from different temperatures
temperature_map = {
'sentiment': 2.0, # Lower for classification
'nli': 3.0, # Higher for complex reasoning
'qa': 4.0, # Highest for generation tasks
}
return temperature_map.get(task, 3.0)
def get_task_alpha(self, task):
# Balance between distillation and hard labels per task
alpha_map = {
'sentiment': 0.5, # More weight on hard labels
'nli': 0.7, # Balanced
'qa': 0.9, # Heavy distillation weight
}
return alpha_map.get(task, 0.7)
6、数据增强和合成数据
蒸馏的效果很大程度上取决于训练数据的多样性和质量。基本蒸馏使用训练教师的相同数据集。增强蒸馏生成合成示例以让学生接触更多教师行为。教师为未标记数据生成标签，大大扩展训练集。这在针对挑战性案例或代表性不足模式的特定任务增强策略中特别有效。
class DataAugmentedDistillation:
def __init__(self, teacher, student, base_dataset):
self.teacher = teacher
self.student = student
self.base_dataset = base_dataset
def generate_synthetic_examples(self, num_examples, augmentation_fn):
"""
Generate synthetic training examples using the teacher
"""
synthetic_data = []
self.teacher.eval()
with torch.no_grad():
for _ in range(num_examples):
# Sample from base dataset and augment
base_example = self.base_dataset[
torch.randint(len(self.base_dataset), (1,)).item()
]
augmented_input = augmentation_fn(base_example)
# Generate teacher predictions
teacher_logits = self.teacher(augmented_input)
synthetic_data.append({
'input': augmented_input,
'teacher_logits': teacher_logits.cpu(),
'source': 'synthetic'
})
return synthetic_data
def hard_example_mining(self, dataloader, percentile=90):
"""
Identify examples where student struggles most
"""
self.student.eval()
self.teacher.eval()
example_difficulties = []
with torch.no_grad():
for batch in dataloader:
inputs, labels = batch
student_logits = self.student(inputs)
teacher_logits = self.teacher(inputs)
# Measure disagreement as difficulty proxy
disagreement = F.kl_div(
F.log_softmax(student_logits, dim=-1),
F.softmax(teacher_logits, dim=-1),
reduction='none'
).sum(dim=-1)
example_difficulties.extend(disagreement.cpu().numpy())
# Return indices of hardest examples
threshold = np.percentile(example_difficulties, percentile)
hard_indices = np.where(
np.array(example_difficulties) >= threshold
)[0]
return hard_indices
7、训练稳定性和优化
蒸馏的优化景观不同于标准监督学习。教师的软标签提供比独热标签更平滑的训练信号，这可以加速收敛，但如果管理不当也会导致不稳定。温度直接影响梯度大小，不当调优可能导致梯度爆炸或梯度消失。
学习率调度对于成功至关重要。常见策略使用预热阶段，其中学习率逐渐增加，让学生在完全蒸馏之前稳定。在主要训练期间，保持中等学习率和高温度。最后，降低温度和学习率的微调阶段完善性能。
class DistillationOptimizer:
def __init__(self, student, initial_lr=1e-4, warmup_steps=1000):
self.student = student
self.initial_lr = initial_lr
self.warmup_steps = warmup_steps
self.global_step = 0
# Use AdamW with weight decay for better generalization
self.optimizer = torch.optim.AdamW(
student.parameters(),
lr=initial_lr,
betas=(0.9, 0.999),
weight_decay=0.01
)
self.scheduler = self._create_scheduler()
def _create_scheduler(self):
# Cosine schedule with warmup
from torch.optim.lr_scheduler import LambdaLR
def lr_lambda(step):
if step < self.warmup_steps:
# Linear warmup
return step / self.warmup_steps
else:
# Cosine decay
progress = (step - self.warmup_steps) / (10000 - self.warmup_steps)
return 0.5 * (1 + np.cos(np.pi * progress))
return LambdaLR(self.optimizer, lr_lambda)
def step(self, loss):
# Gradient clipping for stability
torch.nn.utils.clip_grad_norm_(self.student.parameters(), max_norm=1.0)
self.optimizer.step()
self.scheduler.step()
self.global_step += 1
return self.scheduler.get_last_lr()[0]
def get_temperature_schedule(self, max_steps):
"""
Dynamic temperature scheduling during training
"""
if self.global_step < self.warmup_steps:
# Start with lower temperature during warmup
return 2.0
elif self.global_step < max_steps * 0.8:
# Higher temperature for main distillation
return 4.0
else:
# Reduce temperature for fine-tuning
return 2.0
8、评估指标
评估蒸馏模型需要的不仅仅是准确性。你需要评估压缩比、推理延迟、内存占用和能源消耗。蒸馏效率指标捕获大小减少和性能保留之间的权衡，通常计算为精度保留与压缩比的比率。
class DistillationEvaluator:
def __init__(self, teacher, student, test_loader, device='cuda'):
self.teacher = teacher
self.student = student
self.test_loader = test_loader
self.device = device
def compute_compression_metrics(self):
teacher_params = sum(p.numel() for p in self.teacher.parameters())
student_params = sum(p.numel() for p in self.student.parameters())
compression_ratio = teacher_params / student_params
return {
'teacher_parameters': teacher_params,
'student_parameters': student_params,
'compression_ratio': compression_ratio
}
def measure_inference_speed(self, num_samples=100):
import time
self.teacher.eval()
self.student.eval()
# Sample random inputs
sample_inputs = []
for batch in self.test_loader:
sample_inputs.append(batch[0][:1].to(self.device))
if len(sample_inputs) >= num_samples:
break
# Teacher inference time
teacher_times = []
with torch.no_grad():
for inputs in sample_inputs:
start = time.perf_counter()
_ = self.teacher(inputs)
teacher_times.append(time.perf_counter() - start)
# Student inference time
student_times = []
with torch.no_grad():
for inputs in sample_inputs:
start = time.perf_counter()
_ = self.student(inputs)
student_times.append(time.perf_counter() - start)
speedup = np.mean(teacher_times) / np.mean(student_times)
return {
'teacher_latency_ms': np.mean(teacher_times) * 1000,
'student_latency_ms': np.mean(student_times) * 1000,
'speedup_factor': speedup
}
def compute_agreement_metrics(self):
"""
Measure how well student predictions agree with teacher
"""
self.teacher.eval()
self.student.eval()
total_kl = 0
total_top1_agreement = 0
total_samples = 0
with torch.no_grad():
for inputs, labels in self.test_loader:
inputs = inputs.to(self.device)
teacher_logits = self.teacher(inputs)
student_logits = self.student(inputs)
# KL divergence
kl = F.kl_div(
F.log_softmax(student_logits, dim=-1),
F.softmax(teacher_logits, dim=-1),
reduction='batchmean'
)
total_kl += kl.item() * inputs.size(0)
# Top-1 agreement
teacher_preds = teacher_logits.argmax(dim=-1)
student_preds = student_logits.argmax(dim=-1)
agreement = (teacher_preds == student_preds).float().mean()
total_top1_agreement += agreement.item() * inputs.size(0)
total_samples += inputs.size(0)
return {
'average_kl_divergence': total_kl / total_samples,
'top1_agreement': total_top1_agreement / total_samples
}
def full_evaluation(self):
"""
Comprehensive evaluation of distillation quality
"""
metrics = {}
# Compression metrics
metrics.update(self.compute_compression_metrics())
# Speed metrics
metrics.update(self.measure_inference_speed())
# Agreement metrics
metrics.update(self.compute_agreement_metrics())
# Efficiency score: accuracy preservation per unit compression
metrics['efficiency_score'] = (
metrics['top1_agreement'] * metrics['compression_ratio']
)
return metrics
9、高级技术
最近的进展引入了几种超越传统知识转移的技术。在线蒸馏同时训练教师和学生，教师持续更新而不是保持冻结。这种协同进化可以导致互利学习，学生的进度为教师更新提供信息。自蒸馏将蒸馏应用于相同架构，使用集成预测或不同初始化的模型作为教师，即使没有压缩也可以提高性能。
重生网络是自蒸馏的一种极端形式，其中与教师相同架构的学生通常超过教师的性能。这表明蒸馏提供的不仅仅是压缩——它提供了改进的优化景观和隐式正则化。迭代应用重生蒸馏，其中每一代作为下一代的教师，可以逐步提高性能直到收敛。
class OnlineDistillation(nn.Module):
def __init__(self, teacher, student, teacher_update_freq=10):
super().__init__()
self.teacher = teacher
self.student = student
self.teacher_update_freq = teacher_update_freq
self.step_count = 0
# Initialize teacher with student parameters
self.teacher.load_state_dict(student.state_dict())
# Separate optimizers for teacher and student
self.teacher_optimizer = torch.optim.AdamW(
teacher.parameters(), lr=1e-5
)
self.student_optimizer = torch.optim.AdamW(
student.parameters(), lr=1e-4
)
def train_step(self, inputs, labels, temperature=3.0):
# Student learning from current teacher
with torch.no_grad():
teacher_logits = self.teacher(inputs)
student_logits = self.student(inputs)
criterion = DistillationLoss(temperature=temperature)
student_loss = criterion(student_logits, teacher_logits, labels)
self.student_optimizer.zero_grad()
student_loss.backward()
self.student_optimizer.step()
# Periodically update teacher
self.step_count += 1
if self.step_count % self.teacher_update_freq == 0:
# Teacher learns from student's predictions
with torch.no_grad():
student_logits_detached = self.student(inputs)
teacher_logits = self.teacher(inputs)
teacher_loss = criterion(teacher_logits, student_logits_detached, labels)
self.teacher_optimizer.zero_grad()
teacher_loss.backward()
self.teacher_optimizer.step()
return student_loss.item()
10、结束语
语言模型蒸馏实现了从大型昂贵教师到高效学生的知识转移。技术包括输出分布匹配、中间表示对齐、注意力转移和渐进式多阶段蒸馏。成功取决于温度缩放、课程学习、优化动态和架构映射。
该领域通过在线蒸馏、跨模态知识转移和特定任务压缩方面的创新不断发展。随着语言模型变得更大，蒸馏对于使最先进的自然语言理解为更多人所必需。核心见解很直接：编码在数十亿参数中的知识可以压缩到数百万，同时保留驱动智能行为的本质模式。
原文链接：Language Model Distillation
汇智网翻译整理，转载请标明出处
