知识蒸馏实践指南

知识蒸馏（Knowledge Distillation）是一种模型压缩技术：用一个大而强的”教师”模型指导训练一个小而快的”学生”模型，让学生在参数量和推理成本远小于教师的前提下，尽可能逼近教师的能力。本页从核心直觉出发，完整讲解软标签与温度、蒸馏损失、中间层与特征蒸馏，再到 DistilBERT / TinyBERT / MiniLM 的工程实践。前置阅读可参考模型评估与混合精度训练；如果你更关心小模型选型与部署，可直接跳到小模型选型，蒸馏后的模型常常还需要做微调才能落到具体任务。

核心直觉

想象一位经验丰富的老教师（大模型）带一名小学生（小模型）。考试时，老师不只是告诉学生”答案是 C”（硬标签），而是把完整的心里盘算都说出来：“我觉得 A 有 10% 可能、B 有 5%、C 有 80%、D 有 5%“（软标签）。这份”我觉得”的概率分布里藏着比正确答案本身信息量大得多的东西——它透露了老师认为哪些错误选项”有点像”、哪些”完全不着边际”。Hinton 把这些隐藏在概率分布里的细微判断称为暗知识（dark knowledge）。

硬标签= 只有”对/错”的最终裁决，信息量小，学生只知道”C 是对的”，却不知道 A 和 B 哪个更接近正确。
软标签= 老师对每个选项的置信度分布，信息量大，学生不仅学到”C 是对的”，还学到”A 比 D 更容易和 C 混淆”——这正是泛化能力的来源。
温度= 控制老师”说多详细”的旋钮。温度低时老师嘴很紧，只报”就是 C”；温度高时老师把每个选项的可能性都摊开来讲，暗知识就暴露出来了。
中间层蒸馏= 不仅让学生模仿老师的”最终答案”，还让它偷偷看老师的”草稿纸”（中间层特征），学到老师推理的过程而不只是结论。
特征蒸馏= 进一步让学生模仿老师的”注意力”——老师在做阅读理解时盯着哪些词，学生也跟着盯，从而学到更丰富的表征。

一句话：蒸馏 = 让小模型不仅学会”答案是什么”，还学会”大模型怎么思考”。

为什么软标签信息量更大？——信息论视角

从信息论（Information Theory）的角度看，硬标签是一个 one-hot 向量（只有正确类为 1，其余为 0），它的信息熵（Shannon Entropy，衡量分布的不确定度）为 0，信息量极低。而软标签是一个非退化的概率分布，信息熵大于 0，携带了 Teacher 对所有类别的相对判断。

举个具体的数字感受：假设有 1000 个类别（类似 ImageNet），硬标签告诉你”答案是第 523 类”，只用了 $\log_2 1000 \approx 10$ bit 的信息。但 Teacher 的软标签可能在第 523 类给出 0.7、在第 100 类给出 0.15、在第 789 类给出 0.08……这些”次优选项”的排序本身就是极其宝贵的学习信号。研究表明，一个样本的软标签平均携带的信息量是硬标签的数倍——这就是为什么蒸馏能让 Student 在同等数据量下学得更快、泛化更好。

信息熵（Entropy）： $H(p) = -\sum_i p_i \log p_i$ ，衡量概率分布 $p$ 的不确定度。分布越均匀，熵越大；越尖锐（one-hot），熵越小。软标签的熵远大于硬标签（熵为 0）。

原理详解

Teacher-Student 框架

知识蒸馏的整体流程分两步：

先训练 Teacher：用一个容量大、参数多的模型（如 12 层 BERT、GPT-4）在训练数据上训练到收敛，得到一个性能强大但笨重的教师模型。Teacher 一旦训练好就固定不动，只在后续步骤中提供”指导信号”。
再训练 Student：用一个容量小、参数少的模型（如 6 层 Transformer）作为学生。Student 的训练目标不是从零学习真实标签，而是同时模仿 Teacher 的输出行为——它要让自己的输出分布尽量接近 Teacher 的输出分布。

直观上，Student 的优化空间被 Teacher 的输出”引导”了：真实标签只告诉它正确答案，而 Teacher 的软标签还告诉它错误答案之间的相对关系，相当于给 Student 提供了一个更平滑、信息更丰富的学习信号。Student 因此能用更少的数据、更少的参数达到接近 Teacher 的效果。

软标签与温度（Temperature）

标准的 softmax 把 logits（模型最后一层的原始输出）转成概率：

