摘要

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

介绍

在大规模机器学习中,我们通常在训练阶段和部署阶段使用非常相似的模型,尽管它们的要求非常不同:对于语音和对象识别等任务,训练必须从非常大的数据中提取结构,高度冗余的数据集,但它不需要实时运行,并且可以使用大量计算。然而,部署到大量用户对延迟和计算资源的要求要严格得多。如果能够更容易地从数据中提取结构,我们应该愿意训练非常繁琐的模型。笨重的模型可以是单独训练的模型的集合,也可以是使用非常强的正则化器(如dropout)训练的单个非常大的模型。一旦笨重的模型经过训练,我们就可以使用另一种训练,我们称之为“蒸馏”,将知识从笨重的模型转移到更适合部署的小模型。Rich Caruana及其合作者已经开创了这一策略的一个版本。在他们的重要论文中,他们令人信服地证明,由大量模型集合获得的知识可以转化为单个小模型。

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

一般认为,用于培训的目标函数应尽可能地反映用户的真实目标。尽管如此,当真正的目标是很好地推广到新数据时,通常会对模型进行训练,以优化训练数据的性能。显然,最好是训练模型,使其能够很好地进行概括,但这需要关于正确概括方法的信息,而这些信息通常不可用。然而,当我们将知识从一个大模型提取到一个小模型时,我们可以训练小模型以与大模型相同的方式进行推广。例如,如果笨重的模型能够很好地进行泛化,因为它是不同模型的大集合的平均值,那么以相同方式进行泛化训练的小模型在测试数据上的表现通常会比在训练集合所用的相同训练集上以正常方式进行训练的小模型好得多。

笨重模型的泛化能力转化为小模型的一个明显方法是使用笨重模型产生的类概率作为训练小模型的“软目标”。对于这个转移阶段,我们可以使用相同的训练集或单独的“转移”集。当笨重的模型是简单模型的大集合时,我们可以使用它们各自预测分布的算术或几何平均值作为软目标。当软目标具有高熵时,它们在每个训练案例中提供的信息比硬目标多得多,训练案例之间的梯度变化也小得多,因此小模型通常可以在比原始笨重模型少得多的数据上进行训练,并使用更高的学习率。

对于MNIST这样的任务,笨重的模型几乎总是以非常高的置信度生成正确的答案,关于学习函数的大部分信息都存在于软目标中非常小的概率比率中。这是有价值的信息,定义了数据上丰富的相似性结构(即,它表示哪些2看起来像3,哪些看起来像7),但在传输阶段,它对交叉熵成本函数的影响非常小,因为概率非常接近于零。卡鲁阿纳和他的合作者通过使用Logit(最终softmax的输入)而不是softmax产生的概率作为学习小模型的目标,绕过了这个问题,他们最小化了笨重模型产生的Logit和小模型产生的Logit之间的平方差。我们更通用的解决方案称为“蒸馏”,即提高最终softmax的温度,直到笨重的模型产生一组合适的软目标。然后,我们在训练小模型以匹配这些软目标时使用相同的高温。稍后我们将展示,匹配笨重模型的logits实际上是蒸馏的一个特例。

用于训练小模型的传输集可以完全由未标记的数据组成,也可以使用原始的训练集。我们发现,使用原始训练集效果很好,尤其是如果我们在目标函数中添加一个小项,鼓励小模型预测真实目标,并匹配笨重模型提供的软目标。通常情况下,小模型无法精确匹配软目标,在正确答案的方向上出错会有帮助。

知识蒸馏

在最简单的蒸馏形式中,将知识转移到蒸馏模型中,方法是在转移集中对其进行训练,并对转移集中的每种情况使用软目标分布,该软目标分布是通过使用其softmax中温度较高的笨重模型生成的。训练蒸馏模型时使用相同的高温,但训练后使用的温度为1。

当所有或部分传输集都知道正确的标签时,还可以通过训练提取的模型来生成正确的标签,从而显著改进该方法。一种方法是使用正确的标签来修改软目标,但我们发现更好的方法是简单地使用两个不同目标函数的加权平均值。第一个目标函数是与软目标的交叉熵,该交叉熵是使用蒸馏模型softmax中与从笨重模型生成软目标相同的高温来计算的。第二个目标函数是具有正确标签的交叉熵。这是使用与蒸馏模型在softmax中完全相同的logits进行计算的,但温度为1。我们发现,通常通过对第二个目标函数使用一个相对较低的权重来获得最佳结果。由于软目标产生的梯度大小为1/t2,因此在使用硬目标和软目标时,将其乘以t2非常重要。这确保了在使用元参数进行实验时,如果用于蒸馏的温度发生变化,硬目标和软目标的相对贡献基本保持不变。

软目标

传统的分类问题,模型的目标是将输入的特征映射到输出空间的一个点上,例如在著名的Imagenet比赛中,就是要将所有可能的输入图片映射到输出空间的1000个点上。这么做的话这1000个点中的每一个点是一个one hot编码的类别信息。这样一个label能提供的监督信息只有log(class)这么多bit。然而在KD中,我们可以使用teacher model对于每个样本输出一个连续的label分布,这样可以利用的监督信息就远比one hot的多了。

