CVPR 2022 | 继何恺明的MAE后，MSRA提出更简单的掩码图像建模框架

2022 年 4 月 18 日 PaperWeekly

©作者 | 机智勇敢萌刚刚

本文提出一个用于掩码图像建模（masked image modeling）的简单框架 SmiMIM。作者简化了最近提出的方法，而无需任何特殊设计，如利用离散 VAE 或聚类进行块级别的掩码和分词。为了让掩码图像建模任务能学到更好的表示，作者表示该框架中每个组件的简单设计已经能显示出其优异的学习能力：

1. 采用中等大小的掩码块（如 32），对输入图像进行随机掩码，能使之成为强大的代理任务（pretext task）；

2. 直接回归预测原始像素 RGB 值的效果并不比复杂设计的 patch 分类方法差；

3. 预测头可以像线性层一样轻量，性能并不一定比多层的差。

使用 ViT-B，通过对该数据集进行预训练，该方法在 ImageNet-1K 上实现了 83.8% 的 top-1 微调精度，比之前的最佳方法高出 0.6%。将其应用在约 6.5 亿个参数的更大模型时，SwinV2-H 仅使用 ImageNet-1K 数据可以达到 87.1% 的 top-1 精度。作者还利用这种方法来加速 3B 模型（SwinV2-G）的训练，通过比以前少 40 倍的数据，在四个具有代表性的视觉基线数据集上达到了 SOTA。

论文标题：

SimMIM: A Simple Framework for Masked Image Modeling

作者单位：

微软亚洲研究院

收录会议：

CVPR 2022

论文链接：

https://arxiv.org/pdf/2111.09886.pdf

代码链接：

https://github.com/microsoft/SimMIM

提交时间：

2021.11.18

动机

“What I cannot create, I do not understand.” — Richard Feynman

全文从诺贝尔物理学奖得主理查德·费曼（Richard Feynman）的经典名言展开，介绍掩码信号建模任务是一种具有创造性的学习任务，即遮住输入图像的一部分并尝试预测遮住的信号。

在 NLP 中，遵循这一理念，建立在掩码语言建模任务上的自监督学习方法在很大程度上重塑了这一领域，即通过使用巨大的无标签数据来学习非常大规模的语言模型，并被证明可以广泛推广到 NLP 应用中。而在计算机视觉中，虽然有先驱者利用这一理念进行自监督表示学习，但在发展初期，这方面的工作几乎被对比学习的方法遮挡锋芒。

根据语言和视觉领域的特性，在二者上导致明显差异的原因，作者分析出以下三点：

1. 图像具有更强的局部关系，即相互靠近的像素往往是高度相关的。复制靠近的像素可以很好地完成任务，显然通过语义推理并不容易办到。

2. 视觉信号是原始的、低层次的，而文本分词是由人类产生的高级概念。那么，对低层次信号的预测是否对高层次的视觉识别任务有用？

3. 视觉信号是连续的，而文本分词是离散的。那么，如何利用基于分类的掩码语言建模方法处理连续的视觉信号？

一些近期的工作通过引入特殊的设计尝试弥补模态间的鸿沟，解决该难题，这些工作在很多视觉识别任务上展现了出色的迁移能力。

方法

2.1 整体架构

SimMIM 方法通过掩码图像建模来学习表示，该方法对输入图像信号的一部分进行掩码，并预测在掩码区域的原始信号。该框架由 4 个主要组件组成：

1. Masking strategy：给定一张输入图像，该组件负责选择掩码的区域及实现所选区域的掩码。经过掩码后的图像将用作模型输入。

2. Encoder architecture：编码器提取掩码图像上潜在的特征表示，然后用来预测掩码区域的原始信号。经过学习的编码器可用于不同的视觉任务。在本文中，主要考虑两种典型的vision Transformer架构: vanilla ViT和Swin Transformer。

3. Prediction head：预测头用于潜在的特征表示，表示掩码区域中的原始信号。

4. Prediction target：这个组件定义了要预测的原始信号的形式。它可以是原始像素值，也可以是原始像素的变换。另外，定义了损失类型，包括交叉熵分类损失和l1或l2回归损失。

