知识蒸馏（Knowledge Distillation，KD）论文阅读笔记

一、引言

将一个预训练的大模型进行压缩和轻量化，使得模型能够部署在各种算力资源较少的嵌入式设备上，是如今的一个研究热点之一。知识蒸馏(Knowledge Distillation)是模型压缩的一个重要方法，其中KD方法可以说是知识蒸馏的起源，值得深入研究。

论文：《Distilling the Knowledge in a Neural Network》（NIPS 2014）

二、知识蒸馏方案的提出

1. knowledge transfer想法的提出

一个模型从诞生到投入实际应用，可以分为训练和部署两个主要阶段。在训练阶段，我们无需考虑训练成本、模型尺寸和实时性等要求，可以构建一个大型网络或者将许多简单模型集成到一起，投入大量算力、海量数据进行训练来达到我们的目标。但是，在部署阶段，我们需要考虑到硬件部署环境的算力资源、算法的实时性等，笨重的大模型往往难以满足部署要求。

因此，作者提出了一个方案：能否像老师教学生一样，把笨重的大模型学习到的知识转移给精简的小模型上，来满足我们的部署要求？

这就涉及到一个概念问题：对于要转移的知识应该如何定义？显然，把大模型学习到的权重等参数作为知识是行不通的。

2. knowledge的定义与模型的泛化能力

第一个想法是，以一个多分类模型为例，我们可以把大模型从softmax层输出的对各类别的预测概率/置信度作为soft targets作为要学习的知识提供给小模型训练。这里，soft targets这一概念相对于 非0即1的hard labels 给出。那么为什么不选择hard labels而选择soft targets呢？

类别	soft targets	hard labels
1	0.05	0
2	0.15	0
3	0.8	1

这是因为，我们训练一个模型的真正目的，不是为了让模型在一个给定的数据集上表现得有多好，而是为了让模型具备良好的泛化能力，即在未知数据集上也有很好的表现。

以手写数字分类器为例，在上面的表格中，我假设了一个输入图像的两种输出结果：一种是softmax层输出的置信度结果，另一种是离散的hard labels结果。两种结果表示的是一个意思，但是显然soft targets结果要包含更多的信息：输入的手写数字像3的概率是0.8，像2的概率是0.15，像1的概率是0.05。即：soft targets不仅包含了对正确答案的预测信息，还包含了对其他不正确答案的预测信息，这些信息都隐含了大模型实际学到的**“知识”**，对模型的泛化是非常重要的。此外，在不同的训练样本中，soft targets方案的梯度分散度更低，更soft，因此小模型可以比大模型使用更少的数据、更高的学习率来训练。

3. 蒸馏温度T

但是，使用softmax层直接输出的置信度结果即soft targets也存在缺点。论文在Chapter 1的末尾给出了一个同样是手写数字识别的例子来证明，笨重大模型往往对于正确答案给出很高的置信度，而对于其他不正确答案的置信度几乎趋近于0，导致在knowledge transfer阶段这些有价值的信息对交叉熵损失函数的影响很小。

于是，在原先思想的基础上，作者正式提出了知识蒸馏的方案：提升最终输出的softmax函数值的温度来使得大模型生成的targets变得softer以满足我们的要求；同时，对于小模型的训练我们使用同样的温度，这一过程就叫做“蒸馏”。这样，通过反向传播和梯度下降使蒸馏过程的损失函数最小化，我们使小模型的预测结果更加靠近大模型输出的soft targets，以达到知识转移的效果。

蒸馏时使用的训练集，也叫作转移数据集(transfer set)，可以使用无标签的数据集，也可以使用训练大模型时使用的原始数据集。

三、知识蒸馏的具体方案

1. softmax函数计算公式（温度T）

首先，根据上一节对蒸馏温度T的