knowledge distillation 论文阅读之：ResKD: Residual-Guided Knowledge Distillation

ABSTRACT

• 知识蒸馏是把知识从复杂的 teacher network 迁移到轻量级的 student network 的有效方法，旨在压缩神经网络。
• 然而由于 teacher 和 student 之间参数量的巨大差异，他们的表现也有非常大的差别
• 本文将通过一个新的视角来审视知识蒸馏，利用 teacher 和 student 之间的知识差距（knowledge gap 或者称作 residual：残差）作为引导，来训练一个更加轻量级的 student 网络，我们把它称为 “res-student”
• 我们把 student 网络和 res-student 网络 结合成一个新的 student 网络；这个 res-student 网络可以修正之前 student 网络中的误差
• 这样的利用残差来修正误差的过程可以被一直重复一直达到 accuracy 和 cost 的平衡为止
• 在整个推理的过程中 （inference time）本文提出了一个 sample-adaptive （样本适应）的策略，用来决定 res-student 网络对哪些样本来说是不必要的（对于简单易识别、区分的样本不需要重复太多次残差学习的过程），这样可以很大程度上减少计算开销
• 实验结果表明，我们和其他网络相比达到一个较为有竞争力的结果时，teacher 网络的计算开销仅为其他网络的 18.4%，23.14%，53.59%，53.80%。
• 我们使用验证我们模型的数据集是：CIFAR-10, CIFAR-100, Tiny-ImageNet, ImageNet，最后，我们对我们的方法进行了深入的理论和实证分析。

INTRODUCTION

• 在本文中，我们按照知识蒸馏的常规步骤，首先根据 teacher 网络 $T$ 训练一个student 网络 $S_0$
• 然后，我们训练一个更加轻量化的网络，取名为“res-student” 来学习 teacher 网络和 student 网络之间的 knowledge gap
• student 网络 $S
• $S_0$ 和 res-student 网络 $R_1$ 结合成为新的 student 网络 $S_1$， $R_1$ 用来负责修正 $S_0$ 的误差
• 同样地步骤，可以获得一个比 $R_1$ 还要轻量级的网络 $R_2$， $R_2$ 被用来学习 teacher 网络 $T$ 和新的学生网络 $S_1$ 之间的 knowledge gap （因为 teacher 和 新的 student $S_1$ 之间的 knowledge gap 相比 teacher 网络和 $S_0$ 之间的 knowledge gap 会减小，所以后面 $R_2$ 肯定比 $R_1$ 还要轻量级）
• 上述的过程可以一直重复，一直产生 $S_n$ 我们把这种 residual-guided 方式的知识蒸馏结构称作 “ResKD”，这种用 res-student 来使 student 逼近 teacher network 的方式有点像多项式逼近。在多项式中，高阶的多项式会有更好的逼近效果，但同时也会增大计算量。同样地，高阶的 ResKD-style 的 student 网络也会更加接近于 teacher，但同样计算成本也会相应提高。
• 在实验中，实验者可以通过控制超参数来调节要产生的 $S_n$ 的数量，从而结束程序运行。
• 另外；上述提到的 knowlege gap 现象在样本之间也同样存在。举个例子来说，给定一个 $S_2$ 的 student 网络，我们观察到：对于一些图片，$S_1(=S_0+R_1)$ ，甚至 $S_0$ 就可以让 student 有很高可信度的表现，也就是说残差的部分 $R_2$ 或者 $R_1+R_2$ 的贡献相应的就很小。也就是说，不是对于所有的样本训练的时候都需要 res-student。
• 因此，我们在训练的时候引入了一个 sample-adaptive （样本适应）的策略。对于每个不同的样本，如果 $S_i$ 的可信度足够高，我们就截断不必要的 res-student 部分来节省训练资源。
• 我们在几个 benchmark 上通过实验，实验结果表明，我们和其他网络相比达到较好性能的时候，teacher 网络在 CIFAR-10, CIFAR-100, Tiny-ImageNet, ImageNet 数据集上做的浮点运算仅分别为其他网络的 18.4%，23.14%，53.59%，53.80%。
• 当然，我们还将 ResKD 结构运用到不同的 KD 方法中，最终我们分析了我们的 idea 的有效性并且使用信息量来可视化了 knowledge gap 被逐渐消除的过程。

本文的贡献如下：

1. 设计了一个 residual-guided 学习方法，使用 res-student 来消除 student 网络和 teacher 网络之间的 knowledge gap
2. 引入了 sample-adaptive 策略（样本适应策略），在训练时使得网络结构去适应不同的样本并以此节省运算资源
3. 评估了我们的方法在不同数据集上的表现，并且展示了这个网络每个部分的重要程度。

