蒸馏网络 Distilling the Knowledge in a Neural Network

论文：https://arxiv.org/abs/1503.02531

0 摘要

提高几乎所有机器学习算法性能的一种非常简单的方法是在相同的数据上训练许多不同的模型，然后对它们的预测进行平均[3]。不幸的是，使用整个模型集合进行预测是麻烦的，并且可能在计算上太昂贵而不能允许对大量用户进行部署，尤其是如果各个模型是大型神经网络。Caruana和他的合作者[1]已经证明，有可能将整体中的知识压缩成一个更容易部署的单一模型，并且我们使用不同的压缩技术进一步开发了这种方法。我们在MNIST上取得了一些令人惊讶的结果，我们表明，通过将模型集合中的知识提炼到单个模型中，我们可以显著改善大量使用的商业系统的声学模型。我们还介绍了一种由一个或多个完整模型和许多专业模型组成的新型集合，它们学会区分完整模型混淆的细粒度类。与专家混合不同，这些专业模型可以快速并行地进行训练。

1 简介

许多昆虫具有幼虫形态，其被优化用于从环境中提取能量和营养物，以及完全不同的成体形式，其针对非常不同的旅行和繁殖要求进行了优化。在大规模机器学习中，我们通常在训练阶段和部署阶段使用非常相似的模型，但其实我们对它们的要求非常不同：对于语音和对象识别等任务，训练必须从非常大的、高度冗余的数据集中提取结构，但不会需要实时操作，训练过程可以使用大量的计算。但是，模式部署时对延迟和计算资源有更严格的要求。与昆虫的类比表明，如果能够更容易地从数据中提取结构，我们应该愿意训练非常繁琐的模型。繁琐的模型可以是单独训练的多个模型的集合，也可以是使用非常强的正则化器（例如dropout）训练的单个非常大的模型[9]。一旦繁琐的模型训练完成，我们可以再使用称之为“蒸馏”的训练过程，将知识从繁琐的模型转移到更适合部署的小模型。Rich Caruana及其合作者已经开创了这一战略的一个版本[1]。在他们的重要论文中，他们令人信服地证明，大的模型集合所获得的知识可以转移到一个小型模型中。

可能阻碍对这种非常有前景的方法进行更多概念性研究的问题是我们倾向于使用学习的参数值来识别训练模型中的知识，这使我们很难看到我们如何改变模型的形式但同时保持相同的知识。关于知识的更抽象的观点，使其从任何特定的实例中解放出来，就是它是一种学习的从输入向量到输出向量的映射。对于学会区分大量类别的繁琐模型，正常的训练目标是最大化正确答案的平均对数概率，但学习的副作用是训练的模型将概率分配给所有不正确的答案。即使这些概率非常小，但其中总有一些比其他概率大得多。不正确答案的相对概率告诉我们很多关于繁琐模型如何概括的知识。例如，宝马的图像可能只有很小的机会被误认为是垃圾车，但这个错误仍然比将它被误认为是胡萝卜的可能性高很多倍。

人们普遍认为，用于训练的目标函数应尽可能地反映用户的真实目标。在训练过程中，人们往往以最优化训练集的准确率作为训练目标，但真实目标却是最优化模型的泛化能力。显然如果能以提升模型的泛化能力为目标进行训练是最好的，但这需要正确的关于泛化能力的信息，而这些信息通常不可用。当我们将大型模型的知识蒸馏进小型模型时，我们可以训练小模型以与大型模型相同的方式进行泛化。如果繁琐的模型是因为它是很多不同模型的集成才取得了很好地泛化，那么在同一数据集上，以与繁琐模型相同泛化方式训练的小模型通常比以正常方式训练的小模型具有更好的测试性能。

将繁琐模型的泛化能力转移到小模型的一种显而易见的方法是使用由繁琐模型产生的类概率作为训练小模型的“软目标”。对于此迁移阶段，我们可以使用相同的训练集或单独的“转移”集。当繁琐的模型是较大的简单模型的集合时，我们可以使用其各自预测分布的算术或几何平均值作为软目标。当软目标具有高熵值时，它们为每个训练案例提供比硬目标更多的信息，并且在训练案例之间梯度的变化更小，因此小模型通常可以在比原始繁琐模型少得多的数据上训练并可以使用更高的学习率加快训练过程。

