模型知识蒸馏#

TL;DR#

1. 背景与问题#

只用硬标签训练时，每个样本只指定正确类别。教师模型还会给出其他类别的相对概率；知识蒸馏希望学生利用这份软监督，学习教师的预测行为。本文聚焦分类分布蒸馏，学生可以更小，也可以采用不同结构。[1]

2. 目标与定义#

设教师与学生 logits 为 $z^T,z^S\in\mathbb R^C$，温度 $\tau>0$。定义 $p_i^T=\operatorname{softmax}(z^T/\tau)_i$、$q_i^S=\operatorname{softmax}(z^S/\tau)_i$，教师分布不参与反向。硬标签为 $y$，混合系数 $\alpha\in[0,1]$：

$$\mathcal L_{\text{KD}}=\alpha\tau^2\sum_{i=1}^C p_i^T\log\frac{p_i^T}{q_i^S}+(1-\alpha)\left[-\log\operatorname{softmax}(z^S)_y\right].$$

软分布项在温度 $\tau$ 下计算，硬标签项仍使用温度 1。固定教师时，KL 与软标签交叉熵对学生有相同梯度，因为教师熵不依赖学生参数。

3. 方法与推导#

3.1 为什么乘温度平方#

先对未乘 $\tau^2$ 的 KL 求学生 logit 梯度：

$$\frac{\partial D_{\text{KL}}(p^T\|q^S)}{\partial z_i^S}=\frac{q_i^S-p_i^T}{\tau}.$$

当温度相对 logits 的差异足够大，令 $\bar z=C^{-1}\sum_i z_i$，softmax 的一阶展开为 $q_i\approx C^{-1}+(z_i-\bar z)/(C\tau)$。代回可得：

$$\frac{q_i^S-p_i^T}{\tau}\approx\frac{(z_i^S-\bar z^S)-(z_i^T-\bar z^T)}{C\tau^2}.$$

这解释高温近似下梯度主要按 $1/\tau^2$ 缩小，乘 $\tau^2$ 用来补偿。有限温度下仍要调 $\alpha$；温度越大不表示教师包含的知识越多。[1]

4. 计算过程#

教师 logits 为 $(5,2,0.5)$，类别依次为猫、狗、车。按 softmax 计算：

狗与车的概率比从 $\exp(1.5)\approx4.482$ 变为 $\exp(0.3)\approx1.350$。高温让原来很小的概率更容易参与学习，同时也压平差异；它不会自动把类别解释成有保证的语义相似关系。

若学生在温度 5 时输出均匀分布，则软 KL 约为 $0.07250$，乘 $\tau^2=25$ 后约为 $1.8124$。其对猫 logit 的梯度为 $\tau(q^S_{\text{猫}}-p^T_{\text{猫}})\approx-0.8904$，梯度下降会提高该 logit。以上为公式演算，未包含硬标签项和 $\alpha$ 权重。

5. 比较与边界#

5.1 不同监督信号怎样比较#

DistilBERT 与 TinyBERT 是结合多种监督的模型案例，[2]指标只能连同原任务、模型与评测报告引用。生成数据路线还受教师错误、数据覆盖和筛选方法限制，不能仅由教师规模预测学生收益。

5.2 成本与效果怎样验证#

先固定硬标签基线，再分别改变温度、混合系数与教师信号，记录相同预算下的验证指标。检查教师是否冻结、KL 方向与 reduction 是否正确，以及学生和教师是否在同一输入、位置与词表上比较。

显存不足时按训练诊断区分激活、教师前向与优化器状态；重算和分片见系统专题。教师与学生同时常驻显存只是在线计算的一种实现选择。

参考资料#

1. Distilling the Knowledge in a Neural Network（2015）. §2：软目标、温度与梯度尺度。
2. DistilBERT / TinyBERT（2019）. DistilBERT、TinyBERT：多种监督的具体模型案例。

延伸阅读#

• 训练目标导览：怎样把任务目标写成可优化的损失？