2.2 掩码策略-掩码区域的选择

▲ 图2. 不同掩码策略下的图像。中心区域掩码及块级掩码采用大小为16的补丁，随机掩码策略采用不同的补丁大小（4、8、16和32）。

Patch 对齐的随机掩码策略：对于 Swin Transformer，考虑相同的不同分辨率的补丁大小（4×4 ~ 32×32），默认采用 32×32 的补丁大小。对于 ViT，采用 32×32 作为默认掩码补丁大小。

其他掩码策略：

中心区域掩码策略，让其在图像上随机移动；
块级掩码策略，利用分别为 16x16 和 32x32 的两种掩码块进行掩码。

2.3 预测头

预测头的形式和大小可以是任意的，只要其输入与编码器输出一致，其输出达到预测目标即可。一些早期的工作跟随自编码器使用一个繁琐的预测头（解码器）。在这篇文章中，作者证明了预测头可以做得非常轻，就像线性层一样轻量。作者还尝试较重的头部，如 2 层 MLP、inverse Swin-T，inverse Swin-B。

2.4 预测目标

2.4.1 原始像素值回归

像素值在颜色空间中是连续的。一个直接的想法是通过回归来预测掩码区域的原始像素。一般来说，视觉架构通常生成下采样分辨率的特征图，例如，ViT 为 16×，其他架构为 32×。

为了预测输入图像的全分辨率下的所有像素值：

1. 作者将 feature map 中的每个特征向量映射回原始分辨率，并让该向量负责相应的原始像素的预测。例如，对于 Swin Transformer 编码器生成的 32× 下采样的 feature map，作者使用输出维数为 3072 = 32×32×3 的 1×1 卷积（线性）层来表示 32×32 像素的 RGB 值。对原始图像分别进行 {32×， 16×， 8×， 4×， 2×} 下采样，考虑分辨率较低的目标。

2. 在掩码像素上使用 L1-loss。在实验中还考虑了 L2 和 smooth-L1 损失效果与之类似，默认采用 L1 损失。

实验结果

3.1 实验设置

作者采用 Swin-B 作为消融研究的默认骨干。为了减少实验开销，默认的输入图像大小为 192×192，并将窗口大小调整为 6 以适应改变的输入图像大小。ImageNet-1K 图像分类数据集用于预训练和微调。采用数据增强:随机调整大小裁剪，比例范围为 [0.67,1]，宽高比范围为 [3/ 4,4 /3]，然后进行随机翻转和颜色归一化步骤。SimMIM 组件的默认选项是:一个随机掩码策略，补丁大小为 32×32，掩码率为 0.6；目标图像大小为 192×192 的线性预测头；掩码像素预测的 L1 损失。

3.2 掩码策略

由表 1 可知，不同掩码策略在不同掩码率的微调精度不同。

结论 1：其中简单随机掩码策略最佳精度达到 83.0%，比其他更特别的策略的最佳精度高 0.3%。

结论 2：当采用较大的掩码块（大小为 32）时，这种简单的策略在较宽的掩码率范围（10%-70%）上表现稳定。假设一个大遮罩块的中心像素距离可见像素足够远。因此，即使使用了较低的掩码率（如 10%）或没有掩码周围的所有补丁，它也会迫使网络学习相对较长的连接。

另一种增加预测距离的方法是结论 3：使用更大的掩码率，这也证明了相对较小的 patch 尺寸有利于微调性能。然而，这些较小的 patch 的总体精度不如较大的 patch (32) 的高。进一步将 patch 大小增加到 64 的观测精度下降，可能是由于预测距离太大。

▲ 图3. (a) 掩码率与 AvgDist 关系；(b) 最高的微调性能与 AvgDist 关系

上述观察和分析也可以很好地反映在一个新提出的 AvgDist 度量，该度量测量掩码像素到最近的可见像素的平均欧氏距离。不同掩码策略与不同掩蔽率的 AvgDist 如图 3(a) 所示。从图中可以看出，所有掩码策略的 AvgDist 都随着掩蔽率的增大而平滑增大。对于随机掩码策略，当补丁大小较小（如 4 或 8）时, AvgDist 相对较低，且随着掩码率的增加而增长缓慢。另一方面，当补丁尺寸较大时（如 64），很小的掩码率（如 10%）仍然会产生较大的 AvgDist。中心掩码和块掩码方法产生的 AvgDist 值与补丁大小为 64 的同样高。

