通过掩码图像建模预训练探索轻量级视觉Transformer的实验研究
一种通过掩码图像建模预训练探索强轻量级视觉Transformer的实验研究
学术背景
近年来，自监督学习（self-supervised learning, SSL）在计算机视觉领域取得了显著进展。特别是掩码图像建模（masked image modeling, MIM）预训练方法在大规模视觉Transformer（vision transformers, ViTs）上的成功应用，使得基于这些模型的下游任务性能得到了极大提升。然而，现有的研究主要集中在大型ViTs上，对于轻量级ViTs的预训练方法及其效果的研究相对较少。此外，尽管许多研究致力于设计复杂的轻量级ViTs架构以提高性能，但很少有工作关注如何优化预训练策略来进一步提升现有轻量级模型的表现。本文旨在探讨MIM预训练能否同样有效应用于极其简单的轻量级ViTs，并通过系统性的观察、分析和解决方案来解决这一问题。
论文来源
该论文由Jin Gao、Shubo Lin、Shaoru Wang等多位作者共同完成，他们分别来自中国科学院自动化研究所、中国科学院大学人工智能学院、上海科技大学信息科学与技术学院等多个机构。论文于2024年12月被《International Journal of Computer Vision》接收，并将于2025年正式发表。
研究内容
研究流程
本文采用了一种观察-分析-解决方案的工作流程来进行研究。具体来说，首先系统地观察不同预训练方法在下游微调数据规模上的表现差异；然后分析各层表示相似性和注意力图，揭示MIM预训练在高层学习中的不足；最后提出解耦蒸馏策略，以改善轻量级ViTs的预训练效果。
实验对象及样本量
本研究使用了经过轻微修改的Touvron等人（2021）提出的轻量级ViT-tiny作为实验单元，其参数量为5.7M。此外，还对最近提出的Hierarchical Transformer (Hiera)进行了研究，其参数量为6.5M。实验中涉及的数据集包括ImageNet-1k、ADE20k、LASOT等。
实验过程
- 预训练方法的适应与对比：将多种流行的MIM预训练方法（如MAE、Simmim、BEiT等）以及对比学习（contrastive learning, CL）预训练方法（如MoCo-v3、DINO）和全监督预训练方法应用于轻量级ViTs。
- 基准测试：在ImageNet分类任务上对预训练后的轻量级模型进行微调评估，并进一步评估它们在其他数据集上的迁移性能。
- 线性探测与模型分析：通过线性探测和基于CKA（centered kernel alignment）的层表示相似性分析，以及注意力图分析，揭示不同预训练方法的工作机制。
- 解耦蒸馏策略的提出与验证：提出一种解耦蒸馏策略，通过将重建任务与蒸馏任务分离，进一步改善MIM预训练的效果。
主要结果
- 适当的预训练可以释放轻量级ViTs的巨大潜力：几乎所有比较的预训练方法都能在ImageNet分类任务上优于随机初始化的方法，而MIM预训练在适度的预训练成本下表现出色。
- MIM预训练可以使简单的ViT-tiny在ImageNet上达到最新的轻量级ViT衍生物的性能：通过MIM预训练增强的简单ViT-tiny在ImageNet分类任务上达到了与一些复杂设计的轻量级ViTs相当的性能。
- 轻量级ViTs的自监督预训练难以从“LLM-like”数据扩展中受益：MIM预训练在更大规模的数据集上表现不佳，表明轻量级ViTs的容量限制了其表示质量。
- 尽管MIM预训练在ImageNet上表现良好，但在下游任务上的迁移性能较差：特别是在数据不足的下游任务上，MIM预训练的表现不如CL预训练。
结论
本文通过系统性的观察、分析和解决方案，提出了改进轻量级ViTs的MIM预训练方法。具体来说，通过解耦蒸馏策略，不仅使预训练后的轻量级ViTs在高层学习到与识别相关的语义信息，还能保留MIM预训练带来的有用局部归纳偏差。实验结果表明，这种方法在多个下游任务上取得了显著的性能提升，包括ImageNet分类、ADE20k语义分割、LASOT单目标跟踪等任务。
研究亮点
- 重要发现：适当的预训练可以极大地提升极其简单的轻量级ViTs的性能，使其在ImageNet分类任务上达到最先进的水平。
- 问题意义：解决了轻量级ViTs在预训练方面的瓶颈问题，为未来的轻量级模型设计提供了新的思路。
- 方法创新：提出的解耦蒸馏策略是一种新颖的方法，通过分离重建任务与蒸馏任务，有效地改善了MIM预训练的效果。
其他有价值的信息
除了上述主要内容外，本文还对多种预训练方法进行了全面的基准测试，涵盖了多个下游任务，为未来的研究提供了丰富的参考数据。此外，论文还公开了改进后的代码和原始结果，方便其他研究人员复现和进一步研究。
本文通过系统性的研究，揭示了MIM预训练在轻量级ViTs上的潜力，并提出了一种有效的解耦蒸馏策略，为轻量级模型的设计和优化提供了新的方向。