CVPR 2022 | 这个自蒸馏新框架新SOTA，降低了训练成本，无需修改网络

2022 年 4 月 8 日 机器之心

机器之心专栏

机器之心编辑部

OPP O 研究院联合上海交通大学提出的新的自蒸馏框架DLB，无需额外的网络架构修改，对标签噪声具有鲁棒性，并可大幅节约训练的空间复杂度，在三个基准数据集的实验中达到了 SOTA 性能。

深度学习促进人工智能（AI）领域不断发展，实现了许多技术突破。与此同时，如何在有限硬件资源下挖掘模型潜能、提升部署模型的准确率成为了学界和业界的研究热点。其中，知识蒸馏作为一种模型压缩和增强的方法，将泛化能力更强的「大网络模型」蕴含的知识「蒸馏」到「小网络模型」上，来提高小模型精度，广泛地应用于 AI 领域的全监督、半监督、自监督、域迁移等各个方向。

近日， OPPO 研究院联合上海交通大学将视角聚焦到知识蒸馏的范式本身，提出了新的自蒸馏框架：DLB（Self-Distillation from Last Mini-Batch），模型无需额外的网络架构修改，对标签噪声具有鲁棒性，并大幅节约训练的空间复杂度。此外，在三个基准数据的实验中，模型达到了 SOTA 性能。相关论文「Self-Distillation from the Last Mini-Batch for Consistency Regularization」已被 CVPR 2022 收录。

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

DLB 自蒸馏框架

如何减少蒸馏计算复杂度？

知识蒸馏通常可以划分为三类，即离线蒸馏、在线蒸馏和自蒸馏。其中，自蒸馏具有训练轻量、知识迁移效率高的特点，最近受到更多研究者的重视。

图 1：本文方法与其他自蒸馏方法的比较

但是传统的自蒸馏，例如 Be Your Own Teacher，在模型训练过程中需要对模型结构进行修改。除此以外，训练成本高、计算冗余且效率低下也是自蒸馏需要攻克的难题。

为了解决上述难题，让模型更好地部署到手机等终端设备中，OPPO 研究院和上海交通大学的研究员们提出了 DLB 自蒸馏框架。利用训练时前后 Batch 预测结果的一致性，在无需对模型进行网络结构修改的前提下，就能降低训练复杂度，增强模型泛化能力。

1. 本文的任务

提出更加轻量的自蒸馏方式，降低训练的计算复杂度，提高模型准确率和泛化性。

2. 本文创新与贡献

提出 DLB，通过保存与下个 Batch 部分样本重叠的软目标（soft targets）进行自蒸馏。节省计算内存，并且简化训练流程。
让训练样本的每次前向过程都与一次反向传播过程相关联，从而提升学习效率。
实验分析了 DLB 训练方法的动态影响，发现其正则化效果来源于即时生效的平滑标签带来的训练一致性，为自蒸馏的理论研究提供了实验基础。

DLB 自蒸馏框架训练机制

DLB 在训练阶段的每个 iteration 中，目标网络扮演着「教师」和「学生」的双重角色。其中，教师的作用是生成下一个 iteration 的软目标进行正则化；学生的作用是从前一个 iteration 平滑的标签中蒸馏，并最小化监督学习目标。

数据集定义为，包含 n 个样本的 Batch 定义为：，以图像分类为例，图片先进行数据增强，然后将其输入神经网络，优化预测输出与真值间的交叉熵损失：

上式中 p_i 的表达式如下：

θ为网络参数，K 表示分类类别数，τ表示温度。

为了提高泛化能力，传统的 vanilla 知识蒸馏通过额外优化的 KL 散度损失来迁移预先训练好的 teacher 网络的知识，即：

不同于以往采用预先训练 teacher 模型的方式生成(P_i^τ ) ̃，DLB 采用训练中前一个 Batch 蕴含的信息生成(P_i^τ ) ̃，并将其作为正则化的即时平滑标签。

图 2：DLB 训练方式示意图

如图 2 所示，数据样本在第 t 次迭代时的定义为，神经网络的参数为θ_t。

B_t 和 B_(t-1)是通过使用数据采样器获得，在前向过程后计算 L_CE。每个小批次的一半限制为与上一个 iteration 一致，其余一半与下一个 iteration 一致。之后，前半部分小批次利用上一个 iteration 中生成的动态软目标进行学习。即由 t-1 次迭代的软标签（soften labels）生成。因此，引入的正则化损失公式如下：

存储平滑标签只需要很少的额外内存成本，所以带来的额外计算成本很低。整体损失函数由下式表示：

综上，DLB 算法整体的训练的伪代码如下所示：

实验设置

研究员们采用三个图像分类基准数据集评估性能，包括 CIFAR-10、CIFAR-100TinyImageNet。实验结果都达到了最佳性能，如下表所示：

具体而言，在平均错误率层面， DLB 在 CIAFR-100 改进幅度为 0.83% 至 2.50%，在 CIFAR-10 上为 0.37% 至 1.01%，在 TinyImageNet 上为 0.81% 至 3.17。值得一提的是，DLB 的表现明显优于 Tf-KD 和 PS-KD，这证明了 DLB 在模型泛化提升上的优势。

为了评估 DLB 与基于数据增强正则化方法的兼容性，研究员在 CIFAR-10 和 CIFAR-100 上将 DLB 与 CutMix、CutOut 和 DDGSD 相结合。如下所示，实验表明通过结合 DLB 和基于增强的正则化，可以实现额外的性能增益。