【论文阅读】Distilling the Knowledge in a Neural Network

论文链接
参考博客

知识蒸馏(knowledge distillation)

动机

在机器学习中，我们倾向于使用非常相似的模型进行训练(train)和部署(depolyment)，尽管这两个阶段显然有不同的需求：

• 训练阶段(train)可以使用大型数据集训练深度模型，消耗大量计算资源和时间来获得最佳结果
• 推理阶段(inference)对时间和计算资源占用特别敏感

首先可以明确的一点是，我们需要训练复杂模型来捕捉到数据中足够多的特征。那么为了满足推理阶段对时间的要求，可以使用一个不同的训练方式，将复杂模型学到的“知识‘’(knowledge)转移到一个更适合部署的小模型中。transfer knowledge的过程就叫做知识蒸馏(knowledge distillation)。这里的“知识”，从广义上来说，是从输入向量到输出向量学到的一个映射

使用知识蒸馏的另外一个好处在于，训练模型时我们通常优化模型在训练数据上的表现，但其实我们希望达到的效果是模型对于新来的数据也能很好的处理，即希望模型具有良好的泛化性能。为了训练出具有良好泛化性能的模型，可以使用蒸馏的方法。在将知识从大模型转移到小模型的过程中，小模型可以学习到大模型泛化的方法。比如说，如果一个大模型的泛化性能来自于它是多个不同模型的集合，那么用相同的方法、以提高泛化性能为目的训练的小模型在测试数据上的性能会好过在训练数据上用训练这些不同模型的方法得到的模型

将知识从大模型转移到小模型的一个方式可以是利用大模型产生的各个类上的概率作为soft target来训练小模型。在转移阶段，采用原有的训练数据或者更小的转移数据都可以。因为soft target是介于0-1之间的概率组成的向量，而相应的ground truth（称为hard label）旺旺是one-hot向量，所以soft target包含更多的信息，因此小模型可以在更小的数据集上使用更大的学习率训练

使用方法

对于图像分类问题来说，在将大模型生成的logits送入softmax之前，除以一个比较大的temperature值T，得到更加平滑的分布作为soft targets。训练小模型时也使用这个T，在目标函数中增加一项Ls来减小soft targets和小模型生成的预测值之间的差别。使用新增加的该项与原有的目标函数Lh（hard label与生成的预测值之间的差别）的加权平均作为最终的目标函数，即$Loss=\lambda L_s+(1-\lambda )L_h$。通常初始Ls权重更高，之后慢慢减小

训练结束后，即推断阶段中，T值恢复为1

结构可以用上图表示

实验

MNIST

使用60,000个Minist数据进行训练，测试数据上的错误数量：

模型	结构	结果
Teacher	2个隐藏层，1200个神经元	67 error
Student	2个隐藏层，800个神经元	146 error
Student in KD	2个隐藏层，800个神经元	74 error

实验结果表明，通过知识蒸馏，student能在控制模型大小的同时得到接近teacher的结果

speech recognition

语音识别任务可以通过训练神经网络执行逐帧分类，目标是局部最小化预测值与对应的真值之间的交叉熵

实验的结果显示将一个模型集合蒸馏为一个单独的模型，比直接从相同训练数据学习相同大小的模型工作得好得多