一、研究动机

掩码建模（MIM, MAE）被证明是非常有效的自监督训练方法。然而，如图 1 所示，MIM 对于更大的模型效果相对更好。当模型很小的时候（比如 ViT-T 5M 参数，这样的模型对于现实世界非常重要），MIM 甚至可能一定程度上降低模型的效果。比如用 MAE 训练的 ViT-L 比普通监督训练的模型在 ImageNet 上的分类效果提升 3.3%，但是用 MAE 训练的 ViT-T 比普通监督训练的模型在 ImageNet 上的分类效果降低了 0.6%。

在这篇工作中我们提出了 TinyMIM，其在保持 ViT 结构不变并且不修改结构引入其他归纳偏置（inductive bias）的基础上、用蒸馏的方法迁移大模型上的知识到小模型。

 

 

• 论文地址：https://arxiv.org/pdf/2301.01296.pdf
• 代码地址：https://github.com/OliverRensu/TinyMIM

 

我们系统性的研究了蒸馏目标、数据增强、正则化、辅助损失函数等对于蒸馏的影响。在严格的只用 ImageNet-1K 作为训练数据的情况下（包括 Teacher model 也只用 ImageNet-1K 训练）和 ViT-B 作为模型，我们的方法实现了当前最好的性能。如图所示：

 

把我们的方法（TinyMIM）和基于掩码重建的方法 MAE，以及监督式学习的方法从头开始训练的 DeiT 作比较。MAE 在模型比较大的时候有显著的性能提升，但是在模型比较小的时候提升幅度有限甚至会伤害模型的最终效果。我们的方法 TinyMIM 在不同模型的大小上都有大幅提升。

我们的贡献如下：

1. 蒸馏的目标（Distillation targets）:1）蒸馏 token 之间的关系比单独蒸馏 class token 或者特征图（feature map）更有效；2）用中间层作为蒸馏的目标更有效。

2. 数据增强和模型正则化（Data and network regularization）：1）用带掩码的图片效果更差；2）学生模型需要一点 drop path，但是 teacher 模型不需要。

3. 辅助损失函数（auxiliary losses）：MIM 作为辅助损失函数没有意义。

4. 宏观蒸馏策略（Macro distillation strategy）：我们发现序列化的蒸馏（ViT-B -> ViT-S -> ViT-T）效果最好。