模型蒸馏开山作：Distilling the Knowledge in a Neural Network

Abstract

要提高几乎所有机器学习算法的性能，一个非常简单的方法是在相同的数据上训练许多不同的模型，然后平均它们的预测。不幸的是，使用一整套模型进行预测是很麻烦的，而且计算成本可能太高，以至于无法部署到大量用户，特别是如果单个模型是大型网络的话。Caruana和他的合作者[1]已经证明，将集合中的知识压缩成一个更易于部署的单一模型是可能的，我们使用不同的压缩技术进一步发展了这种方法。我们在MNIST上取得了一些令人惊讶的结果，并且证明了我们可以通过将模型集合中的知识提取到单个模型中来显著地改进一个使用频繁的商业系统的声学模型。本文还提出了一种由一个或多个完整模型和许多特殊模型组成的新型集成，这些模型学习如何区分完整模型混淆的细粒度类。不同于专家的混合，这些专家模型可以快速并行地训练。

1 Introduction

许多昆虫的幼虫形态最适合从环境中提取能量和营养，而成虫形态则完全不同，最适合旅行和繁殖的不同需求。大规模机器学习，我们通常在训练阶段和部署阶段使用非常相似的模型，尽管它们的要求非常不同：对于像语音和对象识别这样的任务，训练必须从非常大、高度冗余的数据集中提取结构，但它不需要实时操作，而且可能会占用大量资源计算。然而，部署到大量用户中，对延迟和计算资源有更严格的要求。他与昆虫的类比表明，如果能更容易地从数据中提取结构，我们应该愿意训练非常繁琐的模型。笨重模型可以是一个单独训练的模型的集合，也可以是一个用非常强的正则化器（如dropout）训练的单个非常大的模型[9]。当笨重的模型被训练好后，我们可以使用一种不同的训练，我们称之为蒸馏，把知识从笨重的模型转移到一个更适合部署的小模型。Rich Caruana和他的合作者已经开创了这一策略的先河[1]。在他们的重要论文中，他们令人信服地证明了一个大的模型集合所获得的知识可以转移到一个小模型上。
我们倾向于在训练过的模型中使用学习过的参数值来识别知识，这使得我们很难看到如何在保持相同知识的同时改变模型的形式。这可能阻碍了对这种非常有前途的方法进行更多的研究。知识的更抽象的观点是，知识是从输入向量到输出向量的学习映射，它使知识从任何特定的实例化中解放出来。或者是学习在大量的类之间进行区分的笨重模型，正常的训练目标是最大化正确答案的平均对数概率，但是学习的副作用是训练的模型将概率分配给所有不正确的答案，即使这些概率很小，其中有些比其他的大得多。不正确答案的相对概率告诉我们很多关于这个笨拙的模型倾向于归纳的方法。例如，一辆宝马的形象被误认为是一辆垃圾车的可能性很小，但被误认为是一根胡萝卜的可能性仍然要高很多倍。
人们普遍认为，用于训练的目标函数应该尽可能地反映用户的真实目标。尽管如此，当真正的目标是将模型很好地归纳为新数据时，模型通常被训练以优化training data的性能。很明显，训练模型使其能够很好地进行概括会更好，但是这需要关于正确的泛化方法的信息，而这些信息通常是不可用的。然而，当我们从一个大模型提取知识到一个小模型时，我们可以训练小模型在与大模型相同的路径上进行泛化。如果繁琐的模型能够很好地进行泛化，例如，它是不同模型的大型集成的平均值，那么以相同方式进行泛化训练的小模型在测试数据上的表现通常要比在训练集成所用的同一训练集上以正常方式训练的小模型要好得多。
将笨重模型的泛化能力转化为小模型的一个明显的方法是用笨重模型产生的类概率作为软目标来训练小模型。或者在这个转移阶段，我们可以使用相同的训练集或者单独的转移集，当这个复杂的模型是一大堆简单模型的集合时，我们可以使用它们各自的预测分布的算术或几何平均值作为软目标。当软目标具有高熵时，它们在每个训练案例中比硬目标提供更多的信息，并且在训练案例之间的梯度差异更小，所以这个小模型通常可以在比原始笨重模型更少的数据上进行训练，并且使用更高的学习率。卡鲁阿纳和他的合作者绕过这个问题通过使用logits (最后softmax的输入)而不是将softmax产生的概率作为学习的目标来学习小模型，然后最小化由繁琐模型产生的对数和由小模型产生的对数之间的平方差。更一般的解决方法，称为蒸馏，是提高最终softmax的温度，直到繁琐的模型产生一个适当的软目标集。然后我们在训练小模型时使用相同的高温来匹配这些软目标。我们稍后会说明，匹配繁琐模型的逻辑实际上是一种特殊的蒸馏方法。
用于训练小模型的转移集可以完全由未标记数据[1]组成，也可以使用原始的训练集。我们发现，使用原始的训练集表现很好，特别是如果我们在目标函数中加入一个小的项，鼓励小模型预测真实目标，同时匹配由笨重的模型提供的软目标。通常情况下，小型模型不能精确地匹配软目标，而在正确答案的方向上出错是有帮助的。

