为什么LLM需要蒸馏?

训练一个超大模型(如700B)然后部署它,成本极其高昂。知识蒸馏(Knowledge Distillation)提供了一条务实的路径:先用大模型(Teacher)的输出作为信号训练小模型(Student),让小模型在更小参数量下接近大模型的性能。

2026年,蒸馏已经成为大模型工程化的标准环节。DeepSeek-V3、Qwen-3等模型都大量使用了蒸馏技术,将超大模型的能力迁移到可部署的尺寸。

蒸馏的理论基础

软标签的信息优势

硬标签(one-hot)只包含"正确答案"的信息,而软标签(softmax概率分布)还包含"错误答案之间的关系"。例如,在分类"猫"时,软标签可能同时给出"狗"的概率0.1——这告诉Student模型"猫和狗在某种特征上是相似的"。

这种"暗知识"(Dark Knowledge)是蒸馏有效性的核心。Teacher模型的输出分布包含了其学到的类别间关系,这些信息在硬标签中完全丢失。

温度参数

温度T控制软标签的"软度":

soft_label = softmax(logits / T)

高温使分布更平滑(暴露更多暗知识),低温使分布更尖锐(接近one-hot)。实践中T通常设置为2-10。

LLM蒸馏的主要方法

1. Logit级蒸馏

最经典的蒸馏方式——Student直接学习Teacher的输出概率分布:

def logit_distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    """
    student_logits, teacher_logits: [batch, seq_len, vocab_size]
    labels: [batch, seq_len]
    T: 温度参数
    alpha: 蒸馏损失权重
    """
    # 蒸馏损失:KL散度
    student_log_probs = F.log_softmax(student_logits / T, dim=-1)
    teacher_probs = F.softmax(teacher_logits / T, dim=-1)
    
    distill_loss = F.kl_div(
        student_log_probs.reshape(-1, student_logits.size(-1)),
        teacher_probs.reshape(-1, teacher_logits.size(-1)),
        reduction='batchmean'
    ) * (T ** 2)  # 梯度缩放补偿
    
    # 任务损失:交叉熵
    task_loss = F.cross_entropy(
        student_logits.reshape(-1, student_logits.size(-1)),
        labels.reshape(-1)
    )
    
    return alpha * distill_loss + (1 - alpha) * task_loss

关键点:

  • KL散度损失需要乘以 T² 来补偿温度对梯度的影响
  • alpha控制蒸馏与任务学习的平衡
  • 需要Teacher和Student的词表对齐

2. 序列级蒸馏(Sequence-Level KD)

不让Student逐token模仿Teacher,而是让Student学习Teacher生成的完整序列。具体做法是先用Teacher生成大量数据,然后用这些数据训练Student:

def sequence_level_distillation(teacher_model, student_model, prompts):
    """序列级蒸馏"""
    # 1. 用Teacher生成数据
    synthetic_data = []
    for prompt in prompts:
        teacher_output = teacher_model.generate(
            prompt, max_tokens=2048, temperature=0.7
        )
        synthetic_data.append({
            'input': prompt,
            'output': teacher_output
        })
    
    # 2. 用合成数据训练Student
    for data in synthetic_data:
        loss = student_model.compute_loss(data['input'], data['output'])
        loss.backward()
        optimizer.step()

序列级蒸馏的优势在于:

  • 不需要Teacher和Student架构对齐
  • 可以离线生成数据,训练时不需要Teacher
  • Student可以学习Teacher的生成风格和格式

3. 特征级蒸馏

除了输出层,还可以对齐Teacher和Student中间层的隐藏状态:

class FeatureDistillation:
    def __init__(self, teacher_model, student_model, layer_mapping):
        """
        layer_mapping: {student_layer_idx: teacher_layer_idx}
        """
        self.teacher = teacher_model
        self.student = student_model
        self.layer_mapping = layer_mapping
        self.projections = nn.ModuleDict({
            str(s_layer): nn.Linear(student_dim, teacher_dim)
            for s_layer, (student_dim, teacher_dim) in ...
        })
    
    def compute_loss(self, student_hidden_states, teacher_hidden_states):
        """计算特征对齐损失"""
        feature_loss = 0
        for s_idx, t_idx in self.layer_mapping.items():
            s_feat = student_hidden_states[s_idx]
            t_feat = teacher_hidden_states[t_idx].detach()
            
            # 投影到相同维度
            s_proj = self.projections[str(s_idx)](s_feat)
            
            # 余弦相似度损失
            feature_loss += 1 - F.cosine_similarity(
                s_proj.flatten(1), t_feat.flatten(1), dim=-1
            ).mean()
        
        return feature_loss

