微软亚太研究院推出TinyMIM，利用知识蒸馏来优化小型ViT

一、研究动机

近年来，掩码建模（Masked Image Modeling, MIM）作为一种自监督训练方法，展示了其在大规模模型中的有效性。然而，研究发现，MIM在小型模型（如ViT-T 5M参数）上反而可能导致性能下降。特别是在ImageNet上，使用MAE训练的ViT-L相比于传统监督学习模型提升了3.3%的分类准确率，但ViT-T的分类准确率却下降了0.6%。因此，微软亚洲研究院提出了TinyMIM，旨在通过知识蒸馏将大模型的知识有效迁移至小模型，而不改变ViT的结构或引入其它归纳偏置。

二、TinyMIM的贡献

TinyMIM的研究重点在于以下几个方面：

• 蒸馏token之间的关系比单独蒸馏class token或特征图更为有效。
• 中间层作为蒸馏目标的效果更佳。

1. 蒸馏目标：

• 带掩码的图片效果较差。
• 学生模型需要适当的drop path，而教师模型则不需要。

1. 数据增强与正则化：

• 将MIM作为辅助损失函数的效果不明显。

1. 辅助损失函数：

• 采用序列化的蒸馏策略（如ViT-B -> ViT-S -> ViT-T）能取得最佳效果。

1. 宏观蒸馏策略：

三、方法

TinyMIM的核心在于系统性地研究影响蒸馏效果的因素，包括特征、关系和输入。具体来说：

• 特征：
• 中间block特征与输出特征的选择影响蒸馏效果。
• 在注意力层和前馈层特征的蒸馏中，不同层次的蒸馏带来不同影响。

• 关系：
• Q、K、V特征之间的关系也可以作为知识蒸馏的目标。

• 输入：
• 探索带掩码的输入是否适合作为知识蒸馏的输入。

四、实验结果

TinyMIM在ImageNet-1K上进行预训练，且教师模型同样在该数据集上预训练。实验结果显示，TinyMIM显著超越了基于MAE的方法，尤其是在小型模型上表现突出。具体而言，超小模型ViT-T的分类准确率达到了75.8%，相比MAE基线提升了4.2%；小模型ViT-S的准确率为83.0%，比之前最佳方法提升了1.4；而在Base尺寸的模型上，TinyMIM分别超越MAE基线和之前最佳模型CAE 4.1和2.0。

此外，TinyMIM在鲁棒性测试中同样表现优异，在ImageNet-A和ImageNet-R上相比MAE-B分别提升了6.4和4.6。

五、消融实验

通过消融实验，研究了不同蒸馏关系、蒸馏策略和中间层的影响。结果表明：

• 蒸馏不同关系：
• 蒸馏QK、VV关系并结合Softmax实现了最佳效果。

• 不同蒸馏策略：
• TinyMIM的关系蒸馏方法在各个模型尺寸上均优于MAE基线、class token蒸馏和特征图蒸馏。

• 蒸馏中间层：
• 蒸馏第十八层特征表现最佳。

六、结论

本文提出了TinyMIM，这是首个成功使小型模型受益于掩码重建建模（MIM）预训练的模型。通过知识蒸馏的方式，TinyMIM在不依赖掩码重建任务的情况下，成功模拟了大模型的特征关系。研究表明，关系蒸馏的效果优于特征蒸馏和类标记蒸馏，为未来的研究奠定了基础。TinyMIM凭借其简单性和强大性能，展现了在小型模型训练中的潜力，值得进一步探索与应用。