对于像MNIST这样的任务，繁琐的模型几乎总是以非常高的置信度产生正确答案，关于学习函数的许多信息存在于软目标中非常小的概率的概率值中。例如，一个真实值为2的一个版本的图像可以以10-6的概率被识别为3和以10-9的概率被识别为7，而对于另一个版本的实际为2的图像，它可以是另一种识别结果。这是有价值的信息，它定义了数据上丰富的相似性结构（即它表示哪个2看起来像3，哪个看起来像7），但它在迁移阶段对交叉熵损失函数的影响非常小，因为概率都是如此接近于零。Caruana和他的合作者通过使用logits（最终softmax的输入）而不是softmax作为学习小模型的目标产生的概率来解决这个问题，并且他们最小化了由繁琐模型和小模型产生的logits之间的平方差异。我们更通用的解决方案，称为“蒸馏”，是提高最终softmax的温度，直到繁琐的模型产生适当柔软的目标。然后，我们在训练小模型时使用相同的高温来匹配这些软目标。后面将表明，匹配繁琐模型的logits实际上是一个特殊的蒸馏情况。

用于训练小模型的迁移集可以完全由未标记的数据组成[1]，或者是原始的训练集。我们发现使用原始训练集效果很好，特别是如果我们在目标函数中添加一个小项，鼓励小模型预测真实目标同时也匹配繁琐模型提供的软目标。通常，小模型不能与软目标完全匹配，并且在正确答案的方向上犯错是有帮助的。

2 蒸馏

神经网络通常通过使用“softmax”输出层来产生类概率，该输出层通过将$z_i$与其他logit进行比较，将针对每个类计算的logit $z_i$转换为概率$q_i$。
$q_i = \frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)}$

其中T是通常设置为1的温度。使用较高的T值可以在类上产生更软的概率分布。

在最简单的蒸馏形式中，通过在迁移集上训练并将迁移集中的每个类别使用软目标分布表示，将知识转移到蒸馏的小模型，软目标分布通过使用在softmax中具有高温的繁琐模型产生。训练蒸馏模型时使用相同的高温，但训练完成后，温度设置为1。

当已知所有或部分迁移集的正确标签时，通过训练蒸馏模型以产生正确的标签，可以显著改善该方法。一种方法是使用正确的标签来修改软目标，但我们发现更好的方法是简单地使用两个不同目标函数的加权平均值。第一个目标函数是具有软目标的交叉熵，该交叉熵是从蒸馏模型的具有高温的softmax函数中获得的，该高温值与繁重模型生成软目标交叉熵时的温度值相同。第二个目标函数是具有正确标签的交叉熵。这是使用与蒸馏模型的softmax中完全相同的logits计算的，但温度值为1。我们发现通常通过在第二目标函数上使用可忽略不计的较低权重来获得最佳结果。由于软目标产生的梯度的大小为$1 / T^2$，因此当使用硬目标和软目标时，将软目标乘以$T^2$是很重要的。这确保了如果在试验元参数时改变了用于蒸馏的温度，则硬和软目标的相对贡献保持大致不变。

2.1 匹配logits是蒸馏的特例

迁移集中的每个例子相对于蒸馏模型的每个logit - $z_i$贡献交叉熵梯度$dC / dz_i$。如果繁琐的模型具有logit $v_i$，其产生软目标概率$p_i$并且迁移训练在温度为T下完成，则该梯度由下式给出：

$\frac{\partial C}{\partial z_i} = \frac{1}{T}(q_i - p_i)=\frac{1}{T}(\frac{e^{z_i/T}}{\sum_je^{z_j/T}}-\frac{e^{v_i/T}}{\sum_j e^{v_i/T}})$

如果温度高于logits的幅度，我们可以近似：

$\frac{\partial C}{\partial z_i} \approx \frac{1}{T}(\frac{1+z_i/T}{N + \sum_j{z_j/T}}-\frac{{1+v_i/T}}{N+\sum_j {v_i/T}})$

如果我们现在假设对于每个迁移样本，logits均值为零，那么$\sum_j z_j = \sum_j v_j$ = 0。上式可以简化为：
$\frac{\partial C}{\partial z_i} \approx \frac{1}{NT^2}(z_i - v_i)$

因此，在高温极限下，蒸馏相当于最小化$1/2(z_i-v_i)^2$，条件是对于每个迁移样本提供零均值logits。在较低的温度下，蒸馏对匹配比平均值更负的标记的关注要少得多。这是有利的，因为这些对数几乎完全不受用于训练繁重模型的代价函数的限制，因此它们可能非常嘈杂。另一方面，非常负的logits可以传达关于由繁重模型获取的知识的有用信息。这些影响中的哪一个主导是一个经验问题。我们发现，当蒸馏模型太小而无法捕获繁琐模型中的所有知识时，中间温度效果最好，这强烈暗示忽略大的负logit可能会有所帮助。

