论文地址：https://arxiv.org/abs/2206.00790
https://github.com/junchen14/LoMaR
Efficient Self-supervised Vision Pretraining with Local Masked Reconstruction，比MAE快3.1倍，比BEiT快5.3倍！KAUST&南洋理工提出基于局部mask重建的高效自监督视觉预训练方法LoMaR，同时提高训练精度和效率！
计算机视觉的自监督学习取得了巨大的进步，改进了许多下游视觉任务，如图像分类、语义分割和目标检测。其中，MAE和BEiT等生成性自监督视觉学习方法表现出了良好的性能。然而，它们的全局掩蔽重建机制对计算的要求很高。
为了解决这个问题，作者提出了局部掩蔽重建（local masked reconstruction，LoMaR），这是一种简单而有效的方法，在一个简单的Transformer编码器上，在7×7块的小窗口内执行掩蔽重建，与整个图像的全局掩蔽重建相比，提高了效率和精度之间的权衡。
大量实验表明，LoMaR在ImageNet-1K分类中达到84.1%的top-1精度，优于MAE 0.5%。在384×384图像上对预训练后的LoMaR进行微调后，可以达到85.4%的top-1精度，超过MAE 0.6%。在MS COCO上，LoMaR在目标检测上比MAE好0.5，在实例分割上比MAE好0.5。LoMaR在预训练高分辨率图像上的计算效率尤其高，例如，在预训练448×448图像上，LoMaR比MAE快3.1倍，分类精度高0.2%。这种局部掩蔽重建学习机制可以很容易地集成到任何其他生成性自监督学习方法中。
本文提出了一种新的模型，称为局部掩蔽重建或LoMaR。该模型将注意力区域限制在一个小窗口内，如7×7的图像块，这足以进行重建。对于那些需要在长序列上操作的任务，在许多NLP领域中也可以看到类似的方法。在视觉领域也探索了小窗口，以提高训练和推理速度。但与之前的视觉Transformer（如Swin Transformer）不同，Swin Transformer为每个图像创建具有固定坐标的移动窗口。本文取而代之的是对几个随机位置的窗口进行采样，这样可以更好地捕获不同空间区域中的对象。
在下图中，作者比较了LoMaR和MAE，并注意到两个主要区别：a）本文对一个区域进行了k×k个patch采样，以进行掩蔽重建，而不是从全部patch中进行重建。作者发现，只需一些局部视觉线索，就足以恢复丢失的信息，而不是从图像中全局25%的可见patch重建遮罩patch。b） 本文将MAE中的重量级解码器替换为轻量级MLP头。将所有图像patch直接输入编码器，包括masked和visible patches。相比之下，在MAE中，只有可见的patch被馈送到编码器。实验表明，这些结构变化为小窗口的局部掩蔽重建带来了更大的性能增益。
经过广泛的实验，作者发现
- LoMaR在ImageNet-1K数据集上可以实现84.1 top-1 acc，比MAE高出0.5 acc。此外，LoMaR的性能可以进一步提高到84.3 acc，在ViT B/8主干上只需预训练400个阶段，与ViT B/16相比，这不会带来额外的预训练成本。在分辨率为384×384的图像上对预训练模型进行微调后，LoMaR可以达到85.4 acc，比MAE高出0.6 acc。
- LoMaR在高分辨率图像预训练中比其他baseline更有效，因为它的计算量对不同的图像分辨率是不变的。然而，其他方法的计算成本是图像分辨率增加的二次方，这导致了昂贵的预训练。比如，对于448×448图像的预训练，LoMaR比MAE快3.1倍，实现了更高的分类性能。
- LoMaR是一种高效的学习方法，可以很容易地集成到任何其他生成性自监督学习方法中。将本文的局部掩蔽重建学习机制安装到BEiT中可以将其ImageNet-1K分类性能从83.2提高到83.4，只消耗最初预训练时间的35.8%。LoMaR在其他任务（如目标检测）上也具有很强的泛化能力。在ViTDet的目标检测框架下，它比MAE的性能高出0.5 。
LoMaR依赖于一堆Transformer块，通过从与MAE类似的损坏图像中恢复缺失的patch来预训练大量未标记图像，但LoMaR在几个关键位置将其与MAE区分开来。上图并排比较了两者。在本节中，作者首先回顾MAE模型，然后描述LoMaR和MAE之间的差异。
Architecture
LoMaR采用了一种简单的编码器-编码器结构，而不是MAE的非对称编码器-解码器。作者将采样区域下所有可见和mask的patch输入编码器。虽然将mask patch输入编码器可能被认为是比仅将mask patch输入解码器的MAE效率更低的操作，但作者发现，在早期阶段输入mask patch可以增强视觉表现，并使其对较小的窗口大小更具鲁棒性。这可能是因为编码器可以在多个编码器层与其他可见patch交互后，将mask patch转换回其原始RGB表示。隐藏层中恢复的mask patch可以隐式地对图像表示作出贡献。因此，本文在LoMaR中保留mask patch作为编码器输入。
Relative positional encoding
LoMaR在MAE中应用相对位置编码（RPE）而不是绝对位置编码。作者应用了上下文RPE，在计算自注意时，它为每个查询i和键j引入了一个可学习的向量。
Implementation
给定一幅图像，首先将其划分为几个不重叠的patch。每个patch线性投影到嵌入中。作者在不同的空间位置随机抽取几个方形的K×K 个patch。然后，将每个窗口中固定百分比的patch归零。然后，将所有patch从每个窗口按顺序提供给编码器。编码器在自注意层中应用可学习的相对位置编码。作者用一个简单的MLP头将编码器输出的潜在表示转换回其原始特征维，然后用归一化的ground-truth图像计算均方误差。
自监督学习（SSL）可以从大量未标记数据的训练中获益。然而，在大规模的预训练下，它们的高计算要求仍然是一个值得关注的问题。在本文的研究中，作者观察到用于生成SSL的局部掩蔽重建（LoMaR）比MAE和BEiT等有影响力的著作使用的全局版本更有效。
LoMaR在图像分类、实例分割和目标检测方面具有良好的泛化能力；它可以很容易地合并到MAE和BEiT中。LoMaR有希望将SSL扩展到更大的数据集和更高的分辨率，以及计算更密集的数据集，如视频。LoMaR的另一个优点在于，当图像patch数量增加时，效率会提高。
主要原因是LoMaR限制了局部窗口内的自注意，其计算复杂度随每幅图像的采样窗口数呈线性增长。此特性可以在高图像分辨率下进行有效的预训练，而对于其他SSL方法来说，这将非常昂贵。它可以使许多视觉任务受益，例如需要在像素级进行密集预测的对象检测或实例分割。尽管LoMaR相对于其他高分辨率图像基线的预训练效率增益很高，但与MAE相比，LoMaR相对于低分辨率图像的效率提高有限。