p_i = exp(z_i) / sum_j exp(z_j)

在标准 softmax（温度 T=1）下，训练好的 Teacher 对正确类的置信度往往接近 1，对其他类的概率接近 0——输出非常”尖锐”，暗知识几乎被掩盖。为了让 Teacher 把暗知识暴露出来，我们在 softmax 前把 logits 除以一个大于 1 的温度 T：

soft_label_i = softmax(z_i / T) = exp(z_i / T) / sum_j exp(z_j / T)

温度 T 的作用：

T=1：标准 softmax，输出尖锐，正确类概率接近 1，其他类接近 0。
T 越大：各类的概率差异被缩小，输出越”柔和”——原来 0.99 和 0.001 的差距，在 T=4 时可能变成 0.6 和 0.15，原本被压成 0 的暗知识被放大了。
T 趋于无穷：所有类概率趋于均匀分布（1/N），信息量反而降低。

蒸馏时通常选用高温（如 T=4 或 T=5）来生成 Teacher 和 Student 的软标签。Student 在训练时也用同样的温度 T 计算自己的软输出，再与 Teacher 的软标签做 KL 散度。推理时把温度恢复为 1 即可正常使用。

softmax：把一组任意实数（logits）归一化为概率分布的函数，保证所有输出非负且和为 1。温度 T 是 softmax 的”柔和度”参数——T 越大，输出越均匀；T 越小，输出越尖锐。

温度对梯度的影响（为什么要乘 $T^2$ ）

KL 散度损失对 Student logits $z_i$ 求梯度时，由于 soft_label 中 $z_i$ 被 $T$ 缩放，梯度会出现 $1/T$ 的因子。具体来说，softmax 的梯度天然带有 $\frac{\partial p_i}{\partial z_j} = \frac{1}{T} p_i (\delta_{ij} - p_j)$ 的形式（ $\delta_{ij}$ 是 Kronecker delta），因此整个 KL 损失的梯度幅度近似为：

$\frac{\partial \mathcal{L}_{KD}}{\partial z_i} \approx \frac{1}{T} (p_i^{student} - p_i^{teacher})$

如果不做补偿，温度越高，梯度越小，Student 学得越慢。Hinton 在原论文中指出：将 KL 损失项乘以 $T^2$ 可以补偿这一缩放，使得无论 T 取多少，蒸馏梯度与硬标签交叉熵的梯度量级保持一致，alpha 权重无需随温度调整。这就是代码中 * (T * T) 的由来。

蒸馏损失函数

Student 训练时的总损失由两部分加权组成：

$\mathcal{L}_{total} = \alpha \cdot \mathcal{L}_{KD}(T) + (1 - \alpha) \cdot \mathcal{L}_{CE}(T=1)$

展开各项：

$\mathcal{L}_{KD} = T^2 \cdot \text{KL}(P_T^{(\tau)} \| P_S^{(\tau)}) = T^2 \sum_i P_T^{(\tau)}(i) \log \frac{P_T^{(\tau)}(i)}{P_S^{(\tau)}(i)}$

$\mathcal{L}_{CE} = -\sum_i y_i \log P_S^{(1)}(i)$

其中 $P_T^{(\tau)}(i) = \text{softmax}(z_i^T / T)$ 是 Teacher 在温度 T 下的软标签， $P_S^{(\tau)}$ 是 Student 的对应输出， $y$ 是 one-hot 真实标签， $P_S^{(1)}$ 是 Student 在标准温度下的输出。

其中：

KL 散度项：衡量 Student 软输出与 Teacher 软标签的差异。具体地， $\text{KL}(P_T \| P_S) = \sum_i P_T(i) \log \frac{P_T(i)}{P_S(i)}$ ，用温度 T 计算双方的软概率。这一项让 Student 模仿 Teacher 的判断分布，吸收暗知识。
交叉熵项（Cross-Entropy，即用对数概率加权的负对数似然）：Student 在标准温度（T=1）下的输出与真实硬标签 y 的交叉熵，即 $-\sum_i y_i \log p_i$ 。这一项保证 Student 不偏离真实任务目标，毕竟 Teacher 本身也不是完美的。
alpha（ $\alpha$ ）：权重系数，取值在 0 到 1 之间，控制”模仿 Teacher”和”学习真实标签”的比例。常用值在 0.5 到 0.9 之间，偏向 Teacher 的暗知识。

