[论文阅读]知识蒸馏（Distilling the Knowledge in a Neural Network）

---

前言

论文名：Distilling the Knowledge in a Neural Network
论文作者：Geoffrey Hinton et.al.
期刊/会议名：NIPS 2014
本文作者：XMU_MIAO
日期：2020/11/26

摘要

  要提高几乎所有机器学习算法的性能，一个非常简单的方法是在同一个数据上训练许多不同的模型并平均它们的预测效果。不幸的是，使用集成模型进行预测是很麻烦的，而且可能计算成本太高，不能部署到大量用户，特别是当单个模型是大型神经网络时。有研究表示，将集成模型中的知识压缩到易于部署的单个模型是有可能的，我们使用不同的压缩技术进一步发展了这种方法。
  我们在MNIST上取得了一些令人惊喜的结果，并且我们表明，通过将集成模型的知识提炼到单个模型中可以显著改进大量应用于商业系统的声学模型。我们还引入了一种新型的由一个或多个完整模型和许多能够学习区分完整模型会混淆的细粒度类的“专家模型”。不同于“专家模型”的混合，这些“专家模型”可以快速地并行训练。

一、Introduction

  许多昆虫的幼年形态是最适合从环境中汲取能量和营养的，而成虫形态则完全不同，更适合旅行和繁殖等不同需求。昆虫的类比表明我们可以训练非常复杂的模型，其易于从数据中提取出结构。这个复杂的模型可以是独自训练模型的集成，也可以是一个用强大正则器如$dropout$训练的单个大模型。一旦复杂模型训练完毕，之后我们可以使用一种不同的训练方式，称之为“蒸馏”，将知识从复杂的模型（称之为$teacher$模型）转移到更易于部署的小模型（称之为$student$模型）中。
  对于$teacher$模型，其能够学习区分大量的类别，正常情况下，训练目标是最大化正确类别的平均对数概率，但这种学习的副作用是训练的模型会将概率分配给错误的类别上，虽然这些概率值可能很小，但一些错误类别比其他错误类别的概率值大很多。比如：一辆宝马车的照片可能有很小的概率被误认为垃圾车，但是仍比被误认为一根萝卜的概率大出很多。在错误类别上的相对概率可以反映出模型是如何进行泛化的。
  将$teacher$模型的泛化能力迁移到$student$模型的明显方法是使用由$teacher$模型的类别概率作为“软目标”用于训练$student$模型。 对于这个迁移阶段，我们可以使用相同的训练集或单独的“迁移”集。当$teacher$模型是小模型的集成时，我们使用小模型各自预测的概率分布的算术或几何平均作为“软目标”。
  当软目标具有高熵值，在训练每一个样本时软目标能够提供比硬目标（$student$模型的$groundtruth$）更多的信息并且训练每一个样本时的梯度差异更小。因此，与$teacher$模型相比，$student$模型训练数据要少得多，使用的学习率也高得多。
  对于像$mnist$一样的任务，$teacher$模型通常能够在正确类别上有高置信度，但关于所学函数的许多信息存在于软目标中极小概率的类别比例中。例如，对于一个版本的数字$2$，可能有$10^{-6}$的概率被认为是数字$3$,$10^{-9}$的概率被认为是数字$7$，而另外一种版本的情况可能相反（每种版本指的是数据集中不同形态是数字$2$）。这是非常有价值的信息，它在数据上定义了丰富的相似性结构（例如，它说明了哪个版本的$2$看起来像$3$,哪个像$7$），但它对迁移阶段的交叉熵代价函数的影响很小，因为这些概率值接近零。 之前的研究解决该问题的方式是使用$logits$（未经过$softmax$函数的值）而不是经过$softamx$函数之后的概率值，他们将$teacher$的$logits$值和$student$的$logits$值的平方差作为最小化目标。
  我们更一般的解决方案，称为“蒸馏”，是提高$softmax$最终值的温度，直到$teacher$模型产生一个合适的软目标集。当训练$student$模型来匹配这些软目标时，我们使用同样高的温度。 稍后说明，之前研究中直接$logits$实际上只是蒸馏的一种特殊情况。

二、Distillation