3 在MNIST上的初步试验

为了了解蒸馏的工作情况，我们在所有60,000个训练案例中训练了一个单个大型神经网络，其中有两个隐藏的1200个整流线性隐藏单元层。如[5]中所述，使用dropout和权重约束强烈地对网络进行正则化。dropout可以被视为一种训练共享权重的指数个大型模型集合的方法。另外，输入的图像在任何方向上最多有两个像素抖动，该网络具有67个测试误差。而较小的具有两层各800个整流线性隐藏单元的网络，没有使用正则化则实现了146个错误。但是，如果较小的网络仅仅通过在20的温度下增加匹配大网络产生的软目标的附加任务来进行正则化，则实现了74个测试误差。这表明软目标可以将大量知识传递给蒸馏模型，包括从迁移训练数据中学习如何泛化的知识，即使迁移集不包含任何变换。

当蒸馏网络在其两个隐藏层中的每一个中具有300个或更多个单元时，所有高于8的温度给出相当类似的结果。但是当每层急剧减少到30个单位时，2.5到4范围内的温度比高温或更低温度的温度要好得多。

然后我们尝试删除转移集中的数字3的所有示例。因此，从蒸馏模型的角度来看，3是一个从未见过的神话数字。尽管如此，蒸馏模型仅产生206个测试错误，其中133个是1010个测试样本中的3。大多数错误是由于对3的学习偏差太低而引起的。如果这个偏差增加3.5（这可以优化测试集的整体性能），蒸馏模型会产生109个错误，其中14个是3。因此，在正确的偏置下，尽管在训练期间从未见过3，但蒸馏模型仍能获得98.6％的针对3的测试准确率。如果转移集仅包含训练集中的7和8，则蒸馏模型的测试误差为47.3％，但当7和8的偏差减少7.6以优化测试性能时，测试误差降至13.2％。

5 在非常大的数据集上训练专家模型集

训练一个模型集合是利用并行计算的一种非常简单的方法，通常的反对意见是模型集合在测试时需要太多的计算，这个问题可以通过蒸馏来处理。对于模型集成还有另一个重要的反对意见：如果单个模型是大型神经网络，并且数据集非常大，那么即使很容易进行并行化处理，在训练时所需的计算量也太多了。

在这一节中，我们给出了这样一个数据集的例子，我们展示了学习专家模型，每个模型都集中在不同的混淆的类子集上，从而减少学习模型集合所需的总计算量。专注于对细微差别进行判断的专家模型的主要问题是他们很容易过拟合，我们描述了如何通过使用软目标来防止这种过拟合。

5.1 JFT数据集

JFT是一个谷歌内部的数据集，它有1亿张带有15000个标签的图像。当我们做本文研究时，谷歌的JFT基线模型是一个深度卷积神经网络[7]，它使用异步随机梯度下降在大量核心上进行了大约6个月的训练。训练使用了两种类型的并行化。首先，使用了数据并行，神经网络在不同的核心上运行，并处理训练集中不同的小批量数据。每个模型副本计算其当前小批量数据的平均梯度，并将此梯度发送到共享的参数服务器，该服务器将返回参数的新值。这些新值反映了参数服务器自上次向所有模型副本发送参数以来接收到的所有梯度。第二，使用了模型并行，每个模型副本通过在每个核心上放置不同的神经元子集而分布在多个核心上。集成学习是第三种并行，但需要大量的计算核心。等待数年来训练一组模型不是一个良好的选择，因此我们需要一种更快的方法来改进基线模型。

5.2 专家模型

当类的数量非常大时，繁琐的模型是一个集合，其中包含一个在所有数据上训练的通用模型和许多“专家”模型，每个专家模型都在从一个非常容易混淆的子类（如不同类型的蘑菇）上进行训练。通过将所有它不关心的类别组合成一个垃圾箱类别，可以使这类专家的Softmax更小。

为了减少过拟合，共享学习低级特征检测器的工作，每个专家模型都是用通用模型的权重初始化的。然后通过训练专家模型对这些权重进行轻微修改，训练集一半来自于其特殊子集，另一半则通过随机从训练集的其余部分进行抽样获得。训练结束后，我们可以通过将垃圾箱类的logit增加到专家类被过采样的概率的对数，从而纠正偏差训练集。

5.3 将类赋予专家模型