交叉熵（Cross-Entropy）：衡量两个分布差异的指标，在分类任务中即 $-\sum_i y_i \log p_i$ 。当 $y$ 是 one-hot 时，它退化为 $-\log p_{correct}$ ——模型对正确类别的预测概率越高（越有信心），损失越小。

KL 散度（KL Divergence）： $\text{KL}(P \| Q) = \sum_i P(i) \log \frac{P(i)}{Q(i)}$ ，度量分布 Q 偏离参考分布 P 的程度。它非负且当 P=Q 时为 0，但不具对称性（ $\text{KL}(P\|Q) \neq \text{KL}(Q\|P)$ ）。交叉熵与 KL 的关系为： $H(P, Q) = H(P) + \text{KL}(P\|Q)$ ，当 P 固定时最小化交叉熵等价于最小化 KL。

KL 散度 vs 反向 KL 散度

标准蒸馏使用的是前向 KL： $\text{KL}(P_T \| P_S)$ 。前向 KL 是”均值搜索（mean-seeking）“的，它鼓励 Student 的分布覆盖 Teacher 所有概率质量较高的区域，即使 Teacher 分布有多个峰，Student 也倾向于平均覆盖。

而在生成式 LLM 蒸馏中（如 MiniLLM），研究者发现反向 KL $\text{KL}(P_S \| P_T)$ 更有效：反向 KL 是”模式搜索（mode-seeking）“的，Student 倾向于集中在 Teacher 分布的主峰上，避免生成低质量的尾巴样本。这对开放式文本生成尤为重要，因为覆盖 Teacher 所有可能输出（包括低质量的长尾）反而有害。

实践中的一个技巧：KL 散度项计算时两个分布都用了温度 T，梯度幅度会随 T 变化（经验上需要乘以 T 的平方来补偿），这是 Hinton 原论文给出的修正，确保温度变化时梯度尺度稳定。

中间层蒸馏

只蒸馏最终输出有一个问题：Student 只能看到 Teacher 的”结论”，看不到”推理过程”。中间层蒸馏（Intermediate Layer Distillation）让 Student 的中间层特征去对齐 Teacher 的中间层特征。

做法是：Teacher 有 N 层，Student 有 M 层（M 通常小于 N），先建立一个层映射关系（比如 Teacher 的第 2、6、10 层对应 Student 的第 1、3、5 层）。对每一对要对齐的层，由于 Teacher 和 Student 的隐藏维度可能不同，需要在 Student 侧加一个线性投影层（adapter / projection），把 Student 的特征维度变换到与 Teacher 一致，再用 MSE 损失让两者逼近：

loss_hidden = MSE(project(student_hidden_i), teacher_hidden_j)

加上最终输出的蒸馏损失，总损失变成：输出蒸馏损失 + 中间层蒸馏损失。TinyBERT 的核心创新就是把中间层蒸馏系统化，证明了对齐中间表征能显著提升小模型的效果。

特征蒸馏

特征蒸馏比中间层蒸馏更精细，关注的是 Teacher 内部的具体信息流，主要包括：

注意力矩阵蒸馏（Attention Distillation）：让 Student 的自注意力分布去模仿 Teacher 的自注意力分布。具体做法是取出 Teacher 每层的注意力权重矩阵（query 和 key 的点积 softmax 后的结果），让 Student 对应层的注意力矩阵去逼近它（用 MSE 或 KL 散度）。这让 Student 学到 Teacher “看哪些词、忽略哪些词”的模式。
隐状态蒸馏（Hidden State Distillation）：即上面提到的中间层特征对齐。
Embedding 层蒸馏：对齐 Teacher 和 Student 的词嵌入输出。

研究表明，注意力蒸馏对迁移到下游任务（如文本分类、问答）尤其有效，因为注意力模式携带了大量关于语言理解的结构化信息。TinyBERT 综合采用了 embedding 蒸馏、注意力蒸馏、隐状态蒸馏和输出蒸馏，是特征蒸馏的代表方案。

DistilBERT

DistilBERT 是 HuggingFace 提出的轻量化 BERT，是工业界最知名的蒸馏实践之一：

结构：6 层 Transformer（Teacher BERT 是 12 层），保留与 BERT 相同的隐藏维度（768）和注意力头数。
参数量：约 6600 万，比 BERT-base（1.1 亿）减少约 40%。
训练方法：主要采用输出层蒸馏（软标签 + 硬标签的混合损失），加上一个”余弦距离”项让 Student 的隐状态对齐 Teacher。训练数据与 BERT 相同的大规模语料。
效果：推理速度提升约 60%，在 GLUE 基准上保留 BERT 约 97% 的能力。
亮点：训练成本低（8 张 V100 上约 3.5 天），部署友好，是 HuggingFace pipeline 的默认小模型之一。

