【论文笔记_知识蒸馏_2019】Variational Information Distillation for Knowledge Transfer

摘要

将在相同或类似任务上预训练的教师神经网络的知识转移到学生神经网络中，可以显著提高学生神经网络的性能。现有的知识转移方法与教师和学生网络的激活或相应的手工制作的特征相匹配。我们提出了一个知识转移的信息理论框架，该框架将知识转移表述为教师和学生网络之间相互信息的最大化。我们将我们的方法与现有的知识转移方法在知识提炼和转移学习任务上进行了比较，结果表明我们的方法一直优于现有的方法。我们通过在CIFAR-10上将知识从卷积神经网络（CNN）转移到多层感知器（MLP），进一步证明了我们的方法在跨异质网络架构的知识转移上的优势。由此产生的MLP明显优于最先进的方法，它达到了与具有单一卷积层的CNN类似的性能。

1.介绍

…
然而，在许多情况下，不可能收集足够大的数据来训练DNN。例如，在许多医疗图像应用程序[24]中，可用数据的数量受到特定疾病患者数量的限制。

处理这种数据缺乏的一种流行的方法是迁移学习[19]，其目标是从源任务迁移知识，以促进在目标任务上的学习。通常，我们认为源任务一般具有大量的可用数据，其中包含学习目标任务的有用知识，例如，来自自然图像分类[23]的知识可能对细粒度鸟类分类[29]有用。Hinton等人[12]提出了在源任务和目标任务上分别训练的DNN之间传递这些知识的师生框架。高层次的想法是为正在接受目标任务(即学生网络)训练的DNN引入额外的正则化，这允许学习存在于DNN中的知识，这些知识是在源任务(即教师网络)上预先训练的。虽然该框架最初是为相同数据集上的DNN之间的知识转移而设计的，但最近的研究[30,31]开始挖掘其潜力，用于更一般的迁移学习任务，即当源数据和目标数据不同时。

许多知识转移方法已经通过各种各样的直觉知识被提出。Hinton et al. [12] and Ba and Caruana[2]提出匹配教师网络和学生网络的最后一层，因为教师网络最后一层的输出提供了比原始标签更多的信息。Romero等人[22]提出将学生网络的中间层与教师网络的相应层进行匹配。最近的研究[3,6,13,30,31]通过匹配精心设计的从教师和学生网络中间层提取的特征/统计量，放松了匹配整个层的正则化，如注意图[31]和最大均值差异[13]。

显然，知识转移背后并没有一个公认的理论。这就造成了理解实证结果和以更有原则的方式开发新方法的困难。在本文中，我们提出了变分信息蒸馏(VID)作为一个朝着这个方向的尝试，在这个方向上，我们将知识转移表述为教师和学生网络之间相互信息的最大化。该框架为知识转移提出了一个可操作的目标，并允许我们量化从教师网络向学生网络转移的信息量。由于互信息难以计算，我们采用变分信息最大化[1]方案来最大化变分下界。所提出的知识转移方法的概念图如图1所示。我们进一步证明了几种现有的知识转移方法[16,22]可以通过选择不同形式的变分下界作为我们框架的具体实现。我们通过经验验证了VID框架，它显著优于现有的方法。我们观察到，在小数据和异构架构的情况下，这种差距尤其大。

图1:所提知识转移方法的概念图。学生网络在保持与教师网络高互信息的同时，通过最小化交叉熵损耗来有效地学习目标任务。通过学习最大化交互信息，估计激活在教师网络中的分布，激发知识的转移。

总的来说，本文贡献如下：
1.在变分信息最大化的基础上，提出了一种基于变分信息最大化的知识转移框架——变分信息蒸馏VID。
2.我们证明VID概括了几种现有的知识转移方法。此外，在各种知识转移实验中，包括在同一数据集或不同数据集上的(异构)dnn之间的知识转移，我们的框架的实施在经验上优于最先进的知识转移方法。
3.最后，我们证明了在cifar10上卷积神经网络(CNN)和多层感知器(MLP)之间的异构知识转移是可能的。我们的方法生成的学生MLP显著优于文献中报道最好的MLP[17,27]。

2.变分信息蒸馏（VID）

在本节中，我们将VID描述为师生框架中知识转移的一般框架。具体来说，考虑在一个目标任务上训练一个学生神经网络，给另一个教师神经网络在一个类似(或相关)的源任务上预先训练。注意，源任务和目标任务可以是相同的，例如，对于模型压缩或知识蒸馏。基本的假设是，教师网络中的各层经过训练，可以表示源任务和目标任务中都存在的给定输入的某些属性。为了成功的知识转移，学生网络必须学会如何将这些属性的知识纳入到自己的学习中。

