浅谈知识蒸馏（Knowledge Distillation）

浅谈知识蒸馏（Knowledge Distillation）

前言：

在实验室做算法研究时，我们最看重的一般是模型精度，因为精度是我们模型有效性的最直接证明。而在公司做研发时，除了算法精度，我们还很关注模型的大小和内存占用。因为实验室模型一般运行在服务器上，很少有运算资源不足的情况，但是公司研发的算法功能最终都是要部署到实际的产品上的，像手机或者小型计算平台，其运算资源是很有限的。所以算法工程师在公司做预研时，算法建模一般都分两部分：先根据需求建模，并尽可能提高模型精度；然后进行模型压缩，在保证算法精度的情况下尽可能减少其参数量。

常用的模型压缩方法： 知识蒸馏、权重共享、模型剪枝、网络量化以及低秩分解。本文我们主要介绍知识蒸馏。
上图为模型剪枝示意图：左边为训练好的大模型，通过剪枝删除掉一切权值趋于零的节点，达到缩减模型参数的效果

知识蒸馏（Knowledge Distillation，KD)：

我一直觉得训练神经网络的过程很像求解线性方程组，用已知的数据及标签（对应求解方程中的xy点对）来拟合一批模型参数（对应方程组的系数矩阵）。一般来讲，在数据量有限的情况下，如果我们的模型过大，就很容易出现过拟合现象，此时我们需要缩减模型参数，或者添加正则项。

但在数据量足够的情况下，网络模型越复杂、参数量越大，训练出的模型性能会越好，而较小的网络却很难达到大网络那么好的效果。要让一个小网络达到和大模型相近的性能，我们就需要换一个思路，让大模型在训练过程中学习到的知识迁移到小模型中，而这个迁移的过程就叫做知识蒸馏（Knowledge Distillation，KD)。

知识蒸馏的开山之作为大佬Hinton发表在NIPS2014文章《Distilling the Knowledge in a Neural Network》。其主要思想是：在给定输入的情况下训练迁移模型（Student Network），让其输出与原模型（Teacher Network）的输出一致，从而达到将原模型学习到的知识迁移至小网络的目标。

上图为知识蒸馏模型训练示意图：左侧为大参数量的原模型（Teacher Network），右侧为小网络（Student Network）

训练过程中，原模型（Teacher）输出 vi 与小模型（Student）输出 zi 之间的一致性约束是知识蒸馏的关键所在，即最小化下式：

对于输出一致性约束，常用的一般为各种距离度量、或者K-L散度等。在神经网络模型中，训练模型就是让模型的softmax输出与Ground Truth匹配；而知识蒸馏任务中，我们需要让Student网络与Teacher网络的的softmax输出尽可能匹配。

下式定义为普通的Softmax函数：

从上面softmax函数的定义式中我们不难看出，它先通过指数函数拉大输出节点之间的差异，然后通过归一化输出一个接近one-hot的向量（其中一个值很大，其他值接近于0）。对于普通的分类等任务，这样的操作没什么问题，但在知识蒸馏中，这种one-hot形式的输出对于知识的体现很有限，并不利于Student网络的学习（容易放大错误分类的概率，引入不必要的噪声）。所以我们通过引入一个温度参数T来将softmax输出的hard分布转化为soft。

加温度参数T后的softmax定义如下：

上述公式可以理解为：将网络的输出除以温度参数T后再做softmax，这样可以获得比较soft的输出向量：向量中每个值介于0~1之间，各个值之间的差异没有one-hot那么大。并且T的数值越大，分布越缓和。

训练过程中，模型总体的损失函数由两部分组成如下所示：

其中，Alpha和Beta为权重参数，Lsoft 为Distill Loss，保证Student网络输出与Teacher网络输出保持一致性，其定义如下：

其中，pj 和 qj 分别为Teacher网络和Student网络在温度T下的softmax输出向量的第j个值。

因为Teacher网络虽然已经经过了预训练，但其输出也会有一定的误差，为了降低将这些误差迁移到Student网络的可能性，在训练时还添加了Lhard ：通过Ground-truth对Student网络的约束损失，定义如下：

其中，cj 为第j个类别的Ground-truth，qj 为Student网络softmax输出向量的第j个值。