2 Distillation

在精馏的最简单形式中，知识转移到精馏模型中，通过在传输集上训练知识，并在传输集中对每种情况使用软目标分布，该传输集是由在其softmax中使用具有高温的繁琐模型生成的。在训练蒸馏的模型时使用相同的高温，但当它训练结束后，它使用的温度为1。当所有或部分传输集的正确标签已知时，通过训练提取的模型来生成正确的标签，这种方法可以得到显著的改进。一种方法是使用正确的标签来修改软目标，但是我们发现更好的方法是简单地使用两个不同目标函数的加权平均值。第一个目标函数是软目标之间的交叉熵，该交叉熵的计算采用与从冗余模型生成软目标相同的高温方法。第二个目标函数是具有正确标签的交叉熵。第二个目标函数是带有正确标签的交叉熵。这是使用与蒸馏模型的softmax中完全相同的logits计算的，但温度为1。我们发现，在第二目标函数上使用一个条件较低的权重通常可以获得最佳结果。由于软目标产生的梯度大小为 $\frac{1}{T^2}$ ，因此在使用硬目标和软目标时，将其乘以 $T^2$ 是很重要的。这就保证了在用元参数实验时改变蒸馏温度时，硬目标和软目标的相对贡献基本不变。

2.1 Matching logits is a special case of distillation

温度越低，蒸馏就越不注意比平均值负得多的logits。由于这些逻辑几乎完全不受用于训练笨重模型的代价函数的约束，因此它们可能非常嘈杂。另一方面，非常负的logits可以传递关于笨重模型所获得的知识的有用信息。这些影响中哪一个占主导地位是一个经验问题。我们发现，当蒸馏模型太小而无法捕捉到数量模型中的所有已知参数时，中间温度的作用最好，这强烈地表明忽略大的负逻辑是有帮助的。

3 Preliminary experiments on MNIST

dropout可以被看作是训练一个指数级的共享权重的模型集合的方法。
软目标可以将大量的知识转移到提取的模型中，包括从转换的的训练数据中学习到的关于如何泛化的知识，即使转移集不包含任何转换。

温度范围在2.5到4之间的效果明显好于更高或更低的温度。
可以预测未在转移集出现的标签。

5 Training ensembles of specialists on very big datasets

训练一个模型集是利用并行计算的一种非常简单的方法，而一个模型集在测试时需要太多的计算量这一常见问题可以通过精馏来解决。然而，对于集成还有一个重要的反对意见：如果个体模型是大型神经网络，并且数据集非常大，那么即使很容易并行化，训练时所需的计算量也太大。
在本节中，我们将给出一个这样的数据集的例子，并且我们将介绍怎样的学习专家模型，即每一个专注于不同的易混淆类子集可以减少学习集成所需的计算总量。专注于精细区分的专家们的主要问题是他们很容易过度拟合，我们描述了如何通过使用软目标来防止这种过度拟合。

5.1 The JFT dataset

JFT是一个内部的谷歌数据集，它拥有1亿张带标记的图像和1.5万个标签。在我们进行这项工作时，用于JFT的谷歌基线模型是一个深度卷积神经网络[7]，在大量核上使用异步随机梯度下降进行了大约6个月的训练。这个训练使用了两种类型的并行性[2]。
首先，在不同的核集上运行了许多神经网络的副本，并处理来自训练集的不同的小批量。每个副本计算其当前小批量的平均渐变，并将此平均梯度发送到分片参数服务器，该服务器将返回参数的新值。这些新值反映了自最后一次向副本发送参数以来参数服务器接收到的所有梯度。第二，每一个复制品通过在每个核心上放置不同的神经元子集而分布在多个核。集成训练是第三种类型的并行，它可以围绕其他两种类型进行，但前提是有更多的核心可用。等待数年来训练一个模型的集合不是一个选择，所以我们需要一个更快的方法来改进基线模型。

