[论文解读] Self-distillation with Batch Knowledge Ensembling Improves ImageNet Classification

Yixiao Ge, Xiao Zhang+7 more|arXiv (Cornell University)|Apr 27, 2021

Advanced Neural Network Applications参考文献 58被引用 16

一句话总结

本文提出批量知识集成（BAKE），一种自蒸馏方法，通过利用同一小批量内样本之间的样本间亲和力，实现在线知识集成，从而提升ImageNet分类性能。该方法仅增加1.5%的FLOPs且不增加参数，即可将Swin-T的Top-1准确率提升0.7%。

ABSTRACT

The recent studies of knowledge distillation have discovered that ensembling the "dark knowledge" from multiple teachers or students contributes to creating better soft targets for training, but at the cost of significantly more computations and/or parameters. In this work, we present BAtch Knowledge Ensembling (BAKE) to produce refined soft targets for anchor images by propagating and ensembling the knowledge of the other samples in the same mini-batch. Specifically, for each sample of interest, the propagation of knowledge is weighted in accordance with the inter-sample affinities, which are estimated on-the-fly with the current network. The propagated knowledge can then be ensembled to form a better soft target for distillation. In this way, our BAKE framework achieves online knowledge ensembling across multiple samples with only a single network. It requires minimal computational and memory overhead compared to existing knowledge ensembling methods. Extensive experiments demonstrate that the lightweight yet effective BAKE consistently boosts the classification performance of various architectures on multiple datasets, e.g., a significant +0.7% gain of Swin-T on ImageNet with only +1.5% computational overhead and zero additional parameters. BAKE does not only improve the vanilla baselines, but also surpasses the single-network state-of-the-arts on all the benchmarks.

研究动机与目标

• 为解决监督图像分类中单热标签的局限性，该局限性阻碍了性能的进一步提升。
• 在保持或提升性能的同时，降低知识蒸馏的计算和内存开销。
• 通过单个网络实现在小批量样本间的在线、轻量级知识集成。
• 通过跨样本知识传播生成的精细化软目标，提升图像分类的泛化能力和鲁棒性。

提出的方法

• BAKE通过在推理过程中传播并聚合同一小批量内其他样本的预测结果，基于样本间亲和力加权实现动态知识集成。
• 样本间亲和力通过当前网络的特征进行估计，确保相似样本对知识集成的贡献更大。
• 知识传播与集成通过迭代方式应用直至收敛，通过高效推理近似实现，避免高计算成本。
• 该方法使用单个网络生成用于自蒸馏的精细化软目标，无需多个教师网络或额外分支。
• 该框架兼容多种架构，包括ResNet和Swin Transformer，并支持使用同步BatchNorm的分布式训练。

实验结果

研究问题

• RQ1 在不增加参数或计算开销的前提下，小批量内样本间的知识集成能否提升蒸馏性能？
• RQ2 实时进行的样本间知识传播如何影响自蒸馏中软目标的质量？
• RQ3 与最先进单网络方法相比，BAKE在ImageNet及下游任务上的准确率能提升多少？
• RQ4 BAKE在测试时鲁棒性基准（如标签噪声和分布偏移）下的表现如何？

主要发现

• BAKE在训练时仅增加3.7%的FLOPs，将ResNet-50在ImageNet上的准确率提升+1.2%（76.8% → 78.0%）。
• BAKE仅增加1.5%的计算开销且不增加额外参数，将Swin-T在ImageNet上的准确率提升+0.7%（81.3% → 82.0%）。
• BAKE在所有评估基准（包括CIFAR-100、TinyImageNet、CUB-200-2011、MIT67和Stanford Dogs）上均持续优于单网络最先进方法。
• 该方法在COCO检测与分割任务上的迁移学习性能得到提升，表明其学习表征具有更好的泛化能力。
• BAKE在ImageNet-R、ImageNet-A和ImageNet-P等扰动数据集上提升了测试时鲁棒性，表明模型鲁棒性得到增强。
• 消融实验表明，通过近似推理实现的无限次迭代知识集成，相比单次迭代集成，性能更鲁棒；且同步BatchNorm在采用每类采样策略的分布式训练中至关重要。