【NeuralPS_2022】Efficient Knowledge Distillation from Model Checkpoints

论文链接：https://arxiv.org/pdf/2210.06458.pdf

代码链接：https://github.com/LeapLabTHU/CheckpointKD

创新点

        训练过程中的中间模型（假设模型训练周期是240$epochs$，中间模型指的是训练周期小于240$epochs$下的网络模型），虽然比完全收敛的模型（经过240$epochs$训练得到的网络模型）准确性低，但是比完全收敛的模型更适合充当教师网络。相同训练轨迹的几个中间模型的弱集合，优于独立训练和完全收敛模型的强集合。本文通过信息瓶颈原理（IB）解释：中间模型的特征拥有更高的关于输入的互信息，从而包含更多的“暗知识”来有效蒸馏，并且进一步提出一种基于最大化任务相关互信息的最优中间教师选择算法。

问题

• 通过实验观察到：训练过程中的非收敛中间模型，通常比最终收敛模型作为更好的教师网络。
• 一个完全收敛的模型往往过于自信，并且可能已经有了非目标类的折叠表示。因此，简单的缩放温度并不能有效的恢复被抑制的信息。
• 有效的$KD$不仅要注重提高教师的表现，还需要重新思考，什么是“黑暗知识”以及如何丰富它。

现象

（a）弱中间模型比完全收敛模型更好的充当教师一职。（b）弱中间模型集合比完全收敛模型集合更好的充当教师一职。

定义

        $I(Y;F)$表示学习特征$F$与目标$Y$之间的互信息。
        $I(X;F)$表示学习特征$F$与输入$X$之间的互信息。

        对于网络本身训练而言：

1. $I(Y;F)$随着网络训练时期的增加单调增加。最大化$I(Y;F)$有助于网络本身训练，以及提升教师网络性能，当对于$KD$（$logit$蒸馏）而言并不总是有必要的。因为基础真实目标$Y$已经包含在$KD$目标函数中。
2. $I(X；F)$在网络训练早期阶段增加，但经过一定时期后逐渐减少。在某种程度上$I(X；F)$可以被看作有效$KD$所需要的一种“暗知识”。

公式

        $T^{full}$表示完全收敛的教师模型。$P_{T^{full}}$表示完全收敛的教师模型经过$softmax$的概率输出。

        $S$表示学生模型。$P_S$表示学生网络经过$softmax$的概率输出。

        $Y_{true}$表示真实标签。

        $KD$蒸馏的损失如下所示，$H$表示交叉熵，$\Gamma$表示温度。

$L_{KD}=\alpha H(Y_{true},P_S)+(1-\alpha )H(P_{T^{full}}^{\Gamma },P_S^{\Gamma })$ (1)

        $T^{\mathrm{int}er}$表示中间教师模型，使用中间教师模型进行$KD$蒸馏时，将公式（1）中$T^{full}$的替换为$T^{\mathrm{int}er}$。

         { T 1 f u l l , T 2 f u l l , . . . T M f u l l } \{ T_1^{full},T_2^{full},...T_M^{full}\} {T1full​,T2full​,...TMfull​}$M(M\ge 2)$表示完整训练模型的集合，它们具有相同的网络结构和训练策略，但初始化参数不同。学生网络学习所有教师模型的平均$softmax$输出。$KD$蒸馏的损失如下所示：

$L_{EKD}=\alpha H(Y_{true},P_S)+(1-\alpha )H(\frac{1}{M}\sum _{i=1}^M{P}_{T_i^{full}}^{\Gamma },P_S^{\Gamma })$(2)

         { T 1 i n t e r , T 2 i n t e r , . . . T M i n t e r } \{ T_1^{{\mathop{\rm int}} er},T_2^{{\mathop{\rm int}} er},...T_M^{{\mathop{\rm int}} er}\} {T1inter​,T2inter​,...TMinter​}$M(M\ge 2)$表示中间模型的集合，使用中间模型的集合进行$KD$蒸馏时，将公式（2）中$T_i^{full}$的替换为$T_i^{\mathrm{int}er}$。

实验

中级教师VS完整教师

        “完整教师VS中级教师”的比较结果”。$T^{full}$代表一个完整的老师。$T^{\mathrm{int}er}$代表一名$KD$表现突出的中级教师。$\ast$代表$KD$表现粗体的最佳中级教师。括号内的数字表示精度的下降$（\downarrow ）$或增加$（\uparrow ）$。

        CIFAR-100采用100$epochs$的中间模型$T^{\mathrm{int}er}$，ImageNet采用60$epochs$的中间模型$T^{\mathrm{int}er}$模型。所有中级教师的培训费用仅为全体教师的一半。

        在CIFAR-100数据集上的训练中间模型$T^{\mathrm{int}er}$的实验。

