MST: Masked Self-Supervised Transformer for Visual Representation

Part 1 问题和挑战

Yann LeCun曾说“如果人工智能是一个蛋糕，则蛋糕的主要成分就是无监督习”。这句话反应了无监督学习在深度学习领域扮演着重要的作用。当前普遍的方法集中在如何设计有效的代理任务来对无标注的输入数据学习到好的视觉表征。在计算机视觉上，目前比较流行的直接有效的方法是对比学习，将训练数据的每个实例当做单一的分类。基于这个实力判别，很多自监督方法在分类任务上得到了有效的提升。他们成功弥补了自监督方法和监督方法的代沟。然而，这项任务仍然具有挑战：

a. 掩码语言模型在自然语言领域得到了广泛的应用。图像是高维特征，多噪声且相比于文本形式复杂。在视觉领域中，图像的主要信息会被随机分到不同的token中，如果这些token被随机masked掉，将会导致很差的表现。这个随机掩码语言模型容易掩盖图像的关键区域的token，这样会导致误判且不适合直接应用于自监督视觉Transformers。

b. 很多自监督方法是利用全局特征学习图像级别预测，对于像素级别预测优化不足。当前自监督学习方法也许对图像分类任务过度拟合，对下游密集任务预测表现效果不好。

Part 2 方法介绍

针对以上提出的问题，我们提出掩码Transformer自监督学习方法，如下图所示。MST创造性的引入了注意力特征图引导掩码策略并利用掩码特征来恢复全局图像特征任务。我们将介绍如何利用注意力特征引导掩码策略帮助掩码语言模型应用到视觉领域。最后我们将介绍网络的结构和实验细节。

图1 MST整体全流程图

1. 自监督方法网络结构

我们根据多重裁剪在多种通用的数据增强下，为每个图像生成多个视图。这样操作可以获得两个标准分辨率的裁剪图片  和  (表示全局视图)和N个低分辨率裁剪样本（表示局部视图）。如图1所示，整体方法是由两个编码来进行编码，教师网络  和学生网络  ，参数分别是  和  。两个编码器都是Transformer的backbone和projection head组成。教师网络的编码器参数  是由学生网络的编码器参数  来动态平均更新。更新公式如公式1所示：

其中：m是动量系数。

给定一个固定的教师网络  ，学生网络  通过最小化交叉熵损失如公式2所示来学习参数  ：

2. 视觉掩码模型的掩码token策略

图2 MST的注意力引导面具策略的插图。与原始随机掩码相比，通过保留图像中的关键区域来改进。从左到右描述图像: (a)输入图像，(b)通过自注意模块获得的注意力图，(c)可能导致关键特征丢失的随机掩码策略，(d) MST的注意力引导掩码策略，仅遮盖非必要区域。事实上，掩码策略是掩蔽Token。

随机掩码策略: 受到自然语言处理领域掩码语言模型策略的启发，我们应用这个随机掩码策略到自监督学习。根据公式3所示，重要区域的token和非重要区域的token有同样的概率可能被掩码。如图2所示，我们可以观察到随机掩码策略会消除重要区域的tokens，导致难以区分输入图像的语义信息。这个随机采样策略会抑制输入图像的重要区域，不利于网络的识别能力。这个策略不适合直接应用到自监督视觉Transformer，如果掩码策略未正确调制，则整体性能会恶化。

其中：m代表的掩码区域，p是掩码的概率（默认为0.15），prob是随机产生的概率值。

注意力引导掩码策略: 我们提出注意力引导掩码策略动态控制掩码区域并降低重要区域被掩码的概率。同时，这个方法并不会增加额外的计算时间，整体方法结构伪代码所示。如公式4所示，我们按照升序对每张图像的不同patch的注意力特征进行排序，并将排序后的注意力特征值作为阈值，则低于的区域作为掩码的候选区域。学生分支模型收到不同patches的重要性，根据概率生成掩码。

其中：Attn代表注意力特征。

遵循BERT方法，掩码区域填充有可学习的掩码嵌入 [MASK]。如公式5所示，注意力引导掩码策略保证高分的patch不被掩码。

注意力引导掩码策略对于预训练模型有两点益处：

a. 模型利用语义信息来学习到不同patches之间的关系。使得模型既保留了图像全局语义信息，同时在图像的局部细节有关注。

b. 我们的策略可以避免掩盖关键区域，同时用可学习的掩码特征，使得模型专注于关键区域。

3. 视觉Transformer掩码解码方法

在掩码语言模型中，利用非掩码区域特征来预测掩码token。区别于原始的掩码语言模型，我们的方法利用非掩码的区域特征来恢复重建原始的输入图片。为了使网络能够在视觉任务上实现像素级恢复，从而增强图像的像素级信息特征和细粒度空间结构的能力。利用卷积的偏置性特征，重建任务利用卷积神经网络作为解码器，使用卷积层和上采样操作交替堆叠。如下公式6展示了重建损失函数：

其中：x代表输入图片，g代表解码器，代表学生分支的编码器，  代表学生分支的参数，  代表解码器参数。

全体损失函数如公式7所示：

其中：λ代表权重系数。

Part 3 实验结果

我们使用不同的transformer结构在ImageNet benchmark上训练了预训练模型，然后验证他们的迁移能力在下游任务上例如目标检测和语义分割，以及不同掩码参数对模型的影响。

1. ImageNet Benchmark的对比