从信息论的角度来看，知识转移可以表示为在教师网络和学生网络的各层之间保留高的相互信息。更具体地说，考虑一个来自目标数据分布p(x)的输入随机变量x和k对层：，其中每对层(T^(K)，^S(K))分别从教师网络和学生网络中选取。通过网络将输入x前馈，得到K对随机变量{(t^(K)， s^(K))}^K k=1，表示所选层的激活，如t^(K) = t^(K)(x)。对随机变量(t, s)之间的互信息定义为:
其中熵H(t)和条件熵H(t|s)来自于联合分布p(t, s)。从经验上看，联合分布p(t, s)是输入分布p (x)中采样的输入x在各层上聚集的结果。直观地说，I(t; s)的定义可以理解为当学生层s已知时，教师在其层t中编码的知识的不确定性的减少。我们现在定义以下损失函数，其目的是为目标任务学习一个学生网络，同时鼓励与教师网络的高相互信息。

其中LS为目标任务的任务特异性损失函数，λk >为引入的超参数，用于各层互信息的正则化。方程(2)需要对学生网络的参数进行最小化。但是，由于互信息的精确计算是困难的。相反，我们提出了每个互信息项I(t;S)，其中定义了一个近似于p(t| S)的变分分布q(t| S):(其中一层的互信息的计算)

（由式3可知，H(t)熵越大，在已知学生网络的情况下教师网络的条件熵越小，即互信息I(t,s)越大效果越好，说明学生网络学到的教师知识越多。）
其中期望高于分布p(t, s)，最后一个不等式是由于KL散度DKL(·)的非负性。这种技术被称为变分信息最大化[1]。最后，对(2)中的每个互信息项I(t^(k)， s^(k))应用变分信息最大化，得到VID，得到以下总损失函数的最小值:

目标~L在学生网络参数和变分分布q(t|s)上共同最小化。注意，熵项H(t)已从方程(3)中删除，因为它相对于要优化的参数是常数。另外，我们可以将目标（4）解释为联合训练目标任务的学生网络和最大化条件似然来适应教师网络中选定层的激活。通过这样做，学生网络获得了恢复教师网络中选定层的激活所需的 "压缩 "知识。

2.1算法表达式
我们通过选择变分分布q(t|s)的形式来进一步指定我们的框架。一般情况下，变分分布q(t|s)采用异方差均值µ(·)、同方差σ的高斯分布，即均值µ(·)是s的函数，而标准差σ不是。接下来，μ(·)和σ的参数化由t对应的层的类型进一步指定。当t对应于教师网络的中间层，其空间尺寸分别表示通道、高度和宽度，即t∈R^C×H×W，我们对变量分布的选择表示如下：

其中t(c,h,w)表示以(c, h,w)为指标的t的标量分量。此外，µ(c,h,w)表示由卷积层组成的神经网络µ(·)的单个单元的输出，使用softplus函数保证方差为正，比如。其中αc∈R是待优化的参数，且ε> 0是为数值稳定性引入的最小方差。典型情况下，可以从层次结构和空间维数与t相似的学生网络中选择s。当两层空间维数相等时，通常采用1 × 1卷积层对µ(·)进行有效参数化。否则，可以使用更大的卷积核或转置卷积来匹配空间维数。

另外考虑t = T^(logit)(x)∈R^N层对应教师网络的logit层的情况。这里，我们对变分分布的选择如下:

其中tn表示向量t的第n个入口，µn表示神经网络µ(·)单个单元的输出，σn再次被softplus函数参数化以强制为正。在对于这种情况，学生网络中对应的层s是倒数第二层S^(pen)，而不是logit层，以匹配两层的层次结构，而又不会对学生网络的输出有太大的限制。此外，我们发现，在实践中，使用简单的线性变换对均值函数进行参数化就足够了，即对于某个权重矩阵W，μ(s)=Ws。

在实验过程中，上述实现取得了令人满意的效果。我们也考虑了采用异方差σ(·)，但其训练效果不稳定，改善程度可忽略不计。其他类型的参数化，如重尾分布或混合密度网络[5]可以用来获得额外的性能。我们把这些想法留给未来的探索。

图2:教师网络中中间层空间维度的变分分布对应的热力图，即。每个图对应(a)原始输入图像，(b, c, d)中对数似然函数log q(t h,w|s)归一化并插值来拟合输入图像空间维数(红色像素对应高概率)，(d)中对数似然变分分布优化的学生网络训练没有使用任何知识转移，(f)表示层的大小t在每个空间维度上的平均值。