另外一个角度的理解,大家可以想象如果只有label这样的一个目标的话,那么这个模型的目标就是把训练样本中每一类的样本强制映射到同一个点上,这样其实对于训练很有帮助的类内variance和类间distance就损失掉了。然而使用teacher model的输出可以恢复出这方面的信息。具体的举例就像是paper中讲的, 猫和狗的距离比猫和桌子要近,同时如果一个动物确实长得像猫又像狗,那么它是可以给两类都提供监督。

综上所述,KD的核心思想在于”打散”原来压缩到了一个点的监督信息,让student模型的输出尽量match teacher模型的输出分布。其实要达到这个目标其实不一定使用teacher model,在数据标注或者采集的时候本身保留的不确定信息也可以帮助模型的训练。

温度系数

在我们正常的训练过程中,我们只会关注概率最高结果与正确结果的差别。这种相似性完全是通过足够数量的样本构建的。

因此Hinton的想法是:如何充分利用大网络中的这种结果?如果只是构建所有类的传统损失函数的话,小概率结果对损失函数的贡献微乎其微。解决的方法无非是:在计算损失函数时放大其他类的概率值所对应的损失值。

Hinton用一个简单的方法解决了这一问题:加入温度系数T。改进后的softmax如下。

论文拆解《Distilling the Knowledge in a Neural Network》 - 图1

MNIST实验

为了了解蒸馏效果如何,我们在所有60000个训练案例中训练了一个包含两个隐藏层的大型神经网络,其中包含1200个校正线性隐藏单元。使用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,但蒸馏模型在测试3中的正确率为98.6%。如果转移集只包含训练集中的7和8,则提取的模型会产生47.3%的测试误差,但当7和8的偏差减少7.6以优化测试性能时,这将下降到13.2%的测试误差。

语音识别实验

我们将研究在自动语音识别(ASR)中使用的深度神经网络(DNN)声学模型的置乱效果。我们表明,我们在本文中提出的蒸馏策略实现了预期的效果,即将一组模型蒸馏成一个单独的模型,其效果明显优于直接从相同训练数据学习的相同大小的模型。

目前,最先进的ASR系统使用DNN将波形特征的(短)时间背景映射为隐马尔可夫模型(HMM)离散状态的概率分布。更具体地说,DNN每次在三音素状态的集群上产生概率分布,然后解码器找到通过HMM状态的路径,这是使用高概率状态和产生在语言模型下可能的转录之间的最佳折中。

虽然可以(并且希望)通过在所有可能路径上边缘化来考虑解码器(以及因此考虑语言模型)来训练DNN,训练DNN执行逐帧分类是很常见的,方法是(局部)最小化网络预测和标签之间的交叉熵,这些标签是通过强制对齐每个观测的地面真值状态序列给出的。

我们使用了一个包含8个隐藏层的架构,每个隐藏层包含2560个校正线性单元,最后一个softmax层包含14000个标签(HMM)。输入是40个Mel缩放滤波器组系数的26帧,每帧提前10毫秒,我们预测第21帧的HMM状态。参数总数约为85M。这是安卓语音搜索使用的声学模型的一个稍微过时的版本,应该被视为一个非常强大的基线。为了训练DNN声学模型,我们使用了大约2000小时的英语口语数据,产生了大约7亿个训练示例。该系统在我们的开发集上实现了58.9%的帧精度和10.9%的字错误率(WER)。

我们训练了10个单独的模型来预测,使用与基线完全相同的架构和训练程序。这些模型是用不同的初始参数值随机初始化的,我们发现这在训练过的模型中创造了足够的多样性,以允许平均值对集合的预测将显著优于单个模型。我们已经探索了通过改变每个模型看到的数据集来增加模型的多样性,但我们发现这不会显著改变我们的结果,所以我们选择了更简单的方法。对于蒸馏,我们尝试了[1,2,5,10]的温度,并对硬目标的交叉熵使用了相对权重0.5。

事实上,我们的蒸馏方法能够从训练集中提取比简单地使用硬标签训练单个模型更有用的信息。通过使用10个模型的集成,超过80%的帧分类精度提高被转移到蒸馏模型,这与我们在MNIST的初步实验中观察到的改进类似。由于目标函数不匹配,集成对WER的最终目标(在23K单词测试集上)的改善较小,但集成对WER的改善再次转移到蒸馏模型。

我们最近意识到,通过匹配已经训练过的较大模型的类概率来学习小型声学模型的相关工作。然而,他们使用一个大的未标记数据集在1的温度下进行蒸馏,当他们都使用硬标签进行训练时,他们的最佳蒸馏模型仅将小模型的错误率减少了大模型和小模型错误率之间差距的28%。

基于大数据的专家系统训练

训练模型集合是利用并行计算的一种非常简单的方法,而集合在测试时需要太多计算的常见问题可以通过使用蒸馏来解决。然而,对集成还有另一个重要的反对意见:如果单个模型是大型神经网络,且数据集非常大,则训练时所需的计算量过大,即使很容易并行化。

