TinyMIM:蒸馏 MIM 预训练模型的实证研究
计算机视觉与模式识别
2023-01-04 v1
摘要
掩码图像建模(MIM)在预训练大型视觉 Transformer(ViT)方面表现强劲。然而,对实际应用至关重要的小型模型无法或仅能边际地从这种预训练方法中获益。在本文中,我们探索蒸馏技术,以将基于 MIM 的大型预训练模型的成功迁移到较小的模型上。我们系统地研究了蒸馏框架中的不同选项,包括蒸馏目标、损失、输入、网络正则化、顺序蒸馏等,揭示出:1)蒸馏 token 关系比基于 CLS token 和特征的蒸馏更有效;2)当学生网络与教师网络的深度不匹配时,使用教师网络的中间层作为目标比使用最后一层表现更好;3)偏好弱正则化;等等。基于这些发现,我们在 ImageNet-1K 分类上相对于从头开始的 MIM 预训练取得了显著的微调精度提升,使用了所有 ViT-Tiny、ViT-Small 和 ViT-base 模型,分别获得 +4.2%/+2.4%/+1.4% 的收益。我们的 base 尺寸 TinyMIM 模型在 AE20K 语义分割中达到 52.2 mIoU,比 MAE 基线高 +4.1。我们的 tiny 尺寸 TinyMIM 模型在 ImageNet-1K 图像分类上达到 79.6% 的 top-1 精度,为同等尺寸和计算预算的小型视觉模型树立了新纪录。这一强劲性能提出了一种开发小型视觉 Transformer 模型的替代途径,即通过探索更好的训练方法,而非像以往大多数工作那样向架构中引入归纳偏置。代码可在 https://github.com/OliverRensu/TinyMIM 获取。
引用
@article{arxiv.2301.01296,
title = {TinyMIM: An Empirical Study of Distilling MIM Pre-trained Models},
author = {Sucheng Ren and Fangyun Wei and Zheng Zhang and Han Hu},
journal= {arXiv preprint arXiv:2301.01296},
year = {2023}
}
备注
Code is available at https://github.com/OliverRensu/TinyMIM