蒸馏：用小模型继承大模型的能力

知识蒸馏是模型压缩领域最优雅的技术——让小模型（学生）学习大模型（教师）的内部表示，而非简单地学习标签。一个好的蒸馏方案可以让7B模型逼近70B模型的效果。

Logits蒸馏：经典方法

原理

教师模型的logits（softmax前的输出）包含了类别间的相似度信息——“软标签"比"硬标签"信息量更大：

class LogitsDistillationLoss(nn.Module):
    def __init__(self, temperature=2.0, alpha=0.5):
        self.temperature = temperature
        self.alpha = alpha  # 蒸馏loss与CE loss的权重比
    
    def forward(self, student_logits, teacher_logits, labels):
        # 蒸馏损失：KL散度
        soft_teacher = F.log_softmax(
            teacher_logits / self.temperature, dim=-1
        )
        soft_student = F.log_softmax(
            student_logits / self.temperature, dim=-1
        )
        distill_loss = F.kl_div(
            soft_student, soft_teacher.exp(),
            reduction="batchmean"
        ) * (self.temperature ** 2)
        
        # 标准交叉熵损失
        ce_loss = F.cross_entropy(student_logits, labels)
        
        return self.alpha * distill_loss + (1 - self.alpha) * ce_loss

温度参数的作用

温度 $T$ 控制软标签的"软度”：

  • $T=1$：标准softmax，概率分布较尖锐
  • $T=2-5$：分布更平滑，类别间关系更明显
  • $T \to \infty$：均匀分布

实践中 $T=2-4$ 效果最好。温度的平方项补偿了梯度缩放——高温softmax的梯度会被 $1/T^2$ 缩小。

在线蒸馏vs离线蒸馏

离线蒸馏：先训练好教师模型，再蒸馏学生模型。简单稳定但教师的错误会被继承。

在线蒸馏：教师和学生同时训练，教师不断更新：

class OnlineDistillation:
    def __init__(self, teacher, student, alpha=0.5):
        self.teacher = teacher
        self.student = student
        self.alpha = alpha
    
    def train_step(self, batch):
        # 教师前向（不更新梯度）
        with torch.no_grad():
            teacher_logits = self.teacher(batch)
        
        # 学生前向
        student_logits = self.student(batch)
        
        # 蒸馏损失
        distill_loss = self._distill_loss(student_logits, teacher_logits)
        ce_loss = F.cross_entropy(student_logits, batch["labels"])
        
        loss = self.alpha * distill_loss + (1 - self.alpha) * ce_loss
        loss.backward()

特征蒸馏：学习中间表示

原理

Logits蒸馏只利用了最终输出，特征蒸馏还利用了中间层的表示：

class FeatureDistillation(nn.Module):
    def __init__(self, teacher_layers, student_layers, projections):
        """
        teacher_layers: 教师模型的层索引列表
        student_layers: 对应的学生模型层索引
        projections: 投影层（将学生维度映射到教师维度）
        """
        self.teacher_layers = teacher_layers
        self.student_layers = student_layers
        self.projections = projections
    
    def forward(self, teacher_features, student_features):
        total_loss = 0
        for t_layer, s_layer, proj in zip(
            self.teacher_layers, self.student_layers, self.projections
        ):
            # 获取中间层特征
            t_feat = teacher_features[t_layer]  # [B, T, d_teacher]
            s_feat = student_features[s_layer]  # [B, T, d_student]
            
            # 投影到相同维度
            s_proj = proj(s_feat)  # [B, T, d_teacher]
            
            # MSE损失
            total_loss += F.mse_loss(s_proj, t_feat)
        
        return total_loss / len(self.teacher_layers)

层选择策略

不是所有层都适合做蒸馏。经验表明：

  • Transformer的FFN输出层效果最好
  • 注意力权重蒸馏效果一般
  • 嵌入层蒸馏帮助不大
  • 最后2-3层最重要

序列级蒸馏：生成任务的蒸馏