参见图2说明训练的VID使用实现基于方程(5)。在这里,我们显示的变化评估对数似变分分布聚合的channels,例如,log q (t h, w | s) = ,给定的输入x(图2)整个VID培训过程。可以观察到，将学生网络逐渐训练成变分分布，从教师网络中估计中间层的密度(图2b、2c和2d)。作为比较，我们还优化了未经知识转移训练的学生网络的变分分布(图2e)。对于这种情况，我们观察到这个变分分布的特殊实例没有获得高的对数似然，这表明教师和学生网络之间的互信息很低。有趣的是，与背景相对应的部分总体上比前景的亮度更高。我们的解释是，与背景相对应的层的输出大部分对应于零激活(图2f)，包含的信息较少，相比于最大化变分分布的对数似然是一个相对容易的目标。

2.2现有工作的联系
infomax原则：我们首先描述了我们的框架和应用于表示学习[28]的著名的infomax原则[18]之间的关系，指出“良好的表示”可能在相应的输入中包含大量的信息。特别是，该原理已经成功地应用于神经网络的半监督学习中，将中间层输入输出之间的互信息最大化，作为对目标任务学习的正则化，例如，学习基于自动编码器[21]重构输入。我们的框架可以类似地看作是修改了infomax原则的半监督学习的实例:教师网络的各个层次包含了目标任务的重要信息，而学生网络的良好表示很可能保留了他们的大部分信息。当我们在式(2)中设置t^(k) = x时，恢复了传统的半监督学习infomax原理。

广义均方误差匹配：接下来，我们将解释如何将现有的基于均方误差匹配的知识转移方法视为所提出框架的一个具体实例。一般来说，在变分分布中，特定选择层R = {(T ^(k)， S^(k))}^K k=1进行知识转移和参数化异方差平均µ(·)，从式(4)导出方法:

注意式(7)对应于教师网络中该层各维上单位方差的高斯分布。Ba和Caruana[2]表明，通过将教师和学生网络的logit层T ^(logit)、S^(logit)的输出相对于均方误差进行匹配，教师和学生网络可以在为同一任务而设计的网络之间传递知识。由式(7)中R = {(T ^(logit)， S^(logit))}，µ(S) = S推导出了这样的公式。之后，Li和Hoiem[16]将这种公式推广到针对不同任务设计的教师和学生网络之间的知识转移。通过在学生网络的倒数第二层S^(pen)之上增加一个额外的线性层来匹配教师网络的logit层T ^(logit)。同理，由式(7)，R = {(T ^(logit)， S^(pen))}，µ(·)为线性变换，即µ(S) = w。接下来，Romero等[22]提出了一种知识转移损失最小化教师和学生网络中间层间均方误差的方法，并引入了卷积层以适应每对匹配层之间不同的维度大小。这是通过选择知识转移的层作为教师和学生网络的中间层，µ(·)为单个1 × 1卷积层对应的线性变换，由式(7)中的正则化项得到的。

这些方法都类似于我们对框架的实现，因为它们都使用高斯分布作为变分分布。然而，我们的方法在两个关键方面有所不同。(a) 允许使用更灵活的非线性函数来处理异方差平均数；(b) 在变异分布中为每个维度建立不同的变异模型。这允许在不浪费模型容量的情况下以更灵活的方式传递相互信息。特别是，为教师网络中的层t的所有维度建立单位方差模型，对学生网络来说可能是高度限制性的。举例来说，教师网络的层可能包括一个激活tn，它包含与学生网络的任务无关的信息，但却需要很多容量来对µn(s)进行回归tn。这将引起过度规则化的问题，也就是说，将学生网络的大部分容量浪费在试图适应这样一个单元上。相反，对这种维度的高同调方差σn建模，使其对整体损失的贡献可以忽略不计，允许人们以有效的方式 "过滤 "掉这种单元。

与特征匹配相比：除了基于均方误差匹配的知识转移方法外，有几篇著作[6,13,30,31]提出了从中间层中提取的手工特征的间接匹配。更具体地说，Zagoruyko和Komodakis[31]提出匹配由层面激活产生的“注意力图”。Huang和Wang[13]后来将注意图推广到匹配激活的最大平均差异。Yim等人[30]提出匹配同一网络中相邻层的Gram矩阵定义的称为求解过程流(Flow of Solution Procedure, FSP)的特征。Chen等人[6]考虑了对教师和学生网络中间层的输入图像进行匹配。这些方法可以看作是巧妙地避免了上述的过度正则化问题，利用专家知识过滤了教师网络中的信息。然而，当特征提取方法不适合特定的知识转移任务时，这些方法可能会导致次优结果，并可能以不可逆的方式丢弃教师网络层的重要信息。

3.实验

