文章总结:Distilling the Knowledge in a Neural Network(蒸馏神经网络)
<font size=3.2>
Abstract:
在机器学习领域,ensemble learning是一种普遍适用的用来提升模型表现的方法, 将通过ensemble learning训练出的模型称为cubersome model, 但在这种情况下,模型通常很大(即参数较多,复杂度较高),以至于在模型部署阶段所需要的计算资源极为昂贵,尤其是对于一些以大型神经网络作为子学习器的集成模型。
其中一种潜在的解决方法是,将cubersome model中的信息(knowledge)压缩到一个单独的模型(single model),将此过程称为distilling(蒸馏)。
1 Introduction
对于目标检测或语音识别等任务,模型训练所花费的大部分时间及计算资源是可以接受的,但当模型部署到生产环境中,对模型所需资源的要求则会严格很多。可以通过distilling(蒸馏)来提取cubersome model中的knowledge从而使模型更易部署。对于实现distilling(蒸馏)的阻碍之一在于,knowledge是一个比较抽象的概念,可以认为cubersome model的knowledge体现或者隐藏在cubersome model中的模型参数中,以至于很难在改变模型结构或参数的同时保留knowledge。
在cubersome model学习处理多分类任务时,普遍使用最大似然函数(最大化对数概率log probability)来作为目标函数, 这么做的一个副作用是,当使用softmax作为神经网络的输出层时,模型会赋值给一些非正确的类别(的概率),即使这些值很小,但在这些非正确分类的类别的概率值中,其中一些会相对更显著,这些relative probalities of incorrect answers是很重要的信息,因为它隐含了cubersome model如何泛化generalize的信息(how the cumbersome model tends to generalize)。
通常,object function应尽量反映任务的真实目标。若希望模型具有较好的泛化generalize能力,则需要information about the correct way to generalize, 通常此信息是unavailable的。我们可以用使cubersome model获得较强泛化能力的相同的方式来训练small model,即集成一些small model,尽管这种方式训练出的模型比单独的small model要强,但相比cubersome model很难保证性能接近。
一种直观的转移cubersome model泛化能力的方法(transfer the generalization ability)是,将cubersome models生成的类别概率向量(class probilities produces by the cubersome model)作为训练small model的soft targets,此时可以使用原始的training set或者transfer set。当soft targets具有较高的熵(entropy)时,其相对于hard target可以对每种分类场景提供更大的信息量,这使得在训练small model的时候可以使用更小的数据量、更大的学习率、(以及也许更简单的结构)。
以MNIST数据集为例,档分类的真实值为2时,hard target为(0, 1, 0, ... , 0),模型的分类结果可能在数字3和7所对应的类别向量的位置的值分别为pow(10, -6)和pow(10, -9),这表示通常数字2的形态,相对于数字7,更接近数字3, 这是valuable information,它揭示了关于样本数据的相似结构(相似度)的信息,在此处,可以理解为此信息可以揭示哪些数字2的样本更像数字3,哪些数字2的样本更像数字7。但问题在于,非正确类别在类别向量中所对应的值为0或一个很小的实数对transfer stage时的cross-entropy的值影响很小。关于此问题的一个解法为,使用logits(the inputs to the final softmax)而不是类别向量probabilities vector(probabilities produced by the softmax)作为训练small model时的targets。训练时的object是最小化cubersome model的logits和small model的logits之间MSE。一个更通用的解法,叫做distillation,它会用到<font color=green>temperature of the final softmax</font>,通过调整<font color=green>temperature</font>,可以使cubersome model产生一个<font color=blue>suitably soft set of targets</font>,此target可以在cross entropy的值的层面上对于one-hot vector和probabilities vector有一个较好的区分,从而使small model得以继续训练。
用来训练small model的<font color=green>transfer set</font>可以由unlabeled data或original training set组成。在训练small model时,促使small model尽可能进行正确分类的同时,让其logits尽可能接近通过cubersome model获得的soft targets。
2 Distillation
在神经网络模型中,<font color=green>softmax function</font>将logits, $z_i$转化为<font color=green>class probabilities vector</font>, 记为($q_1$, $q_2$, ... , $q_n$)。每类所对应的概率记为$q_i$, 计算方法为:
$q_i = exp(z_i/T) / (\sum_{j}exp(z_j/T))$
其中$T$为<font color=green>Temperature</font>,常用取值为1。当使用更大的T时,会生成<font color=blue>softer probability distribution over classes</font>。
最简单的distillation的形式是,使cubersome model以较大的temperature值计算probability distribution,并以此作为transfer set中的soft target来训练small model,通过这种方式来转移cubersome model中的knowledge到<font color=green>distilled model</font>(small model),在训练distilled model时需要使用相同的temperature值,<font color=blue>但在完成训练之后,应在预测阶段使用tmeperature=1</font>。
在将correct labels加入transfer set后可以进一步提升distilled model的表现。此时,训练时的objective function实质两部分的加权平均,第一部分为distilled model和cubersome model在soft targets(使用较高temperature计算的probability distribution vector)的交叉熵(cross entropy),第二部分为二者关于correct label(hard targets)的交叉熵。
通常,把objective function的第二部分权重设置的相对较低可以得到更好的结果,这是因为,上文提到的的objective function的第一部分在求梯度之后会带有$1/T^2$项,此时objective function的第一部分与第二部分在量级(magnitudes)上是不同的,故首先要确保时objective function的第一部分与第二部分对objective function的相对贡献(relative contribution)几乎不变(即控制此变量),方式是首先用$T^2$乘以objective function的第一部分,然后计算两项的加权平均值。
2.1 Matching logits is a special case of distillation
在transfer set中,每个case(unit)都会贡献在某一方向的梯度$dC/dz_i$,对应于distilled model的每个logit,$z_i$。若cubersome model的logits,$v_i$对应其计算的soft target probilities $p_i$(的分子部分),且使transfer learning时temperature的值为$T$,此时对应$z_i$的梯度为:
$\partial C / \partial z_i = (q_i - p_i) / T = (e^{z_i/T}/\sum_je^{z_j/T} - e^{v_i/T}/\sum_je^{v_j/T})/T$
当temperature的值$T$相对logits的量级较大时,由$e^x$的泰勒展开可得:
$\partial C / \partial z_i \approx ((1+z_i/T)/(N+\sum_j{z_j/T}) - (1+v_i/T)/(N+\sum_j{v_j/T}))$
若假设logits对于每个transfer case是零均值的(即假设会对logits所对应的layer做normalization),即$\sum_j{z_j} = \sum_j{v_j} = 0$,此次上式变形为:
$\partial C / \partial z_i \approx (z_i - v_i) / (NT^2)$
因此,当限定temperature取较大的值时且logits layer为0均值时,distillation在logit,$z_i$上可以等价为最小化$1 / 2(z_i - v_i)^2$(此时可最大化$C$在$z_i$上的梯度)。
当temperature的值相对较小时,梯度会相对大,这是很好的现象,因为即使$z_i$与$v_i$很接近,依然可以得到相对显著的梯度,但同时distillation将不会重点关注logits间的matching问题。另一方面,当logit很负(very negative)时,计算probability distribution vector时其所对应项的分子,$exp(z_i/T)$会很小,这也是对distillation很有价值的信息,若distillation将不会重点关注logits,则此部分信息很容易丢失。因此,在temperature过小时distillation momdel将无法捕获cubersome model中所有的知识。
综上,temperature的值较大时有助于帮助distillation model捕获cubersome model中更多的知识,temperature的值较小时可以使logits对应的梯度相对显著,因子temperature应取intermediate value。
3 Preliminary experiments on MNIST
实验中使用的是两个隐含层的神经网络模型,每层1200个结点,激活函数为RELU,训练样本空间为60000,使用dropout和weight-constraints正则化。在验证集上,这个网络有67个分类误差,小一些的且没做dropout的网络(800个结点)有146个分类误差。在采用distilling的方式训练small model后,small model在验证集上的分类误差数降低到74个。当进一步简化small model时(降低到300个节点), 仍然可以得到相似的表现。
</font>