CVPR 2023 | 微软亚研提出TinyMIM，用知识蒸馏改进小型ViT

本文提出了 TinyMIM，它是第一个成功地使小模型受益于 MIM 预训练的模型。

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

一、研究动机

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

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

我们系统性的研究了蒸馏目标、数据增强、正则化、辅助损失函数等对于蒸馏的影响。在严格的只用 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 层，蒸馏不同的层会带来不同的影响。

c.QKV 特征

在 Attention 层内会有 Q，K，V 特征，这些特征用于计算注意力机制，我们也调研了直接蒸馏这些特征。

2）关系

Q，K，V 用于计算注意力图，这些特征之间的关系也可以作为知识蒸馏的目标。

3）输入：是否带掩码

传统的知识蒸馏是直接输入完整的图片。我们的方法为了探索蒸馏掩码建模模型，所以我们也探索了带掩码的图片是否适合作为知识蒸馏时候的输入。

2.2 知识蒸馏方法对比

1）Class Token 蒸馏：

最简单的方法就是类似 DeiT 直接蒸馏 MAE 预训练模型的 class token:

其中指学生模型的 class token，而 指老师模型的 class token。

2）特征蒸馏：我们直接参考了 feature distillation [1] 作为对比

3）关系蒸馏：我们提出了也是本文默认的蒸馏策略

三、实验

3.1 主要实验结果

我们的方法在 ImageNet-1K 上预训练，而且教师模型也是在 ImageNet-1K 预训练。然后我们将我们预训练的模型在下游任务（分类、语义分割）上进行了微调。模型表现如图：

我们的方法显著超过之前基于 MAE 的方法，尤其是小模型。具体来讲，对于超小的模型 ViT-T，我们的方法实现了 75.8% 的分类准确性，相比 MAE 基线模型实现了 4.2 的提升。对于小模型 ViT-S，我们实现了 83.0% 的分类准确性，比之前最好的方法提升了 1.4。对于 Base 尺寸的模型，我们的方法分别超过 MAE 基线模型和以前最好的模型 CAE 4.1 和 2.0。

同时我们也测试了模型的鲁棒性，如图所示：

TinyMIM-B 对比 MAE-B，在 ImageNet-A 和 ImageNet-R 分别提升了 + 6.4 和 +4.6。

3.2 消融实验

1）蒸馏不同关系

同时蒸馏 QK,VV 关系而且在计算关系的时候有 Softmax 实现了最好的效果。

2）不同的蒸馏策略

TinyMIM 这种蒸馏关系的方法实现了比 MAE 基线模型，class token 蒸馏，特征图蒸馏都更好的效果，在各种尺寸的模型上都是如此。

3）蒸馏中间层

我们发现蒸馏第十八层实现了最好的效果。

四、结论

在本文中，我们提出了 TinyMIM，它是第一个成功地使小模型受益于掩码重建建模（MIM）预训练的模型。我们没有采用掩码重建作为任务，而是通过以知识蒸馏的方式训练小模型模拟大模型的关系来预训练小模型。TinyMIM 的成功可以归功于对可能影响 TinyMIM 预训练的各种因素的全面研究，包括蒸馏目标、蒸馏输入和中间层。通过大量的实验，我们得出结论，关系蒸馏优于特征蒸馏和类标记蒸馏等。凭借其简单性和强大的性能，我们希望我们的方法能够为未来的研究提供坚实的基础。

参考

[1] Wei, Y., Hu, H., Xie, Z., Zhang, Z., Cao, Y., Bao, J., ... & Guo, B. (2022). Contrastive learning rivals masked image modeling in fine-tuning via feature distillation. arXiv preprint arXiv:2205.14141.

作者：魏芳芸 来源：公众号【机器之心】

Illustration by IconScout Store from IconScout

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线480+期talk视频，2400+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈

⤵一键送你进入TechBeat快乐星球