特征级蒸馏强制Student学习Teacher的内部表示,而不仅仅是输出行为。这对于对齐Student的推理过程特别有效。

4. 注意力蒸馏

让Student模仿Teacher的注意力模式:

def attention_distillation_loss(student_attn_weights, teacher_attn_weights, layer_mapping):
    """注意力分布蒸馏"""
    loss = 0
    for s_idx, t_idx in layer_mapping.items():
        s_attn = student_attn_weights[s_idx]
        t_attn = teacher_attn_weights[t_idx].detach()
        
        # 对齐注意力头数
        if s_attn.shape[1] != t_attn.shape[1]:
            s_attn = merge_heads(s_attn, t_attn.shape[1])
        
        # MSE损失
        loss += F.mse_loss(s_attn, t_attn)
    
    return loss

注意力蒸馏传递了Teacher"关注哪里"的知识,这在需要复杂推理的任务中特别有价值。

在线蒸馏 vs 离线蒸馏

离线蒸馏

传统的两阶段方法:先训练好Teacher,再蒸馏Student。

  • 优势:Teacher质量高,蒸馏过程简单
  • 劣势:Teacher的训练成本需要单独承担

在线蒸馏

Teacher和Student同时训练,Teacher在训练过程中持续指导Student:

def online_distillation_train_step(teacher, student, batch, optimizer, T=4.0):
    """在线蒸馏单步训练"""
    # Teacher前向(不计算梯度)
    with torch.no_grad():
        teacher_logits = teacher(batch['input'])
    
    # Student前向
    student_logits = student(batch['input'])
    
    # 蒸馏损失
    loss = logit_distillation_loss(
        student_logits, teacher_logits, batch['labels'], T=T
    )
    
    # Student反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # Teacher也更新(用真实标签)
    teacher_loss = F.cross_entropy(
        teacher_logits.reshape(-1, teacher_logits.size(-1)),
        batch['labels'].reshape(-1)
    )
    teacher_loss.backward()
    teacher_optimizer.step()

在线蒸馏的优势是Teacher在训练过程中逐渐变强,Student也能跟上Teacher的进步。

蒸馏的工程实践

数据选择

不是所有数据都适合蒸馏。研究表明,蒸馏效果最好的数据具有以下特征:

  • Teacher在这些数据上的置信度适中(不是太高也不是太低)
  • 覆盖Teacher能力的多个方面
  • 难度分布与学生模型容量匹配

渐进式蒸馏

从简单任务开始,逐步增加难度:

def progressive_distillation(student, teacher, curriculum):
    """
    curriculum: [(difficulty, n_steps), ...] 由易到难
    """
    for difficulty, n_steps in curriculum:
        data = load_data(difficulty=difficulty)
        for step in range(n_steps):
            train_step(student, teacher, data)

多Teacher蒸馏

使用多个Teacher的不同能力来指导Student:

def multi_teacher_distillation(student, teachers, weights, data):
    """多Teacher加权蒸馏"""
    teacher_logits = []
    for teacher in teachers:
        with torch.no_grad():
            teacher_logits.append(teacher(data))
    
    # 加权平均Teacher的logits
    avg_logits = sum(w * l for w, l in zip(weights, teacher_logits))
    
    # 蒸馏
    loss = logit_distillation_loss(student(data), avg_logits, data['labels'])
    return loss

例如,用一个Teacher教数学推理,另一个教代码生成,第三个教通用对话。

2026年的蒸馏前沿

蒸馏+强化学习

将RLHF/RLAIF与蒸馏结合——Teacher用RL优化后蒸馏给Student,让Student直接获得RL对齐的效果而不需要自己做RL:

def rl_distillation(student, rl_teacher, reward_model, data):
    """RL对齐蒸馏"""
    # Teacher用RL生成对齐后的输出
    teacher_output = rl_teacher.generate(
        data, 
        reward_fn=reward_model,
        method='ppo'
    )
    # Student学习Teacher的RL对齐输出
    loss = student.compute_loss(data, teacher_output)
    return loss

思维链蒸馏

将Teacher的思维链(CoT)推理过程蒸馏给Student。Student不仅学习最终答案,还学习Teacher的推理步骤。

自蒸馏

模型自己作为Teacher和Student——先用完整模型训练,然后将知识蒸馏到自身的子网络。这在持续学习中特别有用。

结语

知识蒸馏是大模型从实验室走向生产的桥梁。从Logit级到序列级、从单Teacher到多Teacher、从静态蒸馏到在线蒸馏,蒸馏技术在2026年已经形成了丰富的工具箱。随着模型规模继续增长,蒸馏的价值只会越来越高——它是让AGI能力"下沉"到可部署规模的关键技术。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞,认知差异创造无限可能。