TinyBERT

TinyBERT（华为提出）的两阶段蒸馏是它的核心特色：

预训练阶段蒸馏：在大规模无标注语料上，让 Student 模仿预训练好的 Teacher BERT。这一阶段包含 Transformer 各层的 embedding 蒸馏、注意力蒸馏、隐状态蒸馏，以及最终的预测层蒸馏。
微调阶段蒸馏：在具体下游任务（如 SST-2 情感分类、SQuAD 问答）上，先让 Teacher 在该任务上微调好，再用同样的多层蒸馏方法训练 Student。同时配合数据增强（用 Teacher 或预训练语言模型对训练样本做同义改写）扩充数据，提升 Student 的泛化能力。
结构：4 层 Transformer，参数量约为 BERT 的 1/7，在多个下游任务上明显优于同等规模的基线。
亮点：证明了”两阶段都蒸馏”比”只蒸馏微调阶段”效果好得多，中间层对齐是关键。

MiniLM

MiniLM（微软提出）的思路是深度自注意力蒸馏（Deep Self-Attention Distillation），特点包括：

不强制 Student 模仿 Teacher 的每一层，而是重点蒸馏 Teacher 最后一层的自注意力关系矩阵（self-attention relation matrix，即 value 向量之间的点积），这比单纯蒸馏注意力权重更能捕捉 token 间的关系。
Student 的层数和隐藏维度都可以比 Teacher 小，甚至可以做跨语言、跨模态的蒸馏（比如用英文 Teacher 蒸馏多语言 Student），因为它蒸馏的是与具体语言无关的注意力关系。
MiniLM 系列在多语言理解（XTune、XNLI）等任务上表现突出，是微软多语言模型训练的核心技术之一。
适用范围广：由于不依赖严格的层对齐，Student 的结构设计更灵活。

原理图解

Teacher-Student 蒸馏总框架

中间层与特征蒸馏对齐

动手实践

下面用 PyTorch 展示蒸馏的核心损失计算（软标签 KL 散度 + 硬标签交叉熵的混合）：

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels,
                      T=4.0, alpha=0.7):
    # 学生软输出与教师软标签的 KL 散度（都除以温度 T）
    soft_student = F.log_softmax(student_logits / T, dim=-1)
    soft_teacher = F.softmax(teacher_logits / T, dim=-1)
    # KL 散度，乘以 T^2 补偿温度缩放对梯度的影响
    loss_kd = F.kl_div(soft_student, soft_teacher,
                       reduction="batchmean") * (T * T)
    # 学生标准输出与真实硬标签的交叉熵
    loss_ce = F.cross_entropy(student_logits, labels)
    # 加权总损失：alpha 控制模仿教师与学习真实标签的比例
    return alpha * loss_kd + (1.0 - alpha) * loss_ce

# teacher 模型设为 eval 并冻结，不计算梯度
teacher.eval()
with torch.no_grad():
    teacher_logits = teacher(input_ids)

student_logits = student(input_ids)           # 学生前向传播
loss = distillation_loss(student_logits,
                         teacher_logits, labels, T=4.0, alpha=0.7)
loss.backward()                                # 反向传播，只更新 student
optimizer.step()

实践要点

温度选择：T 通常取 2 到 10 之间，T=4 是 Hinton 论文的经验默认值。T 太小则软标签仍然太尖锐、暗知识暴露不足；T 太大则各类概率趋于均匀、信号被稀释。建议在验证集上对小范围（如 1、2、4、8）做网格搜索。
alpha 权重调参：alpha 决定”学 Teacher”与”学真实标签”的比例。数据充足、Teacher 很强时可以偏向 Teacher（alpha 接近 0.7 到 0.9）；数据噪声大或 Teacher 本身不够准时应降低 alpha。实测中 alpha 的敏感度往往比温度低，可先固定 alpha 再调温度。
数据增强配合蒸馏：TinyBERT 的经验表明，蒸馏时配合数据增强（如回译、同义词替换、用 Teacher 生成更多训练样本）能显著提升 Student 泛化能力，尤其在下游任务数据有限的场景。
中间层蒸馏优先用于迁移任务：如果你的 Student 要迁移到多个下游任务，中间层 + 注意力蒸馏（TinyBERT 式）通常优于只蒸馏最终输出（DistilBERT 式）。只做单一任务且追求最快训练，输出蒸馏即可。
Student 结构设计：Student 不是越小越好，层数减半（如 12 层变 6 层）但保留隐藏维度，是性价比最高的方案；过度压缩维度会导致表征能力不足，蒸馏也救不回来。参考 DistilBERT 的”减层不减宽”策略。
学习率与训练步数：蒸馏时 Student 的学习率通常比从头训练略小，训练步数可以更多——Teacher 提供了更平滑的梯度信号，Student 有更大空间慢慢逼近。配合 warmup 和 cosine 衰减效果更好。