如表1所示为我们的方法与当前主要自监督算法比较。所有的这些方法拥有相同的backbone用于公平比较。我们方法的300-epoch模型实现76.9% 在线性评估的top-1的精度。我们的方法在相同的训练epochs下超越当时最好的自监督方法DINO大约1.7%，甚至接近DINO更长训练策略的表现效果（77.0%的800epoch）。需要强调的是，我们的算法缓解了自监督学习对极长训练时间的需求，并且能够在仅100个epochs的情况下获得不错的效果（75.0%）。

MST是通用的方法可以应用到任何基于Transformer结构的自监督方法。这里我们使用流行的Swin-T作为示例，它具有与DeiT-S相似的参数数量。使用相同的训练epochs，MST的性能比MoBY高出1.8%，这是一种为Swin-T精心设计的自监督学习方法。Swin-T与DeiT-S共享相同的超参数，它仍然可以通过进一步调整来改进。

表1 ImageNet上流行的自监督学习方法的比较

2. 目标检测和实例分割下游任务

在表2中，我们显示了通过不同的自我监督方法和监督训练学习的表示的表现。为了公平比较，所有这些方法都预先训练了100个epochs。我们观察到，我们的方法以42.7%的bbox mAP和38.8%的掩模mAP达到了最佳效果。它比ImageNet监督模型高出1.2%和0.5%，MoBY结果在同一时期分别比ImageNet监督模型高出1.2%和0.5%。结果表明，MST不仅在图像分类任务上表现良好，而且在下游密集预测任务中表现良好。因此它具有很强的迁移能力。

表2 在 MS COCO 上微调的对象检测和实例分段结果

如表所示，它说明了监督方法、DINO和我们的方法在该评估上的比较。我们的方法实现了最高的mIoU74.7%和mAcc82.35%。它的表现优于监督结果（+2.71%mIoU和+2.05%mAcc）和DINO预训练结果（+1.08% mIoU和 +1.03%mAcc）。我们的模型也适用于语义分割任务的迁移。

表3 在城市景观上微调的语义分割结果

3. 不同Mask策略的影响

表4显示了不同掩码策略的影响。我们使用随机掩码策略，注意力引导的掩码策略和无掩码来训练DeiT-S。为了公平比较，所有方法都以相同的概率p进行掩码。可以观察到随机掩码策略的性能下降。此策略可能会抑制识别图像能力（从 73.1 到 63.2）。随机掩码策略可能会破坏原始图像关键区域的tokens，这些tokens对于识别对象可能是必不可少的。被屏蔽的输入可能包含不完整甚至误导性的信息。相反，我们的注意力引导掩码策略的表现稳步提高（从73.1到73.7）。基本区域大多被保留下来，这可能是我们假设的有力证据。

表4 不同掩模策略（DeiT-S）的线性探针结果

4. 不同掩码超参数的影响

表5验证了注意力引导掩码策略下不同掩码超参数的性能。我们按升序对每个图像的不同patch的注意力图进行排序，并将前1/num个patch拆分为被掩码的候选。删除这些候选可以迫使网络从相邻的patch中学习局部特征，从而在不破坏语义的情况下增强对局部上下文进行建模的能力。这些候选根据概率p进行屏蔽。ImageNet上线性评估的Top-1精度如表下表所示。当num设置为8时，任何选择p都可以得到一个可靠的结果，这表明最后的1/8个patch作为掩码候选是相对安全的。

表5 基于注意力的掩码策略的超参数设置

5. 与 BERT 的区别

在表6中，我们使用纯掩码语言模型和DeiT-S在100个epochs下进行实验，结果在相同的实验配置下约为40%。然后我们进一步调整其学习率和其他超参数，最佳结果仅为61%，远低于DINO的10.6%（DINO结果为71.6%），也比监督结果低7.7%（监督结果为68.7%）。表明纯掩码语言模型方法可能不适合计算机视觉任务。此外，我们用对比损失+ BERT解决方案（即DINO+纯掩码语言模型）进行实验，线性结果为71.9%。我们的方法比其结果高出2.0%（73.9%）。结果证明我们的方法比原来的结果方法更好。同时，我们进一步进行了实验，仅将[mask] token替换为纯掩码语言模型策略，线性结果为73.5%，这也落后于我们的结果。这些结果充分展示了MLM对计算机视觉的更好设置，并进一步突出了我们论文的技术贡献。

表6与 BERT 的区别

Part 4 结语

本文探讨了当前视觉自我监督学习的两个问题，即缺乏局部信息提取和空间信息丢失。为了克服上述问题，我们提出了一种基于Transformer的新型自监督学习方法，称为MST。 MST利用注意力引导的掩码策略来捕获patch之间的局部关系，同时保留全局语义信息。需要注意的是，注意力引导的掩码策略是基于从教师模型中提取的Multi-Head自注意特征图，不会造成额外的计算成本。此外，在注意力引导掩模策略下，进一步使用全局图像解码器来恢复图像的空间信息，这对于密集的预测任务至关重要。该方法在多个下游视觉任务中表现出良好的通用性和可扩展性。

Reference:

 [1] Devlin, J., Chang, M.W., Lee, K., Toutanova, K.N.: Bert: Pre-training of deep bidirectional transformers for language understanding. In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers). pp. 4171–4186 (2018)

[2] Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., Joulin, A.: Emerging properties in self-supervised vision transformers. arXiv: Computer Vision and Pattern Recognition

[3] Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.14030 (2021)

[4]Xie, Z., Lin, Y., Yao, Z., Zhang, Z., Dai, Q., Cao, Y., Hu, H.: Self-supervised learning with swin transformers. arXiv preprint arXiv:2105.04553 (2021)