5.2 Specialist Models

当类的数量非常大时，笨重的模型应该是一个集合，其中包含一个针对所有数据训练的泛型模型和许多“专家”模型，每个模型都是基于非常丰富的数据进行训练的，这些数据来自于类的一个非常容易混淆的子集（比如不同类型的蘑菇）。这类专家的softmax可以通过将所有它不关心的类合并到一个垃圾箱类中来变得更小。为了减少过拟合和分担低级特征检测器的学习工作，用泛型模型的权值初始化每个专家模型。然后，通过训练专家，其中一半的例子来自其特殊子集，一半样本来自训练集的其余部分，对这些权重稍加修改。训练时，我们可以通过增加垃圾桶类的logit，乘以特殊类被过度抽样的比例的对数来修正训练集的偏差。

5.3 Assigning classes to specialists

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

5.4 Performing inference with ensembles of specialists

在调查专家模型被提炼出来后会发生什么之前，我们想看看包含专家的团队是如何表现的。

6 Soft Targets as Regularizers

软目标蒸馏还可以缓解过拟合。
我们在JFT数据集上的实验中使用的专家将他们所有的非专业类都压缩到一个垃圾桶类中。我们允许专业人士在所有类别上都有一个完整的softmax，可能有更好的方法来防止他们过度适应，而不是使用早期停止。专家们专门针对其特殊类别中高度丰富的数据进行训练。这意味着它的训练集的有效大小要小得多，而且它有很强的过度拟合的倾向。这个问题不能通过缩小专业类来解决，因为这样我们就会失去建模所有非专业类所获得的非常有用的转移效果。
我们用3%的语音数据进行的实验强烈地表明，如果一个专家用多面手（generalist）的权重初始化，除了用硬目标训练专家外，我们还可以通过对非特殊类进行软目标训练，使其保留几乎所有关于非特殊类的知识。多面手可以提供软目标。我们目前正在探索这种方法。

7 Relationship to Mixtures of Experts

使用在数据子集上受过训练的专家与混合专家[6]有一些相似之处，后者使用一个门控制网络来计算将每个示例分配给每个专家的概率。在专家学习处理分配给他们的示例的同时，门控网络也在学习根据专家对每个示例的相对区分性能来选择将每个示例分配给哪些专家。利用专家的辨别能力来确定学习任务比简单地对输入向量进行聚类并为每个簇分配一个专家要好得多，但这使得训练难以并行化：首先，每个专家的加权训练集不断变化，这取决于所有其他专家；其次，门控网络需要比较不同专家在同一个例子中的表现，以了解如何修正其分配概率。这些困难意味着专家的混合体很少用于他们可能最有益的领域：任务具有包含明显不同子集的巨大数据集。
同时培养多个专家要容易得多。我们首先训练一个多面手模型，然后使用混淆矩阵来定义训练专家的子集。一旦确定了这些子集，专家就可以完全独立地接受培训。在测试时，我们可以使用通用模型的预测来决定哪些专家是相关的，并且只需要运行这些专家。

8 Discussion

我们已经证明，提取对于将知识从一个集合或从一个大的高度正则化的模型转移到一个更小的、经过提炼的模型非常有效。即使用于训练蒸馏模型的传递集缺少一个或多个类的任何示例，在MNIST蒸馏中也能非常好地工作。对于一个深度声学模型，它是Androidvoice搜索所使用的模型的一个版本，我们已经表明，通过训练一个深度神经网络的一个单元，几乎所有的改进都可以被提炼成一个大小相同的神经网络，这是更易于部署的。对于真正大的神经网络来说，即使训练一个完整的集合也是不可行的，但是我们已经证明，一个经过很长时间训练的单一的非常大的网络的性能可以通过学习大量的专家网络来显著提高，每一个网络都能在高度混乱的情况下学会区分不同的类集群。我们还没有证明我们可以将专家的知识提炼回单一的大网络中。