RELATED WORK

Knowledge Distillation

KD 通常是个两步走的方法：teacher 网络先别训练出来，然后 student 网络在 teacher 的指引下完成训练。

• Bucil˘ a 等人首创了将知识从一个笨重的模型转移到一个小型模型的想法
• Hinton 通过 KD 的概念使这种研究流行起来：student 网络 从 teacher 网络训练获得的 soft target 中获得引导。相对于传统的 one-hot 刚性编码，从 teacher network 中产出的 soft target 包含了关于数据的细粒度分布的更多信息；从而使 studnet 网络表现更佳。
• 近年来，国内外的研究工作主要集中在改进信息传播方式或通过优化对蒸馏过程进行严格控制等方面，以便于训练表现更好的 student 网络
• 例如，Peng等人提出，学生网络不仅要在实例级模仿 teacher，还要模仿 teacher 的嵌入空间，使学生具有类内紧性和类间可分性
• 此外，还研究了不同 teacher 的影响。例如，Sau等人提出了一种方法，通过向训练数据注入噪声，扰动 teacher 的logit输出，来模拟多名 teacher 的效果；这样，扰动输出不仅模拟了多个 teacher的设置，而且在softmax层中产生了噪声，从而使蒸馏损失得到了正则化；在许多teacher 的帮助下，student 网络表现进步了很多
• Kang等人[30]使用Neural Architecture Search (NAS)从不同的 teacher 网络那里获取关于 student 网络的结构和参数的知识。
• KD除了经典的图像分类任务外，还可以用于很多不同的领域，如人脸识别，视觉答题，视频任务等

最近，一些 KD 方法被提出有少于或多于两个步骤。

• 首先，KD可以是一阶段策略。Zhang等人提出，可以使用一群网络结构相同的未经训练的 student 网络同时学习目标任务，而不是传统的两阶段 KD 策略。
• 另一个研究方向是侧重于两个以上阶段的 KD 方法【36】和 【37】。在[36]中，teacher-assistan 网络被用来更平滑和有效地（softly and effectively）传递来自 teacher 网络的知识。teacher 先向 assistant 网络传播知识，然后 assistant 网络再向 student 网络传播知识。
• 在[37]中，student 网络最初的架构与 teacher 网络相同。student 网络的最后一块被一个具有简单架构的块所取代，并在第一步中训练最后一块。接下来，用类似的方法替换倒数第二个块，然后训练最后两个块。在这种方式下，所有的 block 经过几个阶段的训练，最终得到一个简单的 student 网络

在这篇论文中，我们主要研究如何利用 teacher 网络 和 student 网络之间的 gap 作为知识。我们使用一系列轻量级的网络，名为res-students，以多阶段的方式学习 knowledge gap。

Ensemble Methods

Ensemble methods 集成方法，为了提高模型的性能而得到了广泛的研究。集成方法是一种通过多个模型平均、多数投票或者其他方法的策略。即：多个具有相同状态的模型被用来提升最终的表现。集成方法的结果至少和每一个集成方法中的模型的结果（集成方法表现>=每一个模型的表现）。
现存有几种研究集成方法的研究思路：

• 引入不同的正则项
• 减少训练时间
• 减少测试时间

为了将 KD 和集成方法相结合，Lan等人提出将几个同质 student 模型的 logits 进行聚合，成为一个集成 teacher ，然后从 teacher 那里提取知识。

与之相反的是，我们首先运用 KD 的方法来建立一个 student 网络和一个 res-student 网络，然后将他们进行结合。而且， student 和 res-student 扮演的角色是不同的，student 从 teacher 网络中获得 knowledge， res-student 关注的则是 knowledge gap，而且我们还引入了 sample-adaptive 的训练策略。

RESIDUAL-GUIDED KNOWLEDGE DISTILLATION

这个部分作者提出了 Res-KD 的模型，先介绍了一下背景，然后讲自己做的工作的详细思路；然后提一下如何使用 sample-adaptive 的训练策略。最后把所有这些结合起来形成了一个完整的网络框架。为了方便起见，把下面公式中用到的符号列在这个表里：

background 背景

一个普通的 KD 过程

• 定义一个 teacher network $T$ 如下：
  

  • $x$ 是一个输入的图
  • $W_T$ 是一个 $T$ 的权重
  • ${\alpha }_T$ 是这个网络结构的参数

• 同样的，student 网络用下面的公式代表：
  