神经网络通常使用$softmax$输出层来生成类概率，该输出层将每一个$logits{z}_i$生成对应的概率$q_i$
$$q_i=\frac{exp(z_i/T)}{{\sum }_{j}exp(z_j/T)}$$其中$T$表示温度，正常情况下设为$1.0$，使用更高的$T$可以生成更加平滑的类概率分布。
  以下举例一个例子进行说明，其中$logits$是由$[0,5]$均匀分布产生的，经过$softmax$函数之后并设置不同温度$T$的柱状图如下图所示：

  如上图所示，红色柱状图是$logits$经过普通$softmax$函数（$T=1.0$），将温度$T$设置到10以上时，原本的概率分布逐渐趋向均匀分布， 而将温度$T$降低时，只有大的概率值被保留下来，而其它小的值几乎都被削弱变得很小，有点马太效应的意思（强的越强，弱的越弱）。
  最简单的蒸馏形式是通过在迁移集训练上$student$模型将知识迁移到$student$模型中，在迁移集中得每个样本使用由将温度$T$设置得很高的$teacher$模型生成软目标分布。在训练$student$模型时使用相同的高温，但在训练结束后，其使用温度$T=1.0$。在训练$student$模型时，我们使用两个不同目标函数的加权平均，形式如下：$$costfunction=CroEntropy(y_s,y_t)+\alpha CrossEntropy(y_s,y)$$其中$CrossEntropy(\ast )$为交叉熵函数，$y_s$表示$student$模型的预测结果， $y_s$表示$teacher$模型的预测结果， $y$是$student$模型的真实标签。第一个交叉熵函数是$student$模型预测结果与软目标的交叉熵，其使用与训练该软目标的$teacher$中一样的温度。第二个目标函数是$student$模型与真实标签的交叉熵，其温度设置为$1.0$。
  我们发现最好的结果是在第二个目标函数上使用较低的权值（即：$\alpha$取小值）。由于软目标产生的梯度缩放了$1/T^2$，因而在同时使用软目标和硬目标时，将软目标的梯度乘以$T^2$是非常重要的。这确保了在实验过程中，如果用于蒸馏的温度发生改变，那么硬目标和软目标的相对贡献大致保持不变。
  应用蒸馏技术的大致流程如下图所示：

• $teacher$模型是一个大或者集成模型，其训练之后得到$\text{transferset}$，即输入到$teacher$模型的输出对
• 将$student$模型在训练软目标时，要在高温环境下进行，即$T$取一个较大的值，同时与它做交叉熵的$teacher$模型的输出也要在相同的温度情况下
• $student$的硬目标则在$T=1$情况下进行，并且总的优化目标是软目标和硬目标的加权和

三、Experiments

3.1 MNIST

在该数据集上，首先在60000个训练样本上训练了一个带有两个隐藏层（每层有1200个单元）的$teacher$网络，该网络使用$dropout$和$weight-constraints$的正则化方法。另外输入的图片在任意方向上抖动了两个像素。该网络取得了67个测试错误的结果，而一个带有两个隐藏层（每层有800个单元）的$student$网络并不带正则方法取得了146个测试错误的结果。但如果这个$student$网络添加了由温度设置为20的大网络软目标任务，则它的的测试错误结果为74个。这表明了软目标能够将大量的知识迁移到$student$网络中，其中包括了从转译数据中学习到的如何泛化的知识，及时迁移数据集中不包含任何的转译。
  当带有两层隐藏层的$student$网络中每个隐藏层中的单元数据超过300个时，所有高于8的温度设置都能得到相似的结果。但是将隐藏层单元量急剧减小为30个时，温度在2.5-4之间的效果优于这个范围之外的温度设置。
  我们尝试将迁移数据集中数字$3$的样本移除。因此从$student$网络的观点看，数字$3$是一个从未见过的神秘数字。即使是这样，$student$网络也只造成了206个测试错误，其中133个是来自测试集中的1010个数字$3$。

3.2 Speech Recognition

对于语音识别的任务，我们训练了十个与$baseline$完全一样的架构，每个架构都采用不同的随机初始化参数，我们发现通过平均这些模型的预测能够制造出模型的多样性。对于蒸馏，我们尝试了一些温度$[1,\text{2},5,10]$，同时设置硬目标权重$\alpha =0.5$，表中加粗的温度是得到表1结果的最优值，表1如下：

使用10个模型的集成的$teacher$模型所取得的准确率提高的80%能够通过知识蒸馏迁移到$student$网络中，这与在$mnist$任务上的初步实验结果相近。 由于目标函数不匹配，$teacher$模型对$\text{WER}$的改进较小，但集$teacher$模型对$\text{WER}$的改进也迁移到了$studeng$模型上。

3.3 Specialists Models on Very Big Dataset

  如果一个模型是非常大的神经网络，尽管很容易并行化，但在训练时需要的计算量还是太多了。在本节中，我们将给出这样一个例子，并展示学习专家模型如何通过关注不同的易混淆类子集来减少学习集成所需的总计算量。专注于细粒度区分的专家的主要问题是他们很容易过度拟合，我们将描述如何通过使用软目标来防止这种过度拟合。 使用的数据集是JFT数据集，JFT是一个内部的谷歌数据集，它拥有1亿张带标记的图像和1.5万个标签。

3.3.1 Specialist Models

当类别数量很大时，将复杂($teacher$)模型作为一个集成是很有意义的，这个集成包含一个对所有数据进行训练的泛型模型和许多专家模型，专家模型在类别中非常容易混淆的子类集上进行训练，专家模型的$softmax$可以通过将所有不关心的类别放进一个“垃圾类别”中变得小得很多。
  为了减少过拟合和分担低级特征检测器的学习工作，每个专家模型都使用泛型模型学习之后的权值进行初始化。然后，通过训练专家模型（其中一半样本来自其自身特殊的训练子集，另外一半从剩余的训练集中随机采样），这些权值被轻微修改。训练结束后，通过将垃圾箱类的$logits$乘以专家模型专属类的过采样比例的对数来矫正有偏训练。