图 3(b) 绘制了微调精度和 AvgDist 测量之间的关系，它遵循山脊形状。微调精度高的条目大致分布在 AvgDist 的 [10,20] 范围内，而 AvgDist 越小或越高的条目表现越差。这表明掩码图像建模中的预测距离应该适中，既不要太大，也不要太小。或许，在掩码预测中，小距离可能会让网络学习到太多的短连接，而大距离可能会让网络很难学习。这些结果也表明，AvgDist 可能是一个很好的指标用于检测掩码建模的有效性。

3.3 预测头

▲ 表2. 不同预测头的消融结果。简单的线性层以更低的训练成本表现得最好

表 2 给出了不同的预测头的影响。虽然通常复杂的预测头产生的损失略低，但在下游 ImageNet-1K 任务的迁移性能较低。这表明：

结论 1：较强的填涂能力并不一定意味着较好的下游性能。这可能是因为优异的性能很大一部分用于预测头，而不会用于下游任务。

结论 2：复杂的预测头带来较高的训练成本。作者证明了 SimMIM 的单个线性层头在微调度量下已经显示出具有竞争力甚至是最优的迁移性能。

这表明，结论 3：如果目标是学习好的特征来进行微调，那么在对比学习方法中对头部设计的重要探索对于掩码图像建模来说可能是不必要的。

3.4 预测分辨率

▲ 表3. 在不同的预测分辨率下的消融结果。中等大小的分辨率（不少于 1/16）都表现良好。

表 3 给出了目标分辨率变化的影响。它表明，在的低分辨率时准确率才会下降，可能是因为这个设置下丢弃了太多的信息。这些结果暗示了下游图像分类任务所需的信息粒度。

3.5 预测目标

▲ 表4. 不同预测目标的消融结果

表 4 比较不同预测目标的效果。结果表明，掩码图像建模的目标不需要与掩码语言建模的分类目标一致。设计与视觉信号本身的性质保持一致的方法效果很好。

▲ 表5. 对不同区域计算损失的消融结果。如果在掩码区域计算损失，则执行纯预测任务。如果对整个图像进行计算，则执行联合预测和重建任务

表 5 比较了两种方法，一种是预测默认设置下的掩码区域，另一种是重建被掩码区域和未被掩码区域。掩蔽区域的预测效果明显优于重建所有图像像素效果。这表明，这两种任务在内部机制上有根本的不同，预测任务可能是一种更有前景的表征学习方法。

3.6 与其他方法比较

▲ 表6. 与其他方法比较结果

表 6 比较了 SimMIM 与使用 ViTB 进行微调和线性探测的其他方法。SimMIM 通过微调达到了 83.8% 的 Top-1 准确率，比之前的最佳方法高出了 0.6%。SimMIM 由于其简单性，保留了最高的训练效率，比 DINO、MoCo v3、ViT 和 BEiT（不包括 dVAE 训练前的时间）分别高出 2.0×、1.8×、~ 4.0× 和 1.5×。

3.7 可伸缩性实验

作者采用不同模型尺寸的 Swin Transformer 进行实验，包括 Swin-B、Swin-L、SwinV2-H、SwinV2-G。为了减少实验开销，在训练前采用较小的的图像尺寸，所有模型都使用 ImageNet-1K 数据集进行训练，但 SwinV2-G 使用更大的 ImageNet-22K-ext 数据集。在微调时，采用较大的图像尺寸。对于 SwinV2-H，选用更大的分辨率。

表 7 列出了 SimMIM 在不同模型尺寸下的结果，并与有监督的同行进行了比较。在 SimMIM 预训练下，所有的 Swin-B、Swin-L、SwinV2-H 的准确率都显著高于有监督的对照组。另外，分辨率为的 SwinV2-H 模型在 ImageNet-1K 上的 top-1 精度达到了 87.1%，是仅使用 ImageNet-1K 数据的方法中精度最高的方法。