对于生成任务（而非分类），logits蒸馏不直接适用。序列级蒸馏让教师先生成回答，学生用这些回答做训练数据：

def sequence_level_distillation(teacher, student, prompts):
    # 1. 教师生成回答
    training_data = []
    for prompt in prompts:
        # 多次采样增加多样性
        for _ in range(3):
            response = teacher.generate(
                prompt, temperature=0.7, top_p=0.9
            )
            training_data.append({
                "prompt": prompt,
                "response": response
            })
    
    # 2. 学生在教师输出上训练
    for data in training_data:
        loss = student.compute_loss(
            data["prompt"],
            data["response"]
        )
        loss.backward()

带拒绝采样的蒸馏

教师生成的回答质量参差不齐，通过拒绝采样保留高质量回答：

def distillation_with_rejection_sampling(teacher, student, prompts, reward_model):
    for prompt in prompts:
        # 生成多个候选
        candidates = [teacher.generate(prompt) for _ in range(4)]
        
        # 用奖励模型筛选
        scores = [reward_model.score(prompt, c) for c in candidates]
        best = candidates[argmax(scores)]
        
        # 只用最好的回答训练学生
        if max(scores) > threshold:
            student.train_on(prompt, best)

多教师蒸馏

集成教师

多个教师模型的蒸馏可以让学生继承不同教师的优势：

class MultiTeacherDistillation:
    def __init__(self, teachers, student):
        self.teachers = teachers  # 多个教师模型
        self.student = student
    
    def distill(self, batch):
        # 各教师独立预测
        teacher_logits = [
            teacher(batch).logits for teacher in self.teachers
        ]
        
        # 加权融合
        weights = [0.4, 0.3, 0.3]  # 根据教师能力分配
        fused_logits = sum(
            w * logits for w, logits in zip(weights, teacher_logits)
        )
        
        # 蒸馏
        distill_loss = kl_div(
            self.student(batch).logits / T,
            fused_logits / T
        ) * T ** 2
        
        return distill_loss

教师选择策略

不同教师擅长不同任务，可以根据输入动态选择教师：

def dynamic_teacher_selection(input, teachers):
    # 根据输入特征选择最合适的教师
    if is_code_task(input):
        return teachers["code"]      # DeepSeek-Coder
    elif is_math_task(input):
        return teachers["math"]      # Qwen-Math
    else:
        return teachers["general"]   # Llama

蒸馏效果与局限

典型效果

从Llama-3-70B蒸馏到Llama-3-8B：

方法MMLUHumanEvalGSM8K
基线（仅SFT）65.272.078.5
+Logits蒸馏68.775.381.2
+特征蒸馏70.176.882.8
+序列蒸馏71.578.284.1
教师模型82.088.592.0

学生模型恢复了约70%的教师能力提升。

蒸馏的局限

  1. 能力天花板：学生模型不可能超过教师模型
  2. 容量瓶颈：如果学生模型太小，即使完美蒸馏也无法达到教师水平
  3. 模态依赖：蒸馏效果在不同任务上不一致
  4. 过拟合风险：学生可能过度模仿教师的具体输出而非泛化能力

实践建议

蒸馏配置推荐

# 推荐的蒸馏配置
distill_config = {
    "temperature": 3.0,          # 生成任务用较高温度
    "alpha": 0.7,                # 蒸馏loss权重
    "distill_layers": [-3, -2, -1],  # 最后3层做特征蒸馏
    "learning_rate": 5e-5,       # 比正常训练略低
    "batch_size": 32,
    "epochs": 3,
    "warmup_ratio": 0.1,
}

数据选择

  • 使用多样化的高质量prompt
  • 包含教师擅长的领域数据
  • 混合一定比例的通用数据防止过拟合

结语

知识蒸馏是大模型民主化部署的关键技术——让小模型也能拥有接近大模型的能力。随着蒸馏技术的精细化发展（特征蒸馏、多教师蒸馏、渐进式蒸馏），学生模型与教师模型的差距正在缩小。在边缘部署和实时推理场景中，蒸馏后的模型往往是唯一可行的方案。