模型蒸馏(Knowledge Distillation)
深度学习模型的效果与规模高度相关——参数量越大,通常性能越强。然而大模型在实际部署时面临一系列问题:
所属专题:机器学习基础 (ml-basics·09)
模型蒸馏(Knowledge Distillation)
一、起源:为什么需要这个技术
深度学习模型的效果与规模高度相关——参数量越大,通常性能越强。然而大模型在实际部署时面临一系列问题:
- 推理成本高:GPT-4、Claude Opus 这类千亿级模型单次推理耗费大量算力,成本难以承受。
- 延迟不可接受:移动端、嵌入式设备、实时对话场景对响应时间要求毫秒级。
- 内存与显存受限:手机、IoT、边缘设备根本装不下大模型。
- 能耗与散热:大规模在线服务的电费和碳排放不可忽视。
问题的本质:我们希望”用小模型的成本,获得接近大模型的效果”。
传统压缩手段(剪枝 Pruning、量化 Quantization)主要在同一个模型内部做减法,而 Hinton 等人在 2015 年的论文 “Distilling the Knowledge in a Neural Network” 提出了一条新路径:让一个小模型去”学习”大模型的行为。这就是知识蒸馏的起点。
在 LLM 时代,蒸馏的重要性进一步凸显:Claude Haiku、Gemini Flash、GPT-4o-mini 等”小而强”的模型,背后都有蒸馏的影子。
二、概述
知识蒸馏(Knowledge Distillation, KD) 是一种模型压缩与知识迁移技术,核心思想是:
用一个已训练好的高性能大模型(教师模型 Teacher)去指导一个小模型(学生模型 Student)的训练,让学生在参数量远小于教师的情况下,尽可能复现教师的行为。
关键要素:
- 教师模型:能力强、体积大,通常已经充分训练好。
- 学生模型:结构更小、参数更少,是最终要部署的模型。
- 迁移信号:不仅仅是”正确答案”,更是教师模型输出中的概率分布、中间特征、注意力模式等丰富信息。
蒸馏本质上是一种从”硬知识”到”软知识”的迁移:传统监督学习只告诉模型”对/错”,蒸馏则告诉模型”这个答案有多可能,其他选项有多接近”。
三、核心原理
3.1 硬标签 vs 软标签
直观理解
传统监督学习使用 硬标签(Hard Label),例如一张猫的图片:
[猫: 1, 狗: 0, 汽车: 0, 飞机: 0]
这种表示只告诉模型”是猫”,信息量极低。
而教师模型输出的 软标签(Soft Label)(softmax 后的概率分布):
[猫: 0.85, 狗: 0.10, 汽车: 0.001, 飞机: 0.0001]
这里蕴含了额外信息:猫和狗更相似,和汽车/飞机差别很大。Hinton 把这种类别间的相似性信息称为 “暗知识”(Dark Knowledge)——它是教师模型在大量数据训练中习得的、无法通过 one-hot 标签表达的归纳偏置。
本质:logits 的”坍缩” vs “保留”
从数学上看,硬标签和软标签的区别,其实就是如何处理教师模型的 logits(softmax 前的原始分数):
| 数学形式 | 信息量 | |
|---|---|---|
| 硬标签 | argmax(logits) → one-hot | 只保留 Top-1,丢弃所有其他信息 |
| 软标签 | softmax(logits / T) → 概率分布 | 保留所有类别的相对关系 |
一图看清 logits 的三种”打开方式”:
教师模型
│
▼
logits: [8.2, 6.1, -2.3, -4.5] ← 原始分数
│
├──► argmax ──► [1, 0, 0, 0] = 硬标签(信息坍缩)
│
├──► softmax(T=1) ──► [0.89, 0.11, ~0, ~0] = 尖锐分布(暗知识被压平)
│
└──► softmax(T=4) ──► [0.55, 0.32, 0.09, 0.04] = 软标签(暗知识可见)
可以看到:
- “暗知识”本来就藏在原始 logits 的相对大小里——猫和狗的 logit 差 2.1,猫和汽车差 10.5,这个结构信息是教师的核心价值。
- T=1 的 softmax 把小 logit 压得接近 0,暗知识虽然存在但”看不见”。
- T>1 相当于把 logits 之间的差距按比例缩小,让原本被压平的类别关系重新显现出来。
所以更准确的表述是:蒸馏的核心,是用教师的 logits 分布本身作为监督信号,而不是把它坍缩成 argmax 后再用。温度参数 T 的作用,下一节详述。
关于”硬标签”的一个歧义
“硬标签”在不同上下文里指代不同来源,容易混淆:
- 传统监督学习中:硬标签 = 数据集的 ground truth(人工标注),并非任何模型的输出。
- 蒸馏中:可能同时存在两种硬标签
- 数据集真实标签(用于损失函数中的 项,保证学生不学偏)
- 教师的
argmax(logits)(若只做伪标签蒸馏,则等价于用教师做数据标注)
真正让蒸馏成为”蒸馏”的,是软标签这一路信号。
3.2 温度参数 T(Temperature)
原始 softmax 输出的分布通常非常”尖锐”:正确类别概率接近 1,其他接近 0,暗知识被压平了。为了放大这些小概率的信号,蒸馏引入 温度系数 T:
其中 是 logits(softmax 之前的原始分数)。
- T = 1:普通 softmax。
- T > 1:分布变得更平滑,小概率类别被放大,暗知识更明显。
- T → ∞:接近均匀分布。
训练时:教师和学生都用同一个较大的 T(如 4、10)来产生软分布。 推理时:学生恢复 T = 1,正常输出。
3.3 损失函数
学生模型的损失通常是两部分加权:
各项含义:
| 项 | 作用 |
|---|---|
| 学生对真实标签的交叉熵,保证不学偏 | |
| 学生输出与教师软分布的 KL 散度,学习暗知识 | |
| 两个损失的权重,通常 0.1 ~ 0.5 | |
| 补偿温度带来的梯度尺度变化,保证两项梯度量级相当 |
3.4 为什么有效
- 软标签是更丰富的监督信号:一次前向传播的信息量远超 one-hot。
- 正则化效应:教师提供的平滑分布相当于对学生做隐式正则,减少过拟合。
- 归纳偏置迁移:大模型学到的”世界结构”(哪些概念相近、哪些特征重要)被传递给小模型。
- 优化路径更平滑:学生不需要从零探索,教师提供了更好的”引路人”。
四、蒸馏机制(常见方法与变体)
按照迁移的知识类型,蒸馏可以分为几大类:
4.1 基于输出的蒸馏(Response-Based / Logit Distillation)
- 代表:Hinton 原始方法(2015)
- 做法:只对齐教师和学生的最终输出概率。
- 优点:实现简单,只需教师的输出即可。
- 缺点:只用了最后一层信息,中间层的丰富特征被浪费。
4.2 基于特征的蒸馏(Feature-Based Distillation)
- 代表:FitNets、TinyBERT、DistilBERT
- 做法:让学生的中间层特征(隐藏层激活值)去对齐教师的中间层。
- 典型损失:MSE(student_hidden, teacher_hidden)
- 适用:CNN、Transformer 等有明显层次结构的模型。
- 难点:教师和学生层数/维度不同时,需要映射函数(如线性投影)对齐。
4.3 基于关系的蒸馏(Relation-Based Distillation)
- 代表:Attention Transfer、Relational KD
- 做法:不对齐单个特征值,而是对齐样本之间的关系(如注意力矩阵、特征距离矩阵)。
- 直觉:教师”如何看待样本间的相似度”比”具体激活值”更本质。
4.4 自蒸馏(Self-Distillation)
- 做法:教师和学生同结构,甚至同一模型自我迭代。
- 场景:用大模型早期检查点指导后期训练,或让深层监督浅层。
- 意外收益:即使教师=学生,蒸馏也能带来性能提升(一种正则化)。
4.5 在线蒸馏(Online Distillation)
- 对比传统离线蒸馏:教师和学生同时训练,互相学习。
- 代表:Deep Mutual Learning、Codistillation
- 优势:无需预训练好的强教师,适合分布式训练。
4.6 无数据蒸馏(Data-Free Distillation)
- 场景:原始训练数据不可获取(隐私、版权等)。
- 做法:通过生成器合成样本,或用教师模型自身产生伪数据来蒸馏。
4.7 LLM 场景下的蒸馏
大语言模型蒸馏有其特殊性,形成了独立的方法体系:
- Sequence-Level KD:让学生模仿教师生成的完整序列(token-by-token 对齐输出分布)。
- On-Policy Distillation:学生自己生成,教师对学生生成的内容打分/纠正,避免”教师生成 vs 学生生成”的分布偏差(Exposure Bias)。
- 黑盒蒸馏(Black-box Distillation):只能访问教师 API(如 GPT-4),无法拿到 logits。此时通常用教师生成的**(问题, 回答)数据对**做监督微调(SFT),俗称”数据蒸馏”。
- 偏好蒸馏(Preference Distillation):用教师的偏好排序(A > B)训练学生,结合 DPO/RLHF。
- 推理链蒸馏(Chain-of-Thought Distillation):让教师生成带推理过程的答案,学生学习完整推理链,而不只是最终答案。
五、蒸馏实战流程
上面讲了原理和方法,这一节从工程角度说明如何真正跑通一次蒸馏。整体分为 6 个阶段:准备 → 设计 → 数据 → 训练 → 评估 → 迭代。
5.1 准备阶段:选定教师与学生
教师模型的选择
- 性能要求:在目标任务上已经充分训练、指标足够高。教师是学生的性能上限。
- 来源:
- 自训练的大模型 checkpoint
- 开源 SOTA(BERT-large、Llama-70B、Qwen-72B 等)
- 商用 API(GPT-4、Claude Opus)——只能做黑盒蒸馏
- ⚠️ 陷阱:Capacity Gap(能力鸿沟)
- 教师和学生规模差距过大时,学生”学不会”,性能反而下降。
- 经验法则:教师参数量控制在学生的 2~10 倍 为宜。
- 若必须蒸馏超大模型 → 小模型,考虑多阶段蒸馏:70B → 13B → 1B,逐级传递。
学生模型的设计
| 思路 | 说明 | 优缺点 |
|---|---|---|
| 同架构缩小 | 减层数、缩小隐藏维度、减注意力头 | 蒸馏稳定、层间对齐容易;但结构受限 |
| 同族异构 | 同类模型但结构不同(如 BERT → ALBERT) | 灵活,可利用参数共享等技巧 |
| 跨架构 | CNN → MLP、Transformer → RNN | 部署友好,但对齐难度大 |
决策建议:新手优先”同架构缩小 50~80%“,最容易跑通。
常见的学生尺寸对照:
- BERT-base (110M) → DistilBERT (66M),6 层 vs 12 层
- BERT-base (110M) → TinyBERT (14M),4 层 + 隐藏维度缩小
- Llama-13B → Llama-1.3B,层数和维度同时减
5.2 设计阶段:确定蒸馏方案
根据能从教师拿到什么,决定蒸馏方式:
| 能拿到的信号 | 推荐方法 | 复杂度 |
|---|---|---|
| 仅 API 文本输出 | 黑盒蒸馏 / 数据蒸馏(教师生成数据 → SFT) | ⭐ |
| 输出层 logits | Logit Distillation(Hinton 经典法) | ⭐⭐ |
| 中间层 hidden state | Feature Distillation(TinyBERT、DistilBERT) | ⭐⭐⭐ |
| 注意力矩阵 | Attention Distillation | ⭐⭐⭐ |
| 完整梯度 | 罕见,一般不用 | ⭐⭐⭐⭐ |
渐进式策略(推荐):
- 先跑通最简单的 Logit Distillation,建立 baseline。
- 效果不够 → 叠加 Feature/Attention Distillation。
- 仍不够 → 引入数据增强、多任务蒸馏。
5.3 数据阶段:准备训练数据
蒸馏对数据量的要求通常大于普通监督训练,因为学生需要”见到教师在各种输入下的反应”。
| 数据来源 | 说明 | 何时使用 |
|---|---|---|
| 原始训练集 | 教师训练时用的完整数据 | 最理想,效果最好 |
| 未标注数据 | 只有输入,教师现场打软标签 | 标注成本高、原始数据不可得 |
| 合成数据 | 教师主动生成 (input, output) | LLM 蒸馏主流方法 |
| 数据增强 | 对原数据做扰动/回译等 | 数据量不足时补充 |
关键工程技巧:教师 logits 离线缓存
# 一次性生成,永久复用
teacher.eval()
cached_logits = {}
for batch_id, batch in enumerate(dataloader):
with torch.no_grad():
cached_logits[batch_id] = teacher(batch).cpu()
torch.save(cached_logits, "teacher_logits.pt")
训练时直接读缓存,速度提升 50%+,也让教师和学生解耦部署。
5.4 训练阶段:核心实现
最基础版本(Logit Distillation)
import torch
import torch.nn.functional as F
def distillation_loss(student_logits, teacher_logits, labels,
T=4.0, alpha=0.3):
"""
T: 温度参数,越大分布越平滑
alpha: 硬损失权重,(1-alpha) 为软损失权重
"""
# 硬损失:学生 vs 真实标签
hard_loss = F.cross_entropy(student_logits, labels)
# 软损失:学生 vs 教师软分布(KL 散度)
soft_loss = F.kl_div(
F.log_softmax(student_logits / T, dim=-1),
F.softmax(teacher_logits / T, dim=-1),
reduction='batchmean'
) * (T * T) # 温度补偿,保证梯度尺度
return alpha * hard_loss + (1 - alpha) * soft_loss
# 训练循环
teacher.eval() # 教师冻结
student.train()
for x, y in dataloader:
with torch.no_grad():
t_logits = teacher(x)
s_logits = student(x)
loss = distillation_loss(s_logits, t_logits, y, T=4.0, alpha=0.3)
loss.backward()
optimizer.step()
optimizer.zero_grad()
进阶版本(多层次蒸馏,TinyBERT 风格)
def multi_level_loss(s_out, t_out,
s_hidden, t_hidden,
s_attn, t_attn,
s_emb, t_emb,
labels, T=4.0):
# 1. 输出层蒸馏
logit_loss = kl_distillation(s_out, t_out, T)
# 2. 中间层 hidden state 蒸馏(MSE)
# 学生层与教师层的映射:如学生第 i 层对齐教师第 2i 层
hidden_loss = sum(
F.mse_loss(proj(s_h), t_h)
for s_h, t_h, proj in zip(s_hidden, t_hidden, projections)
)
# 3. 注意力矩阵蒸馏
attn_loss = sum(
F.mse_loss(s_a, t_a)
for s_a, t_a in zip(s_attn, t_attn)
)
# 4. embedding 层对齐
emb_loss = F.mse_loss(proj_emb(s_emb), t_emb)
# 5. 真实标签
hard_loss = F.cross_entropy(s_out, labels)
return (0.3 * logit_loss + 0.3 * hidden_loss
+ 0.2 * attn_loss + 0.1 * emb_loss + 0.1 * hard_loss)
要点:
- 学生和教师层数不同时,需要指定”哪一层对齐哪一层”的映射(如学生 6 层对齐教师 12 层的第 2/4/6/8/10/12 层)。
- 学生和教师维度不同时,用线性投影层(
nn.Linear)把学生维度映射到教师维度。 - 各项权重通常靠实验调,或用两阶段训练:先只用 feature loss 训练一遍,再加 logit loss 微调。
关键超参数调节
| 参数 | 常用范围 | 调节方向 |
|---|---|---|
| 温度 T | 2 ~ 10 | 学生太”死板”(只学到硬标签)→ 调大 T;学生学偏 → 调小 T |
| α(硬/软权重) | 0.1 ~ 0.5 | 教师很强、真实标签有噪声 → α 小;教师一般 → α 大 |
| 学习率 | 教师训练的 1/2 ~ 1/5 | 学生有教师引导,不需要大步探索 |
| batch size | 越大越稳 | 显存不够 → 梯度累积 |
| 训练轮次 | 通常比从头训练少 30~50% | 蒸馏收敛更快 |
训练稳定性技巧
- 温度退火:训练初期 T 大(如 10),后期逐步降到 1。让学生先学结构、再学细节。
- 权重预热:前几个 epoch α=1(只用硬损失),之后逐步加大软损失比重。
- 教师输出正则化:对教师输出做 label smoothing,防止学生过度自信。
- 梯度裁剪:
clip_grad_norm=1.0,多层次蒸馏尤其重要。
工程优化技巧
- 半精度:教师前向用 fp16/bf16,学生保持 fp32,速度↑ 精度↓ 很小。
- 教师放独立 GPU:多卡场景下教师和学生分卡,避免抢显存。
- DataLoader
num_workers调优:教师前向是瓶颈时,多进程加速数据+教师推理。 - checkpoint 教师中间层:Feature Distillation 需要多层 hidden,一次前向全部返回。
5.5 评估阶段:多维度验证
蒸馏后必须严格评估,只看训练 loss 会误导。
5.5.1 性能指标(Quality)
- 主任务指标:准确率、F1、BLEU、ROUGE 等。
- 性能保留率 = 学生指标 / 教师指标,目标通常 90%+。
- 对比 baseline:学生从头训练 vs 学生蒸馏训练,差值才是”蒸馏的净收益”。
5.5.2 效率指标(Efficiency)
| 指标 | 说明 | 测量方法 |
|---|---|---|
| 推理延迟 | ms/sample | 单样本 100 次取平均 |
| 模型大小 | MB | 磁盘上 checkpoint 大小 |
| 吞吐量 | QPS / tokens per sec | 满负载压测 |
| 显存占用 | MB | nvidia-smi 或 torch.cuda.max_memory_allocated() |
| 能耗 | J/sample | 生产部署才关心 |
5.5.3 分布对齐度(Fidelity,可选)
- KL 散度:学生输出 vs 教师输出的分布距离,越小越接近。
- Agreement Rate:Top-1 预测一致率。目标 ≥ 95%。
- Top-K 重叠:Top-5 预测集合的 Jaccard 相似度。
5.5.4 泛化性与鲁棒性
- OOD 测试:教师训练分布外的数据,看学生是否过拟合教师的特定行为。
- 对抗样本:加噪、扰动,测鲁棒性。
- 长尾类别:小样本类别上学生是否退化。
- LLM 场景:hallucination 检测、指令遵循能力、安全性。
5.6 迭代阶段:诊断与优化
蒸馏很少一次成功,需要根据评估结果迭代。常见问题诊断表:
| 现象 | 可能原因 | 对策 |
|---|---|---|
| 学生性能远不如教师 | Capacity gap 太大 | 换更小的教师;加中间层蒸馏;多阶段蒸馏 |
| 学生比不蒸馏时还差 | α 或 T 设置不当 | 降 α(更信教师);调大 T |
| 训练不稳定,loss 震荡 | 学习率过大、T 过高 | 降 lr;T 退火 |
| 训练集好、验证集差 | 过拟合教师特定行为 | 数据增强;混合真实标签;早停 |
| 学生输出过于自信 | 温度太低,暗知识没学到 | 增大 T;用 label smoothing |
| 收敛非常慢 | 损失权重失衡 | 逐项检查各 loss 数量级,重新调权重 |
| LLM 蒸馏出现幻觉 | 教师生成数据含幻觉 | 数据过滤 + 拒绝采样 + 人工审核 |
| 学生在教师错的样本上也错 | 学到了教师的偏见 | 引入独立评估集;结合真实标签惩罚 |
5.7 LLM 蒸馏的特殊流程
大语言模型蒸馏(如 Llama-70B → Llama-7B)与传统蒸馏差异较大,实际流程通常是:
步骤 1:Prompt 池构造
- 收集 10 万~100 万条多样化 prompt,覆盖目标任务分布。
- 来源:真实用户日志、开源数据集(Alpaca、ShareGPT)、人工设计。
步骤 2:教师批量生成
- 用教师对每条 prompt 生成 response,得到 (prompt, response) 数据集。
- 可采样多个 response(temperature > 0),后续做拒绝采样。
步骤 3:数据清洗
- 过滤:过长/过短、重复、明显错误、涉及敏感内容。
- 拒绝采样:用打分模型或规则筛掉低质量样本。
- 常见比例:生成 100 万 → 保留 20~30 万。
步骤 4:学生 SFT(监督微调)
- 在清洗后的数据上做 next-token prediction。
- 通常 2~3 个 epoch。
- 关键:mask 掉 prompt 部分,只在 response tokens 上计算 loss。
步骤 5(可选):偏好蒸馏
- 教师对 (A, B) 两个回答做偏好排序。
- 学生用 DPO/RLHF 学习偏好。
步骤 6(可选):推理链蒸馏(CoT Distillation)
- 让教师生成带推理过程的回答(“Let’s think step by step…”)。
- 学生学完整推理链,而不是最终答案。
LLM 蒸馏的特点:
- 拿不到教师 logits(尤其是 API 教师)→ 本质是”教师数据 + SFT”。
- Hinton 那套温度 + KL 用得少。
- 数据质量 >> 模型架构,80% 的精力花在数据构造和清洗上。
5.8 快速上手推荐路径
如果你是第一次做蒸馏:
- 跑通别人的例子:HuggingFace
transformers里的 DistilBERT 训练脚本 是最好的入门。 - 换成自己的模型:教师/学生替换为你要用的模型,先做纯 logit 蒸馏。
- 建 baseline:对比”学生从头训练”、“学生 fine-tune”、“学生蒸馏”三种效果,量化蒸馏的净收益。
- 叠加中间层蒸馏:logit 蒸馏效果不够时加。
- 调超参:T、α、lr 是三大调节点。
避免的常见错误:
- ❌ 教师没冻结,误更新参数
- ❌ 忘记
T*T温度补偿,导致梯度过小 - ❌ 只看训练 loss,不做多维度评估
- ❌ 学生和教师用不同的 tokenizer/预处理,导致对齐失败
- ❌ LLM 蒸馏时 loss 算在 prompt tokens 上
六、总结
| 维度 | 关键点 |
|---|---|
| 本质 | 把大模型的暗知识迁移到小模型 |
| 核心工具 | 软标签 + 温度参数 + KL 散度 |
| 迁移对象 | 输出分布 / 中间特征 / 关系结构 / 推理过程 |
| 典型收益 | 参数量减少 5~10 倍,性能保留 90%+ |
| 主要应用 | 模型部署压缩、边缘计算、LLM 小型化、跨模型能力迁移 |
模型蒸馏并不是”简单地让小模型模仿大模型”,而是利用大模型学到的丰富知识结构,作为小模型训练的高质量监督信号。它已成为现代 AI 模型工程的基础技术之一,尤其在 LLM 时代,是 “把大模型能力送到每一台设备上” 的关键路径。
参考文献
- Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531.
- Sanh, V., et al. (2019). DistilBERT, a distilled version of BERT. arXiv:1910.01108.
- Jiao, X., et al. (2019). TinyBERT: Distilling BERT for Natural Language Understanding. arXiv:1909.10351.
- Gou, J., et al. (2021). Knowledge Distillation: A Survey. IJCV.