1. 好老师和坏老师都提高了学生的性能。
2. 所有曲线的峰值都不在最后一点，意味着总有一个中间模型会比完全收敛的模型好。
3. 学生网络的准确率随着教师网络训练的$epochs$的变化先上升后下降。

中间模型集合VS完整模型集合

        “中间模型集合VS完整模型集合”。$E{N}_1$表示$T_1^{full}+T_2^{full}$。$E{N}_2$表示$T_1^{\mathrm{int}er}+T_2^{full}$。$E{N}_3$表示$T_1^{\mathrm{int}er}+T_1^{full}$。集成KD性能的最佳结果是粗体，而集成性能的最佳结果。

1）对于整体性能，$T_1^{full}+T_2^{full}>T_1^{\mathrm{int}er}+T_2^{full}>T_1^{full}+T_1^{\mathrm{int}er}$是一致的。
2）对于蒸馏性能，$T_1^{inter}+T_1^{full}$和$T_1^{inter}>T_2^{full}$具有相似的性能，但明显优于$T_1^{full}+T_2^{full}$。

观察结果
1）中间教师模型的蒸馏性能可能与完全收敛的教师模型相当，甚至更好，但前者的准确率和训练成本较低。
2）一个更强的完整模型集合并不总是能够带来更好的蒸馏效果。中间模型集合比完整模型集合具有更差的集成性能和更低的训练成本，但其有更好的蒸馏性能。

互信息量

        $X$表示输入变量。$F$表示输入变量的任何表示。$Y$表示输出变量。
使用预测精度量化$I(Y;F)$，采用重建损失（在网络最后一个卷积层之后连接一个解码器来重构输入$X$）量化$I(X;F)$。

        信息平面上的互信息曲线。红线代表教师模式。深蓝色的*表示为$T^{full}$，而浅蓝色的表示$T^{inter}$。灰色代表使用交叉熵训练的学生、深蓝色的线代表$T^{full}$蒸馏的学生、浅蓝色的线代表$T^{inter}$蒸馏的学生。黄色区域为观察区域。

现象
        网络与输入的互信息量$I(X;F)$先上升后降低。
        网络与输出的互信息量$I(Y;F)$单调上升。
解释
        对于$I(X;F)$，网络首先吸收输入$X$的信息，然后消除与目标输出$Y$无关的部份信息。
        对于$I(Y;F)$，网络不断积累与目标$Y$相关的信息。

        完整的模型$T^{full}$相较于中间模型$T^{inter}$具有较高的置信度，压缩了更多的$I(X;F)$，但也保留了更多的$I(Y;F)$。较高置信度的模型可能从输入$X$中丢弃了更多的非目标类的信息。保留更多的$I(X;F)$似乎是从$T^{inter}$获得竞争表现的关键因素。

        完整的教师具有大的$I(Y;F_t)$但小的$I(X;F_t)$，等式(5)被重新制定为

$\underset{s}{\min }\{(1+\gamma )I(X;F_s)-(\beta +\gamma )I(Y;F_s)\}$（6）

        中级教师具有大的$I(Y;F_t)$和大的$I(X;F_t)$，等式(5)被重新制定为

$\underset{s}{\min }\{(1-\gamma )I(X;F_s)-(\beta +\gamma )I(Y;F_s)\}$ （7）

        为此，完整的教师模型加速了$I(X;F_s)$的压缩，而中间模型减轻了$I(X;F_s)$的压缩。
        与完整模型相比，一个合适的中间模型具有更多的$I(X;F)$和更少的$I(Y;F)$，这意味着中间模型具有更多的非目标类信息。

标签平滑（LSR）、正则化、知识蒸馏

        标签平滑正则化$（LSR）$可以被认为是一个具有预定义的均匀分布教师的特别$KD$。

$L_{KD}=\alpha H(Y_{true},P_S)+(1-\alpha )H(\mu ,P_S^{\Gamma })$ $(LSR)$

        均匀分布$\mu$有一个小的$I(Y;F_t)$和一个小的$I(X;F_t)$。等式(5)被重新制定为

$\underset{s}{\min }\{(1+\gamma )I(X;F_S)-(\beta -\gamma )I(Y;F_S)\}$ (9)

        其中，$(\beta -\gamma )$必须是正的，以保持优化方向的正确。

        信息平面上的互信息曲线。红色的曲线代表了教师的模型。灰色、深蓝色和浅蓝色的曲线代表学生，正式教师的蒸馏学生，以及中级教师的蒸馏学生。黄色的曲线代表了学生使用LSR进行的训练。黄色区域为观察区域。

        对于蒸馏性能，$T^{\mathrm{int}er}>T^{full}>\mu$。正常的$KD$比标签平滑$（LSR）$向学生模型提供更多的相关信息。