为了获得专家对象类别的分组，我们决定关注我们的完整网络经常混淆的类别。尽管我们可以计算混淆矩阵并将其用作查找此类聚类的方法，但我们选择了一种更简单的方法，它不需要真正的标签来构建聚类。

特别是，我们将聚类算法应用于我们通用模型预测的协方差矩阵，因此，通常一起预测的一组类$S^m$将被用作我们的第m个专家模型的目标。我们将K-means算法的在线版本应用于协方差矩阵的列，并获得了合理的聚类（如表2所示）。我们尝试了几种产生不同的聚类算法，都产生了类似的结果。

5.4 使用专家模型的集成进行推理

在调查专家模型被蒸馏后会发生什么之前，我们想看看专家模型的集成表现如何。除了专家模型，我们总是有一个通用模型，以便我们可以处理我们没有专家模型的类别，进而决定使用哪些专家模型。给定输入图像x，我们分两步进行第一类的分类：

• 第一步，对每一个测试样本，我们根据通用模型找出前n个最可能的类别，称之为类别k。在我们的实验中，我们设定n=1；
• 第二步，我们实验所有的专家模型m，针对混淆类$S^m$，其与k有一个非空交集，并将其称为活动的专家组$A_k$（请注意，此集合可能为空）。然后，我们在所有的类上找到完整的概率分布q，即最小化：
  $KL(p^g,q)+\sum_{m \in A_k}KL(p^m,q)$
  KL表示KL散度，$p^mp^g$表示专家模型和通用模型产生的概率分布。分布$p^m$是m的所有专家类加上一个垃圾箱类的分布，因此当从完全q分布计算其KL散度时，我们将完整q分布分配给m的垃圾箱中的所有类的所有概率相加。

5.5 结论

从训练好的基线全网络开始，专家模型训练速度非常快（JFT的几天而不是几周）。此外，所有专家模型都完全独立地被训练。表3显示了基线系统和基线系统与专家模型相结合的绝对测试精度。使用61个专业模型，整体测试精度相对提高4.4％。我们还报告了条件测试准确性，这是仅考虑属于专家类的示例的准确性，并将我们的预测限制在该类子集中。

对于我们的JFT专家实验，我们训练了61个专家模型，每个模型有300个类别（加上垃圾箱课程）。因为专家模型的类集是相交的，所以我们经常有多个专家模型来覆盖特定的图像类。表4显示了测试集示例的数量，使用专家模型时在位置1处正确的示例数量的变化，以及按专家数量细分的JFT数据集的top1准确度的相对百分比改进。当我们有更多的专家来覆盖特定的类别时，我们对总体趋势感到鼓舞，因为训练独立的专家模型非常容易并行化，因此准确性的提高会更大。

6 使用软目标进行正则化

我们关于使用软目标而不是硬目标的主要主张之一是，可以在软目标中携带许多有用的信息，这些信息不可能用单个硬目标编码。在本节中，我们通过使用少得多的数据来拟合前面描述的基线语音模型的85M参数来证明这是一个非常大的影响。表5显示只有3％的数据（大约20M的例子），用硬目标训练基线模型会导致严重的过度拟合（我们进行了提前停止，因为准确度在达到44.5％后急剧下降），而使用软目标训练的同一模型能够恢复完整训练集中的几乎所有信息（约2％下降）。更值得注意的是，我们不必提前停止：具有软目标的系统简单地“收敛”到57％。这表明软目标是将由所有数据训练的模型发现的规律传达给另一个模型的非常有效的方式。

6.1 使用软目标防止专家模型过拟合

我们在JFT数据集实验中使用的专家模型将所有非专家类折叠成一个垃圾箱类。如果我们允许专家系统在所有类别上拥有完整的softmax，那么可能有一种更好的方法来防止它们过度拟合而不是使用早期停止。专家模型接受有关其特殊类别高度丰富的数据的训练。这意味着它的训练集的有效大小要小得多，并且很容易过度拟合它对应的特殊类别。这个问题不能通过让专家变得更小来解决，因为我们失去了从建模所有非专家类别中获得的非常有用的转移效果。

我们使用3％的语音数据的实验强有力地表明，如果专家模型用通用的权重进行初始化，我们可以通过使用非特殊类的软目标训练它来保留关于非特殊类的几乎所有知识而不是用硬目标训练它之外的特殊类别。软目标可以由通用模型提供。我们目前正在探索这种方法。

7 与专家混合的关系