在本节中,我们给出了这样一个数据集的示例,并展示了学习专家模型如何减少学习集合所需的总计算量,每个模型都关注不同的可混淆的类子集。专注于进行细粒度区分的专家的主要问题是,它们很容易过度拟合,我们描述了如何通过使用软目标来防止这种过度拟合

JFT数据集

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

专家模型

当类的数量非常大时,笨重的模型应该是一个集合,其中包含一个针对所有数据训练的通用模型和许多“专家”模型,每个模型都是针对数据进行训练的,这些数据在非常容易混淆的类子集(如不同类型的蘑菇)的示例中高度丰富。通过将它不关心的所有类合并到一个垃圾箱类中,可以使此类专家的softmax变得更小。

为了减少过度拟合并分担学习低层特征检测器的工作,每个专家模型都用通才模型的权重初始化。然后,通过训练专家,对这些权重进行轻微修改,其中一半样本来自其特殊子集,另一半样本来自训练集的其余部分。训练结束后,我们可以通过增加垃圾箱类的logit来纠正有偏差的训练集,方法是记录专家类的过采样比例。

类分配到专家模型

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

特别是,我们将聚类算法应用于我们的通用模型预测的协方差矩阵,以便经常一起预测的一组类将用作我们的一个专业模型的目标。我们将在线版本的K均值算法应用于协方差矩阵的列,并得到了合理的聚类。我们尝试了几种产生类似结果的聚类算法。

运用专家系统推理

在研究提取专家模型时会发生什么之前,我们想看看包含专家的集合表现如何。除了专家模型,我们总是有一个通用模型,这样我们就可以处理没有专家的类别,这样我们就可以决定使用哪些专家。给定输入图像x,我们分两步进行分类:

第一步:对于每个测试用例,我们根据通用模型找到n个最可能的类。将这组类称为k。在我们的实验中,我们使用n=1。

第二步:然后我们取所有专家模型m,其可混淆类的特殊子集Sm与k有一个非空交集,并将其称为专家活动集Ak(注意,该集合可能为空)。然后,我们找到所有类的全概率分布q

结果

从训练有素的基线完整网络开始,专家模型训练速度极快(JFT只需几天,而不是几周)。此外,所有专家都经过完全独立的培训。有了61个专家模型,总体测试精度相对提高了4.4%。我们还报告了条件测试的准确性,即只考虑属于专家类的示例,并将我们的预测限制在该类子集的准确性。

对于我们的JFT专家实验,我们培训了61名专家模型,每个模型有300个类别(加上垃圾箱类)。因为专家的类集合不是不相交的,所以我们经常有多个专家覆盖一个特定的图像类。当我们有更多的专家覆盖特定的类别时,准确性的提高会更大,这一总体趋势让我们感到鼓舞,因为训练独立的专家模型非常容易并行化。

软目标作为正则项

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

我们在JFT数据集上的实验中使用的专家将他们所有的非专家类合并为一个垃圾箱类。如果我们允许专家在所有课程中使用完整的softmax,可能有一种比提前停止更好的方法来防止他们过度适应。专家在其特殊类中接受高度丰富的数据培训。这意味着它的训练集的有效规模要小得多,而且它在特殊课程上有很强的过度适应倾向。这个问题不能通过让专家变得更小来解决,因为这样我们就失去了建模所有非专家类时得到的非常有用的传递效果。

我们使用3%的语音数据进行的实验强烈表明,如果一个专家是用多面手的权重初始化的,那么除了用硬目标训练之外,我们还可以通过用非特殊类的软目标训练它,使它保留关于非特殊类的几乎所有知识。软目标可以由多面手提供。我们目前正在探索这种方法。

与混合专家系统的关联

使用接受过数据子集培训的专家与混合专家有一些相似之处,后者使用门控网络来计算将每个示例分配给每个专家的概率。在专家学习处理分配给他们的示例的同时,门控网络正在学习根据专家对该示例的相对辨别能力来选择将每个示例分配给哪个专家。使用专家的辨别能力来确定学习的任务,比简单地对输入向量进行聚类并为每个聚类指定一名专家要好得多,但这会使训练难以并行化:首先,每个专家的加权训练集以一种依赖于所有其他专家的方式不断变化,其次,门控网络需要比较同一示例中不同专家的性能,以了解如何修改其分配概率。这些困难意味着,在可能最有益的领域,专家的混合很少被使用:包含明显不同子集的庞大数据集的任务。

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

总结

我们已经证明,蒸馏对于将知识从集合或大型高度正则化模型转移到较小的提取模型非常有效。在MNIST上,即使用于训练蒸馏模型的传递集缺少一个或多个类的示例,蒸馏也能非常好地工作。对于安卓语音搜索使用的深度声学模型,我们已经证明,通过训练一组深度神经网络实现的几乎所有改进都可以提炼成一个同样大小的单一神经网络,部署起来要容易得多。

对于非常大的神经网络,即使训练一个完整的集合也是不可行的,但我们已经证明,通过学习大量的专家网络,一个训练了很长时间的非常大的网络的性能可以显著提高,每个人都学会在一个高度易混淆的集群中区分不同的类。我们还没有证明我们可以将专家们的知识提取回单一的大网络中。