典型应用

DistilBERT 部署：HuggingFace 的 DistilBERT 是工业界最广泛使用的蒸馏 BERT 之一，常用于搜索排序、文本分类、命名实体识别等在线服务，在保持接近 BERT 效果的同时显著降低延迟和服务器成本。
移动端与边缘设备小模型：手机、智能音箱、车载系统上部署 NLP 模型时，参数量和推理延迟是硬约束。通过蒸馏得到的 TinyBERT、MobileBERT 等可以在端侧完成推理，无需上云，兼顾隐私和延迟。
LLM 蒸馏（Alpaca / Vicuna / WizardLM）：Stanford Alpaca 用 GPT-4（Teacher）生成 5.2 万条指令数据，蒸馏出 7B 的 Alpaca 模型；Vicuna、WizardLM 等也类似，用闭源大模型的输出训练开源小模型，让 7B 到 13B 的模型获得接近 GPT-4 在某些任务上的表现。这是当前开源 LLM 生态最主流的”能力下放”路径。
多语言与跨语言模型：微软 MiniLM 用英文 Teacher 蒸馏多语言 Student，在 XNLI 等跨语言基准上取得领先效果，证明了蒸馏可以迁移与语言无关的语义表征。
语音与多模态蒸馏：语音识别（如 DistilWhisper）、图文检索等场景也大量采用蒸馏，把大模型压缩到可以在实时管道中运行的体积。
推荐系统与广告：大规模 CTR 模型常用蒸馏把复杂模型（如深度交叉网络）的知识传递给轻量部署模型，兼顾离线训练效果和在线推理性能。

典型类库与工具

类库	语言	说明
Transformers	Python	HuggingFace 库内置 DistilBERT、TinyBERT、MiniLM 等蒸馏模型，开箱即用
textbrewer	Python	哈工大讯飞联合实验室的通用知识蒸馏工具包，支持中间层与注意力蒸馏
(adapter-less)	Python	微软开源的 MiniLM 训练代码与预训练权重，支持多语言蒸馏
nn.distill (Torch)	Python	PyTorch 生态中可自定义蒸馏损失的自实现方案（本文实践代码即此思路）
DeepSpeed	Python	微软训练框架，内置模型压缩与蒸馏能力，适合大模型场景
TensorRT	Python/C++	NVIDIA 推理加速库，常与蒸馏后的小模型配合做生产部署

术语表

术语	英文	解释
知识蒸馏	Knowledge Distillation	用大模型（教师）的输出指导训练小模型（学生）的模型压缩技术
教师模型	Teacher Model	预先训练好的大而强的模型，在蒸馏中提供指导信号，参数固定不动
学生模型	Student Model	被训练的小而快的模型，目标是模仿教师行为
软标签	Soft Label	教师在高温 softmax 下输出的概率分布，包含各类的相对置信度
硬标签	Hard Label	数据集真实标签，通常是 one-hot 形式，只标明正确类别
暗知识	Dark Knowledge	软标签中隐藏的、教师对错误类别的细微判断信息
温度	Temperature	softmax 前对 logits 的缩放因子，温度越高输出越柔和
KL 散度	KL Divergence	衡量两个概率分布差异的指标，蒸馏中用于让学生逼近教师分布
中间层蒸馏	Intermediate Layer Distillation	对齐教师与学生中间层特征的蒸馏方法，TinyBERT 的核心
注意力蒸馏	Attention Distillation	让学生模仿教师自注意力分布的特征蒸馏方法
特征蒸馏	Feature Distillation	泛指对齐教师内部表征（注意力、隐状态、嵌入）的蒸馏方法
两阶段蒸馏	Two-stage Distillation	预训练阶段和微调阶段都做蒸馏，TinyBERT 的训练策略