• 为了方便表示，我们表示一个网络和这个网络的 logits 层使用相同的表示方法，比如 $S_0$ 既可以表示这个 student 网络，也能表示这个网络的 logits 层的输出

• KD 的目标是要训练 $w_S$ 来使得 student $S$ 的结果尽可能地接近 teacher $T$

传统的 KD 的做法分两步：

<1> 第一步，student 的网络结构 ${\alpha }_S$ 被确定下来，
<2> 第二步，将 $w_S$ 根据下面公式进行优化;

• $x^{(j)}$ 是一张 dataset 中的图片，这个 dataset 中有 $N$ 个训练样本
• $L_{KD}$ 定义如下：
  
  • $y^{(j)}$ 是输入 $x^{(j)}$ 的 label
  • $\sigma （\cdot ）$ 是激活函数 softmax
  • $L_{CE}$ 是传统的 cross-entropy 损失
  • $\tau ,t$ 是超参数标量
  • $L_{T-S}$ 是用于优化 student 网络 logits 输出与 teacher 网络 logits 输出的损失函数（比如在经典的 KD 中，这个损失是 KL 散度）

student 网络通过公式（2） 进行优化 $w_S$ 之后， 根据公式 ：

输出的 logits 的结果可以接近 teacher 网络通过：

获得的 logits 表现。但是我们依然可以观测到 student 网络和 teacher 网络之间的 knowledge gap $\Delta$ 的存在

我们通过下面的研究来减小这个 knowledge gap

Residual-guided Learning

本文提出了一个 residual-guided 形式的网络结构，使用 teacher 和 student 网络之间的 knowledge gap（残差residual）。给定一个 teacher 网络，我们先按照传统的方式训练一个 student 网络；然后我们引入另外一个 student ：res-student 来学习这个残差。而后，结合原始的 student 网络和 res-student 来构建一个新的 student 网络。这个过程可以被一直重复一直到 teacher 和 student 的表现相似为止，整个过程如下图所示：

形式上，设定 teacher 网络的 logits 输出为 $T$， student 网络的 logits 输出为 $S_0$，res-student 网络的 logits 输出为 { R i } i = 1 n \{R_i\}^n_{i=1} {Ri​}i=1n​。我们分阶段训练 $S_0$ 和 { R i } i = 1 n \{R_i\}^n_{i=1} {Ri​}i=1n​

• 在 stage 0：$S_0$ 被 teacher 网络 $T$ 以经典的 KD 方式训练；$T$ 和 $S_0$ 之间的 knowledge gap 可以被表示为：${\Delta }_0=T-S_0$
• 在 stage 1：${\Delta }_0=T-S_0$ 被定义为一个新的 teacher 网络 $T_1={\Delta }_0$，根据这个 teacher 网络来训练 res-student 网络 $R_1$
• 当我们获得 $R_1$，我们可以定义一个新的 student 网络 $S_1=R_1+S_0$
• 重复上面的步骤，在每一个后面的 stage，都可以通过这样的方式来得到每个 stage 的 knowledge gap ${\Delta }_i=T_i-R_i$ ，(当 $i=0$ 时， ${\Delta }_i={\Delta }_0$， $T_i=T$，$R_i=S_0$，所以 ${\Delta }_0=T-S_0$)，这个 knowledge gap 作为下一个阶段的 teacher network $T_{i+1}$，即： $T_{i+1}={\Delta }_i=T_i-R_i$
• 从 $T_{i+1}={\Delta }_i=T_i-R_i$ 这个 teacher network 中产生 res-student $R_{i+1}$；并且获得下一个 student 网络：$S_{i+1}=S_0+R_1+...+R_{i+1}$，最后，使用在所有 $n+1$ 个阶段训练的 student 网络和 res-student 网络
• 将输入图片分别输入到所有的 network 并且对他们的 logits 进行求和最后得到总的 logits（这里用了前面 ensemble 的思想）：
  

NAS-assisted architecture（神经网络结构搜索）

当我们使用 res-student 来缩小和 teacher 表现的差距时，最好是用 NAS 来获得 res-student 的结构，尽量不要通过人工设计的方式（handcrafted）。我们在使用 NAS 在进行神经网络结构搜索的时候同样使用了 KD 的 loss function ；这可以通过在公式（2）中加入 ${\alpha }_S$ 作为可优化参数来体现（公式2中只优化了 $w_S$）。

