知识蒸馏

TIP

AI-generated content; reviewed by a human.

知识蒸馏（Knowledge Distillation, KD）的基本思想是：除了真实标签，还让学生模型拟合教师模型的输出分布。

one-hot 标签只告诉学生“这个样本属于第 3 类”，教师的概率分布还可能提供额外信息，例如“它更容易被误判为第 7 类，而不像第 1 类”。这类额外信息通常称为暗知识（dark knowledge）。它可能编码教师学到的类别混淆关系和泛化规律，因此一个软目标样本能够提供比单个标签更丰富的监督信号。不过，这些信息的质量仍取决于教师模型本身。[1]

问题在于，当教师模型对正确类别的预测非常确信时，其他类别的概率可能都很小。这些小概率之间的相对比例可能包含有用信息，却不一定能在原始概率尺度下对训练产生足够大的影响。

温度（Temperature）可以调节输出分布的平滑程度。增大温度会抬高非目标类别的相对概率，使它们之间的差异更容易进入训练信号。但温度并不是越高越好：当温度趋于无穷大时，分布会趋近均匀分布。

本文从温度 Softmax 出发，推导经典知识蒸馏损失中的 $T^2$ 补偿项，并讨论相关方法在大语言模型蒸馏中的延伸。

为统一符号，设分类任务有 $N$ 个类别，教师模型的 logits 为 $v=(v_1,\dots ,v_N)$，学生模型的 logits 为 $z=(z_1,\dots ,z_N)$。对应的教师概率记为 $p$，学生概率记为 $q$。

1. 温度

带温度的 Softmax 定义为

$$q_i(T)=\frac{\exp (z_i/T)}{{\sum }_j\exp (z_j/T)}.\text{\hspace{1em}(1)}$$

教师模型的软目标采用同样的定义：

$$p_i(T)=\frac{\exp (v_i/T)}{{\sum }_j\exp (v_j/T)}.$$

其中 $T>0$ 为温度参数。

当 $T=1$ 时，就是普通 Softmax。当 $T>1$ 时，分布会变得更加平滑；随着 $T\to \infty$，概率趋向均匀分布。

对于任意正温度，Softmax 都保留 logits 的大小顺序，因此

$$\arg \underset{i}{\max }{q}_i(T)=\arg \underset{i}{\max }{z}_i.$$

温度改变的是概率分布的尺度，而不是预测类别的排序。对于蒸馏而言，适当提高温度可以让非目标类别的概率不再过于接近零，使它们之间的相对关系更容易影响学生模型的训练。

2. 损失函数

令真实标签为 $y$，通常用 one-hot 向量表示。经典知识蒸馏损失可写为

$$L=\alpha T^2\mathrm{KL}\left(p(T)\Vert q(T)\right)+(1-\alpha )\mathrm{CE}\left(y,q(1)\right).\text{\hspace{1em}(2)}$$

其中：

• 第一项是软目标蒸馏损失，用温度 $T$ 软化教师和学生的概率分布。
• 第二项是真实标签交叉熵，使用 $T=1$ 的学生预测。
• $\alpha \in [0,1]$ 用于平衡两项损失。
• $T^2$ 是经典蒸馏方法中用于补偿高温梯度衰减的因子。[1]

KL 散度定义为

$$\mathrm{KL}(p\Vert q)=\sum _i{p}_i\log \frac{p_i}{q_i}=\sum _i{p}_i\log p_i-\sum _i{p}_i\log q_i.$$

由于教师分布 $p$ 固定，第一项 ${\sum }_i{p}_i\log p_i$ 与学生 logits 无关。因此，对学生 logits 求导时，KL 散度与交叉熵

$$-\sum _i{p}_i\log q_i$$

具有相同的梯度。

注意，$\alpha$ 的合适取值取决于任务、模型和数据，不能认为所有蒸馏任务都存在统一的最佳范围。即使显式损失系数相同，两项损失对参数更新的实际贡献也不一定相同。

3. 为什么要乘 $T^2$

3.1. 梯度

记软目标蒸馏损失为

$$L_{\mathrm{KD}}=\mathrm{KL}(p\Vert q),$$

