微软亚洲研究院
本文提出了 TinyMIM，它是第一个成功地使小模型受益于 MIM 预训练的模型。
一、研究动机
掩码建模（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）效果最好。
二、方法
我们系统性的调研了蒸馏的目标，输入的图片，蒸馏目标模块。
2.1 影响蒸馏效果的因素
1）特征：
a. 中间 block 特征和输出特征
当 i=L 时，指的是 Transformer 输出层的特征。当 i< L 时，指的是 Transformer 中间层的特征。
b. 注意力（Attention）特征和前馈层（FFN）层特征
Transformer 每一个 block 有 Attention 层和 FFN 层，蒸馏不同的层会带来不同的影响。