3.3.2 Assigning Classes to Specialists

  为了让专家模型派生出对象类别的分组，我们决定将重点放在整个网络经常混淆的类别上。尽管可以计算混淆矩阵并将其用作寻找聚类（cluster），但我们选择了一种更简单的方法，它不需要真正的标签来构造聚类。
  我们对泛型模型预测结果的协方差矩阵应用了一种聚类算法，因而经常被预测在一起的类别$S_m$将被用作一个专家模型$m$的目标。将在线版本的$K-means$算法（基础算法可参考博客k-means algorithm）应用于协方差矩阵的列，得到的合理的聚类，如表2所示，尝试了几种聚类算法都产生相似的结果。

3.3.3 Performing Inference with Ensembles of Specialists

  在研究对专家模型进行蒸馏之前，先看看包含专家模型的集成模型性能如何。除了专家模型，集成模型中还有一个泛型模型，这样就可以处理专家模型无法处理的类别。对于输入图像$\text{x}$，$top-1$分类器两步流程如下：

• 第一步：对于每一个测试用例，根据泛型模型找到$n$个最有可能的类别，本文中设置$n=1$。
• 第二步：取容易混淆的特殊子集$S_m$与$k$有非空交集的所有专家模型$m$，并将这些非空交集称为专家模型的活动集合$A_k$（该集合有可能是空的），然后找到所有类的全概率分布并最小化：$$KL(p^g,q)+\sum _{m\in A_k}KL(p^m,q)$$其中$KL$表示$KL$散度，$p^m$和$p^g$分别表示专家模型和泛型模型的预测结果的概率分布。概率分布$p^m$是在所有专家类$m$加上一个垃圾类的分布，因而当计算概率分布$q$的$KL$散度时，将$q$分配$m$中的垃圾类的所有类概率加起来。

3.3.4 Results

  专家模型训练得非常快，在JFT数据集上只需要几天时间而不是几周，并且所有专家模型都是完全独立训练的。表3展示了$baseline$的测试精度和加上专家模型之后的测试精度：

使用61个专家模型进行集成，模型的相对测试精度增加了$4.4\%$，另外在表4中还展示了只考虑专家模型训练集中类别并将预测限制在类别的子集的准确性上：

  在$\text{JFT}$专家实验中，训练了61个专家模型，每个模型有300个类别（加上垃圾类别），由于专家模型预测的类别不是不相交的，因而经常会有多个专家模型覆盖同一个类别。表4展示了测试集示例的数量，在使用专家时，$top1$正确示例数量的变化，以及$\text{JFT}$数据集$top1$准确度的相对提高百分比，这取决于涵盖该类的专家数量。当一个类别有多个专家模型覆盖时，准确性的提升会更大，这是一个普遍的趋势，这十分令人欣喜，因为独立训练专家模型非常容易并行化。

四、Using Soft Targets to Prevent Specialists from Overfiting

关于使用软目标而不是硬目标的主张之一是，软目标可以携带更多有用的信息，而这些信息不能用单一的硬目标进行编码。在这一节中，通过使用少量数据来拟合之前提到的$\text{85M}$参数量语音识别模型上来证明。表5显示了使用$3\%$的数据量（大概$\text{20M}$个样本）来训练硬目标的$baseline$模型造成了严重的过拟合（提早暂停了训练，而测试精度掉到了$44.5\%$）而相同模型使用软目标能够恢复训练集的大部分信息（只下降了$2\%$）。值得注意的是，使用软目标训练没有必要提前暂停训练，模型的测试精度就收敛到了$57\%$。

以上的分析表明，软目标是交流在所有数据上训练的模型所发现规律给另一个模型的一种非常有效的方式。

五、Discussion

  本文中我们展示了蒸馏在将一个集成模型或一个大的正则化模型中的知识迁移到一个小的$student$模型中效果良好。在MNIST任务上，即使用于训练$student$模型的迁移集缺少一个或多个类的例子，蒸馏的效果也非常好。对于一个应用在安卓语音搜索中的深度声学模型，我们表明训练一个集成的深度模型的几乎所有性能提高都能够提炼到一个大小相同的更易于部署的单一神经网络上。

总结

  本文是知识蒸馏的早期文章之一，知识蒸馏是将$teacher$模型（通常的大模型或集成模型）的知识迁移到
$student$模型（通常是小模型），具体做法是将$teacher$模型的学习到的类分类概率分布作为$student$模型学习的软目标，$student$模型最终的目标函数是软目标与自身训练的硬目标的加权和。 作者在$mnist$任务和声学模型上进行实验，结果显示知识蒸馏的效果良好，能够将知识进行很好的迁移。文中作者还介绍了一种集成专家模型，能够对易于混淆的类别进行区分，并具有良好的并行性。本文知识蒸馏主要有以下几个要点需要注意：

• 软目标需要在高温情况下进行（$T$取较大值），而硬目标则为常规的$softmax$函数（$T=1.0$）
• 软目标和硬目标的加权和作为最终优化目标时，软目标梯度更新时需要乘以$T^2$
• 知识蒸馏有一定的防止过拟合的作用
• 知识蒸馏起到了与网络剪枝一样的效果，压缩了模型的大小并且保持了模型的性能