其中 $p=p(T)$、$q=q(T)$。教师分布不随学生 logits $z$ 变化。

因为 KL 散度与交叉熵对学生 logits 的导数相同，所以

$$\begin{array}{rl}\frac{\partial L_{\mathrm{KD}}}{\partial z_i}& =-\sum _k{p}_k\frac{\partial \log q_k}{\partial z_i}\\ & =-\sum _k\frac{p_k}{q_k}\frac{\partial q_k}{\partial z_i}.\end{array}$$

带温度的 Softmax 导数为

$$\frac{\partial q_k}{\partial z_i}=\frac{1}{T}{q}_k({\delta }_{ki}-q_i),$$

其中 ${\delta }_{ki}$ 是 Kronecker delta：

$${\delta }_{ki}=\{\begin{array}{ll}1,& k=i,\\ 0,& k\ne i.\end{array}$$

将 Softmax 导数代入梯度表达式：

$$\begin{array}{rl}\frac{\partial L_{\mathrm{KD}}}{\partial z_i}& =-\frac{1}{T}\sum _k{p}_k({\delta }_{ki}-q_i)\\ & =-\frac{1}{T}\left(p_i-q_i\sum _k{p}_k\right).\end{array}$$

由于教师概率满足 ${\sum }_k{p}_k=1$，得到

$$\boxed{\frac{\partial L_{\mathrm{KD}}}{\partial z_i}=\frac{1}{T}(q_i-p_i)}\text{\hspace{1em}(3)}$$

梯度中已经出现了一个 $1/T$。但这还不是全部：高温下，教师和学生的概率差 $q_i-p_i$ 也会按 $1/T$ 衰减。

3.2. 高温下的概率分布

为了看清概率差如何随温度变化，下面推导 $T\to \infty$ 时的近似表达式。

首先，Softmax 对 logits 的整体平移不变。对任意常数 $c$，

$$\begin{array}{rl}{q}_i(z+c1)& =\frac{e^{(z_i+c)/T}}{{\sum }_j{e}^{(z_j+c)/T}}\\ & =\frac{e^{z_i/T}}{{\sum }_j{e}^{z_j/T}}\\ & =q_i(z).\end{array}\text{\hspace{1em}(4)}$$

因此，可以分别将教师和学生 logits 减去各自的均值，而不改变它们对应的概率分布。为简化符号，下文仍将中心化后的 logits 记为 $z_i$ 和 $v_i$，于是

$$\sum _j{z}_j=0,\sum _j{v}_j=0.$$

下面令 $T\to \infty$，并假设 logits 和类别数 $N$ 固定。定义

$$\epsilon =\frac{1}{T}\underset{j}{\max }\max (|z_j|,|v_j|).\text{\hspace{1em}(5)}$$

这样有

$$\left|\frac{z_j}{T}\right|\le \epsilon ,\left|\frac{v_j}{T}\right|\le \epsilon ,$$

且当 $T\to \infty$ 时，$\epsilon \to 0$。

对 $|x|\le \epsilon$，根据带拉格朗日余项的一阶泰勒展开，

$$e^x=1+x+\frac{x^2}{2}{e}^{\xi },|\xi |\le |x|.\text{\hspace{1em}(6)}$$

所以余项满足

$$\left|\frac{x^2}{2}{e}^{\xi }\right|\le \frac{{\epsilon }^2}{2}{e}^{\epsilon }.$$

将其用于分子：

$$e^{z_i/T}=1+\frac{z_i}{T}+R_i,|R_i|\le \frac{{\epsilon }^2}{2}{e}^{\epsilon }.\text{\hspace{1em}(7)}$$

再看分母。由于 logits 已中心化，线性项之和为零：

$$\begin{array}{rl}\sum _j{e}^{z_j/T}& =\sum _j\left(1+\frac{z_j}{T}+R_j\right)\\ & =N+R_N,\end{array}\text{\hspace{1em}(8)}$$

其中

$$R_N=\sum _j{R}_j,|R_N|\le \frac{N{\epsilon }^2}{2}{e}^{\epsilon }.$$

把分母中的 $N$ 提出来：

$$q_i=\frac{1}{N}\cdot \frac{1+\frac{z_i}{T}+R_i}{1+\frac{R_N}{N}}.\text{\hspace{1em}(9)}$$

记

$$u=\frac{R_N}{N}.$$

由上式可知 $u=O({\epsilon }^2)$。当 $\epsilon$ 足够小时，有 $|u|<1$，因此可以使用

$$\frac{1}{1+u}=1-u+O(u^2).\text{\hspace{1em}(10)}$$

这里的 $-u$ 本身也是 $O({\epsilon }^2)$，不能误认为它是常数项；而 $u^2=O({\epsilon }^4)$，属于更高阶项。因此

$${\left(1+\frac{R_N}{N}\right)}^{-1}=1+O({\epsilon }^2).\text{\hspace{1em}(11)}$$

代回式 (9)，得

$$\begin{array}{rl}{q}_i& =\frac{1}{N}\left(1+\frac{z_i}{T}+R_i\right)\left(1+O({\epsilon }^2)\right)\\ & =\frac{1}{N}+\frac{z_i}{NT}+O({\epsilon }^2).\end{array}\text{\hspace{1em}(12)}$$

对教师分布采用同样的展开：

$$p_i=\frac{1}{N}+\frac{v_i}{NT}+O({\epsilon }^2).$$

两式相减，常数项抵消：

$$\boxed{q_i-p_i=\frac{z_i-v_i}{NT}+O({\epsilon }^2)}\text{\hspace{1em}(13)}$$

这说明在高温极限下，教师与学生的概率差的主项也按 $1/T$ 衰减。

这里的渐近结论是在 logits 和类别数固定的条件下成立的。余项中的常数可以依赖这些固定值。

3.3. 合并两部分结果

将式 (13) 代入式 (3)：

$$\begin{array}{rl}\frac{\partial L_{\mathrm{KD}}}{\partial z_i}& =\frac{1}{T}(q_i-p_i)\\ & =\frac{z_i-v_i}{N{T}^2}+O\left(\frac{{\epsilon }^2}{T}\right).\end{array}\text{\hspace{1em}(14)}$$

由于 $\epsilon =O(T^{-1})$，所以

$$\boxed{\frac{\partial L_{\mathrm{KD}}}{\partial z_i}=\frac{z_i-v_i}{N{T}^2}+O(T^{-3})}$$

因此，未乘温度补偿因子的蒸馏梯度，其主项按 $T^{-2}$ 衰减。原因有两个：

1. Softmax 求导本身带有一个 $1/T$。
2. 高温下，教师和学生的概率差 $q_i-p_i$ 的主项也按 $1/T$ 衰减。

对蒸馏损失乘上 $T^2$ 后，

$$\begin{array}{rl}\frac{\partial (T^2{L}_{\mathrm{KD}})}{\partial z_i}& =T^2\frac{\partial L_{\mathrm{KD}}}{\partial z_i}\\ & =\frac{z_i-v_i}{N}+O(T{\epsilon }^2).\end{array}\text{\hspace{1em}(15)}$$

因为 $\epsilon =O(T^{-1})$，余项为 $O(T^{-1})$。于是，在高温近似下，梯度主项不再随温度增大而衰减。

这就是经典知识蒸馏中使用 $T^2$ 的原因：补偿高温带来的主导梯度缩放，使改变温度时软目标损失与硬目标损失的相对贡献大致保持稳定。 [1]

不过，$T^2$ 并不能保证有限温度下两项损失的实际梯度贡献完全相等。实际训练效果仍然取决于教师和学生的预测分布、数据、模型状态、损失权重以及损失的归约方式等因素。

4. Softmax 的导数

上一节使用了带温度 Softmax 的导数。下面从头推导该公式。

令

$$S=\sum _j{e}^{z_j/T},q_k=\frac{e^{z_k/T}}{S}.$$

4.1. 当 $k=i$ 时

分子和分母都依赖 $z_i$，使用商法则：

$$\begin{array}{rl}\frac{\partial q_i}{\partial z_i}& =\frac{\frac{1}{T}{e}^{z_i/T}S-e^{z_i/T}\frac{1}{T}{e}^{z_i/T}}{S^2}\\ & =\frac{1}{T}{q}_i(1-q_i).\end{array}\text{\hspace{1em}(16)}$$

4.2. 当 $k\ne i$ 时

此时分子 $e^{z_k/T}$ 不依赖 $z_i$，只有分母发生变化：

$$\begin{array}{rl}\frac{\partial q_k}{\partial z_i}& =\frac{0\cdot S-e^{z_k/T}\frac{1}{T}{e}^{z_i/T}}{S^2}\\ & =-\frac{1}{T}{q}_k{q}_i.\end{array}\text{\hspace{1em}(17)}$$

4.3. 合并表达式

用 Kronecker delta 将两种情况合并：

$$\boxed{\frac{\partial q_k}{\partial z_i}=\frac{1}{T}{q}_k({\delta }_{ki}-q_i)}\text{\hspace{1em}(18)}$$

其中

$${\delta }_{ki}=\{\begin{array}{ll}1,& k=i,\\ 0,& k\ne i.\end{array}$$

验证可知：

• 当 $k=i$ 时，式 (18) 变为 $\frac{1}{T}{q}_i(1-q_i)$，与式 (16) 一致。
• 当 $k\ne i$ 时，式 (18) 变为 $-\frac{1}{T}{q}_k{q}_i$，与式 (17) 一致。

因此，式 (18) 正是带温度 Softmax 的雅可比矩阵元素。

5. 经典蒸馏与大语言模型蒸馏

前面的公式描述的是经典的、基于输出概率分布的知识蒸馏。在这一设置下，教师向学生提供软目标，学生通过 KL 散度拟合教师分布，并可以结合真实标签训练。[1]

在大语言模型中，“蒸馏”还可能指其他形式的知识迁移。例如，教师可以生成推理过程或回答，再用这些生成数据对较小的学生模型进行监督微调。这种做法传递的是教师生成的文本行为，不一定直接计算每个 token 上完整的教师概率分布与学生概率分布之间的 KL 散度。

近年来的一些代表性工作包括：

• DeepSeek-R1：通过强化学习提升推理能力，并利用教师模型生成的推理数据训练较小模型。[2]
• On-Policy Distillation of Language Models：研究如何在学生自身生成的序列上获取教师反馈，缓解学生训练时的输入分布与实际生成分布不一致的问题。[3]
• MiniLLM：提出面向大语言模型的蒸馏方法，使用反向 KL 散度改变学生拟合教师分布的方式。[4]

这些工作说明，大模型蒸馏不只有一种实现方式。前向 KL、反向 KL、教师生成数据以及学生自生成序列，解决的问题和训练目标都可能不同。因此，不能把所有大模型蒸馏方法都视为式 (2) 的直接应用。

6. 参考资料

经典知识蒸馏

1. Hinton, G., Vinyals, O., & Dean, J. Distilling the Knowledge in a Neural Network. arXiv:1503.02531, 2015.
   https://arxiv.org/abs/1503.02531

   本文公式 (1)–(3) 和 $T^2$ 温度补偿解释的核心参考文献。 重点参考该文关于软目标、温度以及高温下梯度缩放的讨论。

大语言模型蒸馏与推理能力迁移

2. Guo, D. et al. DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement learning. Nature, 645, 633–638 (2025).
   https://doi.org/10.1038/s41586-025-09422-z

   原始预印本：DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948, 2025.
   https://arxiv.org/abs/2501.12948

   可用于了解推理模型训练及通过生成的推理数据向较小模型迁移能力的做法。它并非式 (2) 中 KL 散度与 $T^2$ 补偿的直接推导来源。

3. Agarwal, R. et al. On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. International Conference on Learning Representations (ICLR), 2024.
   https://arxiv.org/abs/2306.13649
   ICLR 论文页面

   重点关注学生自生成序列、教师反馈，以及训练分布与实际生成分布之间的关系。

4. Gu, Y., Dong, L., Wei, F., & Huang, M. MiniLLM: Knowledge Distillation of Large Language Models. International Conference on Learning Representations (ICLR), 2024.
   https://arxiv.org/abs/2306.08543
   ICLR 论文页面

   重点关注反向 KL 蒸馏目标及其与经典前向 KL 蒸馏的区别。