• 我们使用和 STACNAS[44] 相同的搜索空间
• 一个 STACNAS-style 的神经网络有一系列的 “块” 组成，我们把这些 “块” 称作 “单元（cell）”
• 我们定义了一组候选操作 $O$ 用于这些单元内，并从中为体系结构选择最佳操作。
• 细胞和整个网络的构造可以描述为一个有向无环图(DAG)，如下图所示：
  
  • 一个单元格包含 $N$ 个节点的序列 $f_1,f_2...,f_N$，每一个都是一堆特征图
  • 这些节点由有向边连接
  • 每个边 $(i,j)$ 代表一些将 $f_i$ 变成 $f_j$ 的 操作 $o_k^{(i,j)}\in O$
  • 在一个具体单元 cell 中，每个节点的计算是其所有前面操作过程的加权和:
    
    • ${\alpha }_k^{(i,j)}$ 是在计算 $f_j$ 时作用于 $f_i$ 的操作 $o_k^{(i,j)}\in O$ 的权重；我们使用与STACNAS[44]相同的训练策略来获得最终合适的 res-student 网络。

Termination condition（结束条件）

当一个 ResKD student $S_i$ 的结果已足够接近 teacher 网络 $T$，我们可以让这个 residual-guided 的过程停止。这里我们定义了一个新的概念 $Energy$ 用来帮助衡量 teacher 和 student 网络之间的差距。

对于一个确定的实例 $x^{(j)}$，$S_i$ 网络的 $Energy$ 是：

• $\sigma （\cdot ）$是 softmax 激活函数

对于一个确定的数据集 $D$，$S_i$ 网络的 $Energy$ 为：

通过这种方式，一个网络的 $Energy$ 代表了它在一个 dataset 上的可信度；当 $S_i$ 网络的 $Energy$ 到达 teacher 网络的 90%（举个例子），这个时候，我们可以设置 $n=i$ 并且停止 residual-guided 的训练过程

在实验中，我们在一个验证集 $D_v$ 上计算 $Energy$，$D_v$ 的数据是从训练集 $D$ 中规整地取样出来的。整个算法的步骤如下：

Sample-Adaptive Inference（样本适应性推理）

当我们完成了 residual-guided 的学习过程，我们有 $Sn=S_0+{\sum }_{i=1}^n{R}_i$

• $R_i$ 被设计用来减小 $S_{i-1}$ 和 $T$ 之间的 knowledge gap ${\Delta }_{i-1}$
• 然而对于每一个样本，${\Delta }_{i-1}$ 依然是有区别的；例如，如果样本很容易识别，
  ${∆}_0$ 也可以是非常小的，如果样本比较难，即使 ${∆}_2$ 也可以是相当大的；换句话来说：对于简单的样本，$S_0$ 来逼近 teacher 网络的结果已经很好了，对于复杂的样本，需要用更多的 res-student 的共同效果来逼近 teacher 网络的结果。

基于以上的观察，我们在对每一个样本做 inference 的时候采用了一种 sample-adaptive 的策略：类似于在之前有关于 $Energy$ 的想法，我们统一从训练集 $D$ 中抽取一个验证集 $D_v$，为最终 ResKD student $S_n$ 计算一个 $Energy$ 阈值:

• 当 $S_i$ 对于单个具体的样本 $x^{(j)}$ 的 $Energy$ （可以表示为： E n e r g y ( S i , { x ( j ) } ) Energy(S_{i},\{x^{(j)}\}) Energy(Si​,{x(j)})）高于 $T{H}_{energy}$，我们就把后面的 res-student 项（$R_{i+1},...,R_n$）截断舍去
• 因此，给定样本 $x^{(j)}$，我们定义 $S_n^{(i)}$ 为其在推理时使用的学生网络:
  
  • 对于样本 $x^{(j)}$ 我们用 $S_L$ 而不是 $S_n$
  • 例如，让 $S_n=S_3$ ；给定一个样本 $x^{(j)}$，如果能量 $Energy(S_1),x^{(j)})>T{H}_{energy}$
    这意味着我们只用到了 $S_0$ 和 $R_1$，那么 $L=0$，并且把 $R_2$ 和 $R_3$ 的项截断舍去。

整个 sample-adaptive 的处理过程在下图的算法2 中展示：

The whole Framework

当面对一个 KD 问题时

• 第一步是在 $T$ teacher 的指导下，训练一个 student 网络 $S_0$，它可以手工制作 (handcrafted) 或搜索(NAS)
• 接下来，我们开始使用我们的 residual-guided 学习策略来寻找 res-students；我们以 KD 的方式在 $T_i$ 的引导下训练 res-student $R_i$
• 这样的 residual-guided 过程可以重复，直到 $S_i$ 达到了与 $T$ 相当的 $Energy$。在推理时我们应用我们的样本自适应策略(sample-adaptive strategy)。