使用经过数据子集训练的专家与专家混合[6]有一些相似之处，后者使用门控网络来计算为每个专家分配每个例子的概率。在专家模型学习处理分配给他们的示例的同时，门控网络正在学习根据该示例的专家的相对判别性能来选择将每个示例分配给哪些专家。使用专家的判别性能来确定学习的分配比简单地聚类输入向量并为每个群集分配专家要好得多，但这使得训练难以并行化：首先，每个专家的加权训练集的变化依赖于所有其他专家;其次，门控网络需要比较同一个例子中不同专家的表现，以了解如何修改其分配概率。这些困难意味着专家的混合很少用于可能最有益的制度：具有包含明显不同子集的大型数据集的任务。

并行化多个专家的训练要容易得多。我们首先训练一个通用模型，然后使用混淆矩阵来定义专家训练的子集。一旦定义了这些子集，就可以完全独立地训练专家模型。在测试时，我们可以使用通用模型中的预测来确定哪些专家是相关的，只需要运行这些专家。

8 讨论

我们已经证明，蒸馏可以很好地将知识从集合或大型高度正规化模型转移到较小的蒸馏模型中。在MNIST蒸馏工作非常好，即使用于训练蒸馏模型的转移装置缺少一个或多个类的任何实例。对于Android声音搜索使用的深度声学模型，我们已经证明通过训练深度神经网络集合所实现的几乎所有改进都可以被提炼成相同大小的单个神经网络。该模型部署起来要容易得多。

对于非常大的神经网络来说，即使是训练一个完整的整体也是不可行的，但我们已经证明，通过学习大量的数据，可以显著提高已经训练了很长时间的单个真正大网络的性能。专家网络，每个网都学会在高度混乱的集群中区分各类。我们还没有表明我们可以将专家的知识提炼回单个大网络。

以下内容摘抄于：https://blog.csdn.net/shuzfan/article/details/53839372

—————————— 背景介绍 ——————————

大家都想要得到一个又好又快的模型，但是实际情况往往是模型越小则性能越差。文献[1]中提出了一种策略：大模型学习到的知识可以通过“提取”的方法转移到一个小模型上 。

所以，本文的宏观策略就是先训练一个性能很好的大模型，然后再据此学习得到一个性能不差的小模型。

—————————— 方法介绍 ——————————

因为整体方法比较简单，涉及到的推导和证明也比较少，所以这里直接给出方法步骤：

(1)训练一个大模型

使用普通的SoftmaxWithLoss训练一个性能好但比较复杂的网络，记作A；
（PS. 该方法只适用于softmax分类）

(2)获取soft target label

将A中的softmax替换为下面的新的soft_softmax, 然后将训练集跑一遍，并将每张训练图片的网络输出结果保存下来。

(1)式中T=1时，退化成传统的softmax，T无穷大时，结果趋近于1/C，即所有类别上的概率趋近于相等。T>1时，我们就能获得soft target label。

这里解释两点：

——为什么称之为soft target label?

举一个栗子，假如我们分三类，然后网络最后的输出是[1.0 2.0 3.0]，我们可以很容易的计算出，传统的softmax(即T=1)对此进行处理后得到的概率为[0.09 0.24 0.67]，而当T=4的时候，得到的概率则为[0.25 0.33 0.42]。
可以看出，当T变大的时候输出的概率分布变得平缓了，这就称之为soft。

PS.不知道怎么算概率的，可以看下我的另外一篇博客caffe层解读系列-softmax_loss

——为什么要使用 soft target label?

通常我们使用softmax进行分类的时候，我们的label都是one shot label，比如我们分三类：猫、虎和猪，那么一张猫的图片它的label就是[1 0 0]。这种标注方式意味着每一类之间都是独立的，完全没有任何联系。但是事实上，猫和虎的相似度应该高于猫和猪的相似度，这种丰富的结构信息，one shot label(hard target)描述不了。此外，正因为one shot label的hard，导致了我们学习得到的概率分布也相对hard。显然对于一张猫的图片强行学习一个[1 0 0]的分布，其难度要比学一个[0.65 0.3 0.05]的分布大得多。

(3)训练一个小网络

重新创建一个小的网络，该网络最后有两个loss，一个是hard loss,即传统的softmaxloss，使用one shot label；另外一个是soft loss，即T>1的softmaxloss，使用我们第二步保存下来的soft target label。

—————————— 参考文献 ——————————

《Buciluǎ C, Caruana R, Niculescu-Mizil A. Model compression[C]//Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining. ACM, 2006: 535-541.》https://www.cs.cornell.edu/~caruana/compression.kdd06.pdf