我们通过将VID与目前最先进的图像分类知识转移方法进行比较，来证明所提出的知识转移框架的性能。我们将VID应用于两个不同的位置:(a)教师和学生网络中间层(VID- I)之间的VID， (b)教师网络的logit层和学生网络的倒数第二层(VID- LP)之间的VID。为了进行比较，我们考虑了以下知识转移方法:原始知识蒸馏(KD)[12]、不遗忘学习(LwF)[16]、基于提示的转移(FitNet)[31]、基于激活的注意转移(AT)[31]和基于多项式核的神经选择性转移(NST)[13]。请注意，我们将FitNet作为训练学生网络[31]的正则化方法，而不是[22]中首次提出的分期训练过程。我们比较了相同数据集和不同数据集之间知识转移的方法，通常分别称为知识蒸馏和转移学习任务。

在所有的实验中，我们基于VID-I、FitNet、AT和NST选择了相同对的中间层进行知识转移。同样，LwF和VID-LP也使用了相同的知识转移层对。所有方法的所有超参数都是根据验证集中的性能选择的，即20%的训练集。我们仔细地选取超参数的候选值集，这样就包含了原著中提出的所有值。这次演出是三次重复演出的平均值。关于实验的更多细节包括在补充材料中。该算法的实现将在短期内公开。

表2:CIFAR-100数据集上从教师网络(WRN-40-2)到不同深度d、宽度w因子的学生网络(WRN-d-w)知识蒸馏的实验结果(测试精度)。

3.1知识蒸馏
我们首先比较了传统知识蒸馏任务下的知识转移方法，在传统知识蒸馏任务下，学生网络与教师网络在同一任务下进行训练。通过将大型教师网络中的知识提炼成小型学生网络，我们可以加快预测的计算速度。为此，我们进一步研究了两个问题:在小数据环境下，我们是否能从知识转移中获益?通过减少学生网络的规模，我们会损失多少性能?注意，我们没有评估VID-LP和LwF的性能，因为它们是为迁移学习设计的。当应用KD、VIDLP和LwF时，具有相似的性能。

减少训练数据：知识转移是一项耗费大量计算的任务。在整个培训数据集上建立一个培训前的教师网络，我们探讨了使用一小部分培训集进行知识转移的可能性。我们通过在CIFAR-10[15]上应用知识蒸馏，用四种不同大小的训练数据来证明简化训练集的效果。我们对教师网络(WRN-40-2)和学生网络(WRN-16-1)采用宽残差网络(WRN)[15]，教师网络在CIFAR-10的整个训练集上进行预训练。知识蒸馏应用于4个不同规模的训练集:5000(全尺寸)，1000,500,100个数据点每类。

我们比较了VID-I与KD, FitNet, A T和NST。我们还提供教师网络(教师)和没有任何知识转移(学生)培训的学生网络的性能作为基准。我们选择了四对类似于[31]的中间层，每对中间层都位于一组残块的末端。我们使用三个1 × 1的卷积层实现VID-I，其隐藏通道大小为输出通道大小的两倍。结果如表1所示。我们的方法VID-I，在所有体系中始终优于其他知识转移方法。数据集越小，表现差距越大，例如，即使每个class提供100个数据点给学生网络，VID-I的准确率也只下降10.26%。在没有知识转移的情况下下降了31.88%，最佳基线下降了15.52%，即KD + AT。

改变学生网络的规模：学生网络的规模在知识转移的速度和性能之间提供了一种权衡。我们评估了知识转移方法在不同规模的学生网络上的表现。教师网络(WRN-40-2)在CIFAR-100的整个培训集上进行了预培训。在CIFAR-100的整个训练集上训练一个具有WRN-402、WRN-16-2、WRN-40-1、WRN-16-1四种规模选择的学生网络。我们将VIDI与KD, FitNet, A T和NST以及教师和学生基线进行比较。中间层的选择与之前的实验相同。

结果如表1所示。Furlanello等[9]也注意到，在所有知识转移方法下，与教师网络相同规模的学生网络的表现都优于教师网络。有人观察到，VID-I始终优于FitNet、AT和NST，它们对应于相同的知识转移层的选择。除了在学生网络结构与教师网络结构相同的情况(即WRN-40-2)，两种方法结合可以得到最好的性能外，KD的性能也优于KD。

3.2迁移学习

3.3从CNN到MLP的知识转移

4.结论

在本工作中，我们提出了一个VID框架，通过最大化两个神经网络间互信息的变分下界来实现有效的知识转移。我们的算法的实现是基于高斯观测模型，并被经验证明在蒸馏和迁移学习任务优于其他基准。使用更灵活的识别模型，如[14]，以精确地最大化互信息，以及互信息的可选估计，如[4]，都是未来感兴趣的想法。