[论文解读] MST: Masked Self-Supervised Transformer for Visual Representation

Zhaowen Li, Ziyang Chen+9 more|arXiv (Cornell University)|Jun 10, 2021

Multimodal Machine Learning Applications参考文献 35被引用 29

一句话总结

MST 引入一种基于注意力引导的掩码策略，用于 Transformer 学习局部和全局的视觉表征，并由全局图像解码器帮助保持空间结构，在 ImageNet 线性评估上实现强性能，在密集预测任务上也具有竞争力的结果。

ABSTRACT

Transformer has been widely used for self-supervised pre-training in Natural Language Processing (NLP) and achieved great success. However, it has not been fully explored in visual self-supervised learning. Meanwhile, previous methods only consider the high-level feature and learning representation from a global perspective, which may fail to transfer to the downstream dense prediction tasks focusing on local features. In this paper, we present a novel Masked Self-supervised Transformer approach named MST, which can explicitly capture the local context of an image while preserving the global semantic information. Specifically, inspired by the Masked Language Modeling (MLM) in NLP, we propose a masked token strategy based on the multi-head self-attention map, which dynamically masks some tokens of local patches without damaging the crucial structure for self-supervised learning. More importantly, the masked tokens together with the remaining tokens are further recovered by a global image decoder, which preserves the spatial information of the image and is more friendly to the downstream dense prediction tasks. The experiments on multiple datasets demonstrate the effectiveness and generality of the proposed method. For instance, MST achieves Top-1 accuracy of 76.9% with DeiT-S only using 300-epoch pre-training by linear evaluation, which outperforms supervised methods with the same epoch by 0.4% and its comparable variant DINO by 1.0\%. For dense prediction tasks, MST also achieves 42.7% mAP on MS COCO object detection and 74.04% mIoU on Cityscapes segmentation only with 100-epoch pre-training.

研究动机与目标

• 通过捕捉局部补丁上下文同时保留全局语义来推动视觉自监督学习的改进，以便更好地迁移到密集预测任务。
• 开发一种由多头自注意力引导的掩码方案，以有选择地掩盖补丁而不破坏关键区域。
• 引入一个全局图像解码器以恢复空间信息并提升下游密集预测任务的性能。
• 展示该方法在 ImageNet 分类以及下游任务如目标检测和语义分割上的有效性。

提出的方法

• 提出 MST，一种掩蔽自监督 Transformer 框架，使用来自教师网络自注意力映射的注意力引导掩蔽策略。
• 对低注意力的补丁对应的 token 进行掩蔽，同时保留重要区域，使用基于注意力阈值的伯努利过程进行条件化。
• 使用全局图像解码器对掩蔽与未掩蔽的 token 进行恢复，以重建原始图像并促进空间保真度。
• 在自监督交叉熵损失（对比学习式目标）与用于图像重建的恢复损失的组合下进行训练，由动量教师-学生设置和标准自监督增强引导。
• 使用基于 CNN 的解码器从掩蔽的 Transformer 表征中恢复图像像素，确保与下游密集预测任务的兼容性。
• 在 ImageNet 上使用 DeiT-S 和 Swin-T 作为骨干网络，通过线性与 k-NN 探针进行评估，并评估向 MS COCO 目标检测和 Cityscapes 语义分割的迁移性。

实验结果

研究问题

• RQ1 在自监督预训练过程中，基于注意力引导的掩蔽策略能否在掩蔽信息量较低的补丁的同时保护关键图像区域？
• RQ2 引入全局图像解码器以重建空间信息，是否在不增加训练时间的情况下提升密集预测任务的性能？
• RQ3 MST 在 ImageNet 线性精度和下游密集预测基准与最先进的自监督方法相比的表现如何？
• RQ4 该方法在不同的 Transformer 架构（如 DeiT-S、Swin-T）和训练预算下是否具有鲁棒性？

主要发现

• MST 在 DeiT-S、300 轮训练的 ImageNet 线性评估中实现 76.9% 的 top-1 准确率，超越 DINO 且接近更长训练的结果。
• 在 MS COCO 目标检测上，MST 在 100 轮训练下达到 42.7% bbox mAP 和 38.8% mask mAP，超过 ImageNet 监督基线。
• 在 Cityscapes 语义分割中，MST 获得 74.04% mIoU 和 82.35% mAcc，在可比设置下击败了有监督和 DINO 预训练的基线。
• 基于注意力引导的掩蔽策略显著提升线性探针性能（73.7% 对 63.2% 的随机掩蔽）并展示出对掩蔽超参数的鲁棒性。
• 在消融实验中，使用带有适当设置的注意力引导掩蔽策略在 100/300 轮训练中保持较强的性能，并且投影头中的 BN 处理细节可能影响收益。