无需多个模型也能实现知识整合？港中文MMLab提出「烘焙」算法，全面提升ImageNet性能-轻识

视学算法专栏

转载自：机器之心

作者：葛艺潇

来自港中文 MMLab 的研究者提出一种烘焙（BAKE）算法，为知识蒸馏中的知识整合提供了一个全新的思路，打破了固有的多模型整合的样式，创新地提出并尝试了样本间的知识整合。

知识蒸馏算法方向的研究发现：整合多个模型可以生成更准确的训练监督，但需要以额外的模型参数和明显增加的计算成本为代价。为此，港中文 MMLab 提出了一种新的「烘焙」算法，有效整合同批次内不同样本间的知识以优化软标签，仅需一个网络即可实现知识整合。该「烘焙」方法在任意网络架构的训练中即插即用，以最少的时间成本实现有效的大幅性能提升，研究者在 ImageNet 等多个常见的图像分类基准上进行了算法验证。

论文链接：https://arxiv.org/pdf/2104.13298.pdf
项目主页：https://geyixiao.com/projects/bake
代码链接：https://github.com/yxgeee/BAKE

图像分类是计算机视觉最基本的任务之一，为众多下游视觉任务提供了有效的预训练模型。深度学习领域有大量的算法致力于提升图像分类的性能，尤其是在最流行的 ImageNet 基准上。近期一些研究指出，不够完善的人为标注成为了阻碍监督训练的分类模型性能进一步提升的关键问题。具体来说，人为标注的单标签（一图一类）无法准确描述图像中的丰富内容。

为此，许多研究者提出利用知识蒸馏（Knowledge Distillation）算法自动生成 “软性” 多标签作为有效的训练监督，弥补单标签系统带来的缺陷。一些研究发现，整合多个教师（图 1(a)）或学生模型（图 1(b)）的预测可以生成更鲁棒的训练监督，进一步提升模型的性能，这类算法被称为整合蒸馏（Ensemble Distillation）算法 [1,2,3]。虽然这类算法实现了 SOTA 性能，但它们依赖于额外的网络模型和参数，无疑大大增加了训练时计算和显存的成本。

图 1

除了知识蒸馏系列算法，还有一类算法被称为标签精炼（Label Refinery）算法 [4]。他们往往利用一个预训练标注器为 ImageNet 进行重新标注，该标注器一般为在更大规模的数据集上训练的较深的网络模型，这不仅需要增加额外的模型，还依赖于额外的超大规模数据集和更多的训练资源，在实际应用中不够灵活。

港中文 MMLab 的研究者提出了一种新型「烘焙（BAKE）」训练机制，如图 1(c)所示），该方法整合批次内不同样本间的知识在线优化蒸馏目标，即将同一批次中样本知识进行加权传播和汇总，而无需多个网络。BAKE 首次实现了自蒸馏框架下的知识整合，以最少的训练时间开销和零额外网络参数为代价，持续有效地提高了各种网络架构和数据集的分类性能。例如，使用 BAKE 训练的 ResNet-50 在 ImageNet 上的 top-1 分类准确率显著提升了 1.2%，而相比基线模型训练所增加的计算开销仅为 3.7％。

该论文提出的 BAKE 算法与已有的自蒸馏、整合蒸馏、标签精炼算法的主要区别如下表所示：

方法

传统整合蒸馏算法 [1,2,3] 往往整合的是多个网络模型对于单个样本的预测，与之不同的是，BAKE 整合单网络模型对于批次内多样本的预测。直观来看，视觉上相似的样本应当具有近似的预测。基于该假设，BAKE 依据同批次内其他样本对锚样本的相似度，进行知识的加权传播和汇总，以形成准确的软标签，作为蒸馏训练的监督，如下图所示。

样本间的知识传播基于不同样本与锚样本之间的相似性，所以首先需要计算一个亲和度矩阵 A，也就是计算图像编码器（encoder）输出的特征 {f} 之间的距离。亲和度矩阵需要去除对角线

，也就是同一样本的相似度，并在每行进行 softmax 归一化，使得每一行的和为 1，即

。基于亲和度矩阵 A，可以对除锚样本之外的其他样本的预测进行加权传播，

。并与锚样本本身的预测概率进行加权和，从而获得软标签作为蒸馏目标

。至此，对批次内样本间知识进行了一次传播，并获得了一次传播后的软标签。

往往基于亲和度矩阵对样本预测做多次传播可以获得更鲁棒准确的软标签，

。研究者利用近似预测对传播无限次后的软标签进行了估计，

。

基于上述知识整合后的软标签 Q，可以使用 KL 散度损失函数进行蒸馏训练。

训练的伪代码如下，具体实现请参阅 GitHub repo。

实验

BAKE 以最小的计算开销改进了多种网络架构的训练，并且无需额外的网络辅助。下表给出了在 ImageNet 上的 top-1 分类准确率，「Vanilla」表示使用常规交叉熵损失的基准训练。

BAKE 不光有效提升了基准模型的训练，也超越了所有单网络下最先进的训练机制，如下表所示。

研究者还对 BAKE 生成的软标签进行了可视化，如下图所示。其中，每列的四个样本来自于同一批次，请注意为了简洁这里只对 top-3 的类别进行了展示。

研究者还检验了利用 BAKE 训练的分类模型在下游任务中的表现，发现在目标检测和实例分割中均可获得稳定的性能提升：

更多性能分析、鲁棒性测试、消融研究结果请参阅原论文。

[1] Zhiqiang Shen, Zhankui He, and Xiangyang Xue. Meal: Multi-model ensemble via adversarial learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 4886–4893, 2019.

[2] Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive representation distillation. In International Conference on Learning Representations, 2020.

[3] Qiushan Guo, Xinjiang Wang, Yichao Wu, Zhipeng Yu, Ding Liang, Xiaolin Hu, and Ping Luo. Online knowledge distillation via collaborative learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11020–11029, 2020.

[4] Sangdoo Yun, Seong Joon Oh, Byeongho Heo, Dongyoon Han, Junsuk Choe, and Sanghyuk Chun. Re-labeling imagenet: from single to multi-labels, from global to localized labels. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.

转载请联系原公众号获得授权

投稿或寻求报道：content@jiqizhixin.com

点个在看 paper不断！