[论文解读] A Closer Look at Self-Supervised Lightweight Vision Transformers

Shaoru Wang, Jin Gao+3 more|arXiv (Cornell University)|May 28, 2022

Advanced Neural Network Applications被引用 8

一句话总结

本文研究了轻量级视觉Transformer（ViTs）的自监督预训练，表明经过MAE正确预训练的原始ViT-Tiny（570万参数）在ImageNet上实现了79.0%的top-1准确率，性能可与设计复杂的最先进模型相媲美。研究揭示，当下游数据充足时，预训练更有利于低层特征学习，并提出一种在MAE预训练过程中应用的蒸馏策略，显著提升了数据稀缺任务的性能。

ABSTRACT

Self-supervised learning on large-scale Vision Transformers (ViTs) as pre-training methods has achieved promising downstream performance. Yet, how much these pre-training paradigms promote lightweight ViTs' performance is considerably less studied. In this work, we develop and benchmark several self-supervised pre-training methods on image classification tasks and some downstream dense prediction tasks. We surprisingly find that if proper pre-training is adopted, even vanilla lightweight ViTs show comparable performance to previous SOTA networks with delicate architecture design. It breaks the recently popular conception that vanilla ViTs are not suitable for vision tasks in lightweight regimes. We also point out some defects of such pre-training, e.g., failing to benefit from large-scale pre-training data and showing inferior performance on data-insufficient downstream tasks. Furthermore, we analyze and clearly show the effect of such pre-training by analyzing the properties of the layer representation and attention maps for related models. Finally, based on the above analyses, a distillation strategy during pre-training is developed, which leads to further downstream performance improvement for MAE-based pre-training. Code is available at https://github.com/wangsr126/mae-lite.

研究动机与目标

• 评估自监督预训练在轻量级视觉Transformer（特别是ViT-Tiny）上的有效性，与最先进架构进行对比。
• 探究自监督预训练为何以及如何提升下游任务性能，特别是在数据稀缺场景下的表现。
• 分析ViT各层在预训练和微调过程中的表征行为，识别影响下游性能的关键因素。
• 提出并验证一种在基于MAE的预训练过程中应用的蒸馏策略，以进一步提升轻量级ViT的性能。

提出的方法

• 在图像分类和密集预测任务上，对ViT-Tiny应用多种自监督预训练方法（MoCo-v3，对比学习；MAE，掩码图像建模）进行基准测试。
• 通过深入分析层特征表示和注意力图，研究预训练如何在不同网络深度上塑造特征学习过程。
• 发现当下游数据充足时，低层特征更为关键；而在数据稀缺设置下，高层特征则占据主导地位。
• 发现MAE预训练通过集中注意力图引入了局部性归纳偏置，这有助于性能提升。
• 提出一种在MAE预训练过程中引入蒸馏的策略，由更大的教师网络指导学生模型ViT-Tiny，以提升特征质量，尤其在高层特征中效果显著。
• 在ImageNet、Flowers、Aircraft、CIFAR100、COCO目标检测与分割等基准上验证该方法，并通过消融实验分析蒸馏组件的影响。

实验结果

研究问题

• RQ1 自监督预训练能否弥合原始轻量级ViT与设计复杂的最先进架构之间的性能差距？
• RQ2 ViT各层的表征学习行为在预训练与微调阶段有何差异？其表现如何依赖于下游数据规模？
• RQ3 为何MAE预训练能带来更好的下游性能？它引入了哪些归纳偏置？
• RQ4 在MAE预训练过程中引入知识蒸馏能否进一步提升轻量级ViT的性能？

主要发现

• 经过MAE预训练的原始ViT-Tiny在ImageNet上实现了79.0%的top-1准确率，性能与此前最先进模型相当或更优，尽管后者采用了复杂的设计。
• 当有充足标注数据时，预训练ViT的低层对下游性能贡献更大；而在数据稀缺设置下，高层特征变得更为重要。
• MAE预训练通过集中注意力图引入了局部性归纳偏置，该特性与泛化能力提升及性能增益密切相关。
• 所提出的MAE预训练阶段蒸馏策略显著提升了下游性能，尤其在数据不足的分类任务以及密集预测任务（如COCO目标检测与分割）中效果突出。
• 该蒸馏方法提升了高层特征的质量，表现为表征学习能力增强，并与教师网络输出的对齐度更高。
• 线性探测评估结果证实，该方法所获得的表征在多种下游任务上均具备良好的泛化能力。