作者 | 小小 整理 | NewBeeNLP
知识蒸馏是一种模型压缩方法,是一种基于“教师-学生网络思想”的训练方法,由于其简单,有效,在工业界被广泛应用。这一技术的理论来自于2015年Hinton发表的一篇神作:Distilling the Knowledge in a Neural Network[1]
Knowledge Distillation,简称KD,顾名思义,就是将已经训练好的模型包含的知识(”Knowledge”),蒸馏(“Distill”)提取到另一个模型里面去。今天,我们就来简单读一下这篇论文,力求用简单的语言描述论文作者的主要思想。在本文中,我们将从背景和动机讲起,然后着重介绍“知识蒸馏”的方法,最后我会讨论“温度“这个名词:
虽然在一般情况下,我们不会去区分训练和部署使用的模型,但是训练和部署之间存在着一定的不一致性:
因此,模型压缩(在保证性能的前提下减少模型的参数量)成为了一个重要的问题。而「模型蒸馏」属于模型压缩的一种方法。
插句题外话,我们可以从模型参数量和训练数据量之间的相对关系来理解underfitting和overfitting。AI领域的从业者可能对此已经习以为常,但是为了力求让小白也能读懂本文,还是引用我同事的解释(我印象很深)形象地说明一下:
模型就像一个容器,训练数据中蕴含的知识就像是要装进容器里的水。
上面容器和水的比喻非常经典和贴切,但是会引起一个误解: 人们在直觉上会觉得,要保留相近的知识量,必须保留相近规模的模型。也就是说,一个模型的参数量基本决定了其所能捕获到的数据内蕴含的“知识”的量。
这样的想法是基本正确的,但是需要注意的是:
知识蒸馏使用的是Teacher—Student模型,其中teacher是“知识”的输出者,student是“知识”的接受者。知识蒸馏的过程分为2个阶段:
在本论文中,作者将问题限定在「分类问题」下,或者其他本质上属于分类问题的问题,该类问题的共同点是模型最后会有一个softmax层,其输出值对应了相应类别的概率值。
如果回归机器学习最最基础的理论,我们可以很清楚地意识到一点(而这一点往往在我们深入研究机器学习之后被忽略): 机器学习「最根本的目的」在于训练出在某个问题上泛化能力强的模型。
而现实中,由于我们不可能收集到某问题的所有数据来作为训练数据,并且新数据总是在源源不断的产生,因此我们只能退而求其次,训练目标变成在已有的训练数据集上建模输入和输出之间的关系。由于训练数据集是对真实数据分布情况的采样,训练数据集上的最优解往往会多少偏离真正的最优解(这里的讨论不考虑模型容量)。
而在知识蒸馏时,由于我们已经有了一个泛化能力较强的Net-T,我们在利用Net-T来蒸馏训练Net-S时,可以直接让Net-S去学习Net-T的泛化能力。
一个很直白且高效的迁移泛化能力的方法就是使用softmax层输出的类别的概率来作为“soft target”。
上图: Hard Target 下图: Soft Target
「为什么?」 softmax层的输出,除了正例之外,负标签也带有大量的信息,比如某些负标签对应的概率远远大于其他负标签。而在传统的训练过程(hard target)中,所有负标签都被统一对待。也就是说,KD的训练方式使得每个样本给Net-S带来的信息量大于传统的训练方式。
举个例子来说明一下: 在手写体数字识别任务MNIST中,输出类别有10个。
假设某个输入的“2”更加形似”3”,softmax的输出值中”3”对应的概率为0.1,而其他负标签对应的值都很小,而另一个”2”更加形似”7”,”7”对应的概率为0.1。这两个”2”对应的hard target的值是相同的,但是它们的soft target却是不同的,由此我们可见soft target蕴含着比hard target多的信息。并且soft target分布的熵相对高时,其soft target蕴含的知识就更丰富。
两个”2“的hard target相同而soft target不同
这就解释了为什么通过蒸馏的方法训练出的Net-S相比使用完全相同的模型结构和训练数据只使用hard target的训练方法得到的模型,拥有更好的泛化能力。
先回顾一下原始的softmax函数:
但要是直接使用softmax层的输出值作为soft target, 这又会带来一个问题: 当softmax输出的概率分布熵相对较小时,负标签的值都很接近0,对损失函数的贡献非常小,小到可以忽略不计。因此”温度”这个变量就派上了用场。
下面的公式时加了温度这个变量之后的softmax函数:
训练Net-T的过程很简单,下面详细讲讲第二步:高温蒸馏的过程。高温蒸馏过程的目标函数由distill loss(对应soft target)和student loss(对应hard target)加权得到。示意图如上。
: Net-T的logits
: Net-S的logits
: Net-T的在温度=T下的softmax输出在第i类上的值
: Net-S的在温度=T下的softmax输出在第i类上的值
: 在第i类上的ground truth值,
, 正标签取1,负标签取0.
: 总标签数量
.
, 其中
,
.
, 其中
.
的必要性其实很好理解: Net-T也有一定的错误率,使用ground truth可以有效降低错误被传播给Net-S的可能。打个比方,老师虽然学识远远超过学生,但是他仍然有出错的可能,而这时候如果学生在老师的教授之外,可以同时参考到标准答案,就可以有效地降低被老师偶尔的错误“带偏”的可能性。
「讨论」
另外,
。因此在同时使用soft target和hard target的时候,需要在soft target之前乘上
的系数,这样才能保证soft target和hard target贡献的梯度量基本一致。
直接match logits指的是,直接使用softmax层的输入logits(而不是输出)作为soft targets,需要最小化的目标函数是Net-T和Net-S的logits之间的平方差。
「直接match logits的做法是」
「的情况下的特殊情形。」
由单个case贡献的loss,推算出对应在Net-S每个logit
上的gradient:
当
时,我们使用
来近似
,于是得到
如果再加上logits是零均值的假设
那么上面的公式可以简化成
也就是等价于minimise下面的损失函数
【问题】 我们都知道“蒸馏”需要在高温下进行,那么这个“蒸馏”的温度代表了什么,又是如何选取合适的温度?
随着温度T的增大,概率分布的熵逐渐增大
在回答这个问题之前,先讨论一下「温度T的特点」
时的特例,
时,概率分布比原始更“陡峭”,
时,概率分布比原始更“平缓”。
, 此时softmax的值是平均分布的;(ii)
,此时softmax的值就相当于
,即最大的概率处的值趋近于1,而其他值趋近于0)
携带的信息的倾向
温度的高低改变的是Net-S训练过程中对负标签的关注程度: 温度较低时,对负标签的关注,尤其是那些显著低于平均值的负标签的关注较少;而温度较高时,负标签相关的值会相对增大,Net-S会相对多地关注到负标签。
实际上,负标签中包含一定的信息,尤其是那些值显著「高于」平均值的负标签。但由于Net-T的训练过程决定了负标签部分比较noisy,并且负标签的值越低,其信息就越不可靠。因此温度的选取比较empirical,本质上就是在下面两件事之中取舍:
总的来说,T的选择和Net-S的大小有关,Net-S参数量比较小的时候,相对比较低的温度就可以了(因为参数量小的模型不能capture all knowledge,所以可以适当忽略掉一些负标签的信息)
[1]
Distilling the Knowledge in a Neural Network: https://arxiv.org/abs/1503.02531
[2]
深度压缩之蒸馏模型 - 风雨兼程的文章 - 知乎: https://zhuanlan.zhihu.com/p/24337627>
[3]
知识蒸馏Knowledge Distillation - 船长的文章 - 知乎: https://zhuanlan.zhihu.com/p/83456418
[4]
knowledge-distillation-simplified: https://towardsdatascience.com/knowledge-distillation-simplified-dd4973dbc764
[5]
knowledge_distillation: https://nervanasystems.github.io/distiller/knowledge_distillation.html
- END -