代码：https://github.com/implus/UM-MAE

论文：https://arxiv.org/abs/2205.10063

 

导读

最近，Masked AutoEncoder (MAE) 通过非对称编解码器设计引领了视觉自监督领域的研究趋势，显著优化了预训练效率和微调精度。值得注意的是，非对称结构的成功依赖于 Vanilla Vision Transformer (ViT) 的“全局”属性，即其自注意力机制可以解释离散图像块的任意子集。然而，目前还不清楚如何在 MAE 预训练中采用先进的基于 Pyramid 的 ViT（例如 PVT、Swin），因为它们通常在“局部”窗口中引入算子，从而难以处理部分视觉的随机序列token。在本文中，作者提出了Uniform Masking（UM），成功地为基于 Pyramid 的具有局部性的 ViT 进行 MAE 预训练（简称为“UM-MAE”）。具体来说，UM 包括一个统一采样（US），它从每个 2×2 网格中严格采样 1 个随机patch，以及一个二次masking（SM），它随机掩蔽一部分（通常是 25%）已经采样的区域作为可学习的标记。US 在多个非重叠本地窗口中保留等效元素，从而顺利支持流行的基于 Pyramid 的 ViT；而 SM 是为更好的可转移视觉表示而设计的，因为 US 降低了阻碍语义学习的像素恢复预任务的难度。

 

贡献

自监督学习 (SSL) 使用辅助任务从大规模未标记数据中挖掘不依赖标签的自主监督信息，并学习可转移到下游任务的表示。 最初，SSL 通过 GPT 和 BERT中的“掩码自动编码”方案在自然语言处理 (NLP) 领域显示出巨大潜力。通过学习根据上下文推断数据的删除部分，这些技术创新了 NLP 研究的新范式。

受 BERT 的启发，计算机视觉领域近期对自监督学习也产生了极大的兴趣，其中 Masked AutoEncoder (MAE)成为最具代表性的方法之一，它显著优化了预训练阶段的效率和微调阶段的精度，成为了引领SSL在视觉领域的新研究趋势。MAE 中最具影响力的设计之一是非对称编码器-解码器架构。 与接收所有输入序列的解码器部分不同，编码器部分仅将可见图像补丁（通常仅占总数的 25%）作为输入。 有趣的是，这样的设计不仅显着降低了预训练的复杂度，而且确保了下游任务上的微调性能。

值得注意的是，非对称结构的成功依赖于 Vanilla Vision Transformer (ViT) 的“全局”属性，即其自注意力机制可以对序列化图像块的任意子集进行推理（见图 1（a））。然而，Vanilla ViT 的“全局”属性是一把双刃剑：当将 Vanilla ViT 转移到具有相当大输入分辨率的下游视觉任务（例如，在目标检测中约为 10002）时，它的存储需求对于 GPU 来说是一个挑战。尽管 ViTDet 试图在微调期间部分限制 Vanilla ViT 的部分局部窗口自注意力，但其最优架构仍未知。因此，为了解决上述问题，基于 Pyramid 的 ViTs成为了首选方案，然而，如何针对Pyramid ViT进行MAE中的非对称结构训练仍是一个问题，因为它们通常在“局部”窗口中引入算子，从而难以处理部分视觉的随机序列token。

为解决这一问题，在本文中，作者提出了Uniform Masking（UM），成功地为基于 Pyramid 的具有局部性的 ViT 进行 MAE 预训练（简称为“UM-MAE”）。具体来说，UM 包括一个统一采样（US），它从每个 2×2 网格中严格采样 1 个随机patch，以及一个二次masking（SM），它随机掩蔽一部分（通常是 25%）已经采样的区域作为可学习的标记。US 在多个非重叠本地窗口中保留等效元素，从而顺利支持流行的基于 Pyramid 的 ViT；而 SM 是为更好的可转移视觉表示而设计的，因为 US 降低了阻碍语义学习的像素恢复预任务的难度。

 

方法

本文提出的Uniform Masking能够将密集的图像Token转化为稀疏Token序列，同时确保它们的一致分布，这一方法包含了两个阶段：Uniform Sampling、Secondary Masking。

 

Uniform Sampling

如图1（d）所示，UM首先使用Uniform Sampling将输入图像块进行采样，其采样策略可以概括为：约束每个2*2窗口内均包含一个可见的输入图像块，窗口中其他三个部分均进行掩盖。

这一策略使得UM可以与诸多Pyramid ViT进行结合。

如图2所示，PVT引入了Spatial-Reduction Window(SRW) 以降低原始自注意力块的复杂性。从具有均匀采样块的输入图像开始，作者展示了 PVT 的前三个阶段，并在上层管道中将其典型的空间缩减超参数依次标记为 {8, 4, 2}。作者证明，对于 PVT，通过同时将原始可见输入重组为其紧凑形式并相应地将空间缩减窗口（即 {4,2,1}）的边缘大小减半，有效元素在相应的局部窗口之间是等效的两条管道。因此，US与PVT兼容，当使用重组后的紧凑2D图像作为输入时，编码器部分的输入元素减少了75%。

Swin 与 PVT 有两个主要区别：（1）Local self Attention Window (LSAW) 大小通常跨阶段固定； (2) 它在连续的块中具有局部自我注意窗口 (SLSAW) 的偏移设置，它引入了相邻非重叠窗口之间的连接。 基于这些差异，作者推断 Swin 在预训练时对窗口大小和输入图像尺度的选择有更多限制：需要 16 · 2n ×16 · 2n(n ∈ N) 大小的窗口（和输入图像） 以确保在考虑移位偏移为 8 · 2n × 8 · 2n 的移位情况时的等价性，如图 3 所示。在上述约束下，有效元素在两条管道上的对应（移位）局部窗口之间是等价的 ，类似于 PVT。

 

Secondary Masking

为了解决 US 带来的退化问题，本文进一步提出了 Secondary Masking (SM) 策略，该策略在 US 已经采样的可见块中执行二级随机掩码，如图 4 (c) 至 (d) 所示。 与 US 阶段完全丢弃掩码补丁不同，SM 使用共享掩码令牌保留掩码补丁，以兼容具有局部性的基于金字塔的 ViT（参见第 3.1 节）。 因此，简单的操作增加了语义恢复预任务的难度，它使网络专注于学习高质量的表示不完全的视觉上下文，没有严重依赖相邻的低级像素。

 

UM-MAE模型

上图展示了UM方法应用于Pyramid ViT（例如 PVT 和 Swin）时的模型设计（即 MAE 式管道）。

 

实验

如下图所示，作者首先证明了 UM-MAE 显著提高了基于 Pyramid 的 ViT 的预训练效率（例如，它加速并将 GPU 内存减少了 ~2 倍），但在下游任务中保持了具有竞争力的微调性能。例如，使用 HTC++ 检测器，仅在 ImageNet-1K 中在 UM-MAE 下自监督的预训练 Swin-Large 主干甚至可以胜过在 ImageNet-22K 中监督的主干。

根据SimMIM[42]的设置，作者对UM-MAE预训练了800个epoc用，并在IN1K上进行了100个epoch的微调。表4报告了top-1的精度，表明本文的UM-MAE在大规模模型上保持了竞争性能。

.表5表明，仅在IN1K中在UM-MAE下预先训练的主干甚至可以超过在IN22K中自我监督的主干，仅使用基线训练时期的一半。

在图8中，我们观察到它们都大致恢复了接近原始图像的语义细节，而SimMIM的结果有时可能过于平滑，例如青蛙的右长腿。

内容中包含的图片若涉及版权问题，请及时与我们联系删除