1. 首页
  2. 教育
  3. 深度学习

深度学习:知识蒸馏的全过程

知识蒸馏的核心思想是通过迁移知识,从而通过训练好的大模型得到更加适合推理的小模型。本文作者介绍了知识蒸馏的全过程,以及引用Hinton等人的实验结果解释说明,并提出了知识蒸馏的一些延伸工作方向。

知识蒸馏的核心思想是通过迁移知识,从而通过训练好的大模型得到更加适合推理的小模型。本文作者介绍了知识蒸馏的全过程,以及引用Hinton等人的实验结果解释说明,并提出了知识蒸馏的一些延伸工作方向。

0.写在前面

有人说过:“神经网络用剩的logits不要扔,沾上鸡蛋液,裹上面包糠…”这两天对知识蒸馏(KnowledgeDistillation)萌生了一点兴趣,正好写一篇文章分享一下。这篇文章姑且算是一篇小科普。

1.从模型压缩开始

各种模型算法,最终目的都是要为某个应用服务。在买卖中,我们需要控制收入和支出。类似地,在工业应用中,除了要求模型要有好的预测(收入)以外,往往还希望它的「支出」要足够小。具体来说,我们一般希望部署到应用中的模型使用较少的计算资源(存储空间、计算单元等),产生较低的时延。

深度学习:知识蒸馏的全过程

深度学习的背景下,为了达到更好的预测,常常会有两种方案:1.使用过参数化的深度神经网络,这类网络学习能力非常强,因此往往加上一定的正则化策略(如dropout);2.集成模型(ensemble),将许多弱的模型集成起来,往往可以实现较好的预测。这两种方案无疑都有较大的「支出」,需要的计算量和计算资源很大,对部署非常不利。这也就是模型压缩的动机:我们希望有一个规模较小的模型,能达到和大模型一样或相当的结果。当然,从头训练一个小模型,从经验上看是很难达到上述效果的,也许我们能先训练一个大而强的模型,然后将其包含的知识转移给小的模型呢?如何做到呢?

*下文统一将要训练的小模型称为新模型,将以及训练的大模型称为原模型。

RichCaruana等人在[1]中指出,可以让新模型近似(approximate)原模型(模型即函数)。注意到,在机器学习中,我们常常假定输入到输出有一个潜在的函数关系,这个函数是未知的:从头学习一个新模型就是从有限的数据中近似一个未知的函数。如果让新模型近似原模型,因为原模型的函数是已知的,我们可以使用很多非训练集内的伪数据来训练新模型,这显然要更可行。

这样,原来我们需要让新模型的softmax分布与真实标签匹配,现在只需要让新模型与原模型在给定输入下的softmax分布匹配了。直观来看,后者比前者具有这样一个优势:经过训练后的原模型,其softmax分布包含有一定的知识——真实标签只能告诉我们,某个图像样本是一辆宝马,不是一辆垃圾车,也不是一颗萝卜;而经过训练的softmax可能会告诉我们,它最可能是一辆宝马,不大可能是一辆垃圾车,但绝不可能是一颗萝卜[2]。

2.为什么叫「蒸馏」?

接续前面的讨论,我们的目标是让新模型与原模型的softmax输出的分布充分接近。直接这样做是有问题的:在一般的softmax函数中,自然指数先拉大logits之间的差距,然后作归一化,最终得到的分布是一个argmax的近似(参考我之前的文章:浅谈Softmax函数),其输出是一个接近one-hot的向量,其中一个值很大,其他的都很小。这种情况下,前面说到的「可能是垃圾车,但绝不是萝卜」这种知识的体现是非常有限的。相较类似one-hot这样的硬性输出,我们更希望输出更「软」一些。

一种方法是直接比较logits来避免这个问题。具体地,对于每一条数据,记原模型产生的某个logits是,新模型产生的logits是,我们需要最小化

文献[2]提出了更通用的一种做法。考虑一个广义的softmax函数

其中是温度,这是从统计力学中的玻尔兹曼分布中借用的概念。容易证明,当温度趋向于0时,softmax输出将收敛为一个one-hot向量(证明可以参考我之前的文章:浅谈Softmax函数,将替换为即可);温度趋向于无穷时,softmax的输出则更「软」。因此,在训练新模型的时候,可以使用较高的使得softmax产生的分布足够软,这时让新模型的softmax输出近似原模型;在训练结束以后再使用正常的温度来预测。具体地,在训练时我们需要最小化两个分布的交叉熵(Cross-entropy),记新模型利用公式产生的分布是,原模型产生的分布是,则我们需要最小化

在化学中,蒸馏是一个有效的分离沸点不同的组分的方法,大致步骤是先升温使低沸点的组分汽化,然后降温冷凝,达到分离出目标物质的目的。在前面提到的这个过程中,我们先让温度升高,然后在测试阶段恢复「低温」,从而将原模型中的知识提取出来,因此将其称为是蒸馏,实在是妙。

当然,如果转移时使用的是有标签的数据,那么也可以将标签与新模型softmax分布的交叉熵加入到损失函数中去。这里需要将式乘上一个,这是为了让损失函数的两项的梯度大致在一个数量级上(参考公式),实验表明这将大大改善新模型的表现(考虑到加入了更多的监督信号)。

3.与直接优化logits差异相比

由公式,对于交叉熵损失来说,其对于新模型的某个logit的梯度是

由于与是等价无穷小(时**)**,易知,当充分大时,有

假设所有logits对每个样本都是零均值化的,即,则有

所以,如果:1.非常大,2.logits对所有样本都是零均值化的,则知识蒸馏和最小化logits的平方差(公式)是等价的(因为梯度大致是同一个形式)。实验表明,温度不能取太大,而应该使用某个适中的值,这表明忽略极负的logits对新模型的表现很有帮助(较低的温度产生的分布比较「硬」,倾向于忽略logits中极小的负值)。

4.实验与结论

Hinton等人做了三组实验,其中两组都验证了知识蒸馏方法的有效性。在MNIST数据集上的实验表明,即便有部分类别的样本缺失,新模型也可以表现得很不错,只需要修改相应的偏置项,就可以与原模型表现相当。在语音任务的实验也表明,蒸馏得到的模型比从头训练的模型捕捉了更多数据集中的有效信息,表现仅比集成模型低了0.3个百分点。总体来说知识蒸馏是一个简单而有效的模型压缩/训练方法。这大体上是因为原模型的softmax提供了比one-hot标签更多的监督信号[3]。

知识蒸馏在后续也有很多延伸工作。在NLP方面比较有名的有YoonKim等人的Sequence-LevelKnowledgeDistillation等。总的来说,对一些比较臃肿、不便部署的模型,可以将其「知识」转移到小的模型上。比如,在机器翻译中,一般的模型需要有较大的容量(capacity)才可能获得较好的结果;现在非常流行的BERT及其变种,规模都非常大;更不用提,一些情形下我们需要将这些本身就很大的深度模型集成为一个ensemble,这时候,可以用知识蒸馏压缩出一个较小的、「便宜」的模型。

文章地址:https://arxiv.org/abs/1606.07947

另外,在多任务的情境下,使用一般的策略训练一个多任务模型,可能达不到比单任务更好的效果,文献[3]探索了使用知识蒸馏,利用单任务的模型来指导训练多任务模型的方法,很值得参考。

补充

鉴于评论区有知友对公式有疑问,简单补充一下这里梯度的推导(其实就是交叉熵损失对softmax输入的梯度,LOL)。

*这部分有一点繁琐,能接受公式的读者可以跳过。

由链式法则,有

注意到是原模型产生的softmax输出,与无关。

后一项比较容易得到,因为,所以

则是一个维向量

前一项是一个的方阵,分类讨论可以得到。参考公式,记,由除法的求导法则,输出元素对输入的偏导是

注意上面右侧加方框部分,可以进一步展开

这样,代入公式,并且将括号展开,可以得到

左侧方框内偏导可以分类讨论得到

带入式,得到

所以形式如下

代入式,可得

所以有公式,。

参考

[1]Caruanaetal.,ModelCompression,2006

[2]Hintonetal.,DistillingtheKnowledgeinaNeuralNetwork,2015

[3]KevinClarketal.,BAM!Born-AgainMulti-TaskNetworksforNaturalLanguageUnderstanding

责任编辑:xj

原文标题:知识蒸馏是什么?一份入门随笔

文章出处:【微信公众号:深度学习自然语言处理】欢迎添加关注!文章转载请注明出处。

免责声明:本站旨在传递信息,不代表有传资讯的观点和立场。本站遵循行业规范,如转载您的文章未标注版权,请联系我们(QQ:78799268)改正。本站的原创文章,如若转载,请注明出处:http://www.ainoline.cn/jiaoyu/shenduxuexi/15944.html

发表评论

电子邮件地址不会被公开。 必填项已用*标注

联系我们

在线咨询:点击这里给我发消息

邮件:78799268@qq.com

工作时间:周一至周五,9:30-18:30,节假日休息

QR code