MixMAE(MixMIM):用于分层视觉变压器有效预训练的混合和掩码自编码器 论文阅读

论文: MixMAE (arxiv.org)

代码: Sense-X/MixMIM: MixMIM: Mixed and Masked Image Modeling for Efficient Visual Representation Learning (github.com)

摘要: 本文提出MixMAE (Mixed and mask AutoEncoder)，这是一种简单而有效的预训练方法，适用于各种层次视觉变压器。现有的分层视觉变形的掩码图像建模(MIM)方法用一个特殊的[MASK]符号替换随机子集的输入标记，旨在从损坏的图像中重建原始图像标记。然而，我们发现使用[MASK]符号大大降低了训练速度，并导致预训练-微调不一致，这是由于屏蔽比很大(例如SimMIM中为60%)。另一方面，MAE根本没有在其编码器中引入[MASK]令牌，但不适用于分层视觉变压器。为了解决这个问题并加速分层模型的预训练，我们用另一幅图像的可见标记替换一幅图像的屏蔽标记，即创建混合图像。然后，我们对混合输入的两幅原始图像进行双重重建，显著提高了效率。虽然MixMAE可以应用于各种分层变压器，但本文探索使用具有大窗口尺寸的Swin Transformer并扩展到巨大的模型尺寸(达到600M参数)。实验结果表明，MixMAE可以有效地学习高质量的视觉表征。值得注意的是，使用swwin - b /W14的MixMAE通过600 epoch的预训练在ImageNet-1K上达到了85.1%的top-1准确率。此外，它在其他6个数据集上的传输性能表明，MixMAE比以前流行的MIM方法具有更好的FLOPs /性能权衡。代码可从https://github.com/Sense-X/MixMIM获得。

1. 介绍

以自我监督的方式利用未标记的视觉数据来学习表征是有趣但具有挑战性的。继BERT[14]在自然语言处理中，预训练MIM在学习各种下游视觉任务的视觉表示方面取得了巨大成功[4,21,44,50,51]，包括图像分类[13]、目标检测[31]、语义分割[56]、视频分类[20]和运动控制[50]。虽然这些最先进的方法[4,21]在vanilla Vision Transformer (ViT)上取得了优异的性能[16,45]，但如何有效地预训练分层ViT以在宽视觉任务上获得进一步的效率[9,12,33,36,46]仍然是一个悬而未决的问题。一般来说，现有的MIM方法用一个特殊的[MASK]符号替换一部分输入令牌，旨在恢复原始图像补丁[4,51]。然而，使用[MASK]符号会导致两个问题。一方面，预训练中使用的[MASK]符号在微调阶段没有出现，导致预训练与微调不一致[14]。另一方面，预训练网络在处理信息较少的[MASK]符号上浪费了大量的计算量，使得预训练过程效率低下。当使用较大的掩蔽比时，这些问题变得更加严重[4,21,44,51]。例如，在SimMIM[51]中，预训练时使用了60%的掩蔽率，即60%的输入令牌被[MASK]符号替换。因此，SimMIM需要相对较多的epoch(即800)进行预训练。此外，由于高掩蔽率导致预训练微调不一致，限制了SimMIM在下游任务上的性能。相比之下，MAE[21]通过丢弃编码器中的掩码令牌并仅在轻量级解码器中使用[MASK]符号而不存在上述问题。MAE采用香草ViT[16]作为编码器，它可以通过自注意操作高效地处理部分输入。然而，该设计也限制了MAE在分层ViTs上的应用，因为分层ViTs不能处理任意长度的1D令牌序列[36,46]。在这项工作中，我们提出了MixMAE，这是一种广义预训练方法，它利用了SimMIM[51]和MAE[51]的优点，同时避免了它们的局限性。特别地，给定训练集中的两个随机图像，MixMAE以随机混合蒙版为输入，生成混合图像，并训练分层ViT重构两幅原始图像，学习视觉表征。从一个图像的角度来看，不是用特殊的[MASK]符号替换图像的掩码符号，而是用另一个图像的可见符号替换掩码符号。MixMAE采用编码器-解码器设计。编码器是一个分层ViT，并对混合图像进行处理以获得两个部分掩码图像的隐藏表示。在解码之前，隐藏的表示是未混合的，并使用[MASK]令牌填充。根据MAE[21]，解码器是一个小的ViT来重建两个原始图像。我们在图1中演示了所建议的MixMAE。

图1所示。MixMAE概述在预训练中，用随机混合蒙版混合两幅图像，生成混合图像。MixMAE将混合后的图像作为输入，对两幅原始图像进行重构。在解码之前，标记嵌入是未混合的，并填充了两个原始图像的双重重建的掩码标记 

MixMAE可以广泛应用于预训练不同层次的vit，如Swin Transformer[36]、Twins[9]、PVT[46]等。由于使用了层次结构，我们可以很自然地将预训练的编码器应用于目标检测和语义分割任务。从经验上看，在相似的模型大小和FLOPs下，MixMAE在广泛的下游任务上始终优于BEiT[4]和MAE[21]，包括iNaturalist[23]和Places[55]上的图像分类，COCO[31]上的对象检测和实例分割，以及ADE20K上的语义分割[56]。通过在编码器中放弃使用[MASK]令牌，MixMAE在各种分层vit[9,36,46]上表现出比SimMIM[51]更好的预训练效率。

2. 相关的工作

受BERT[14]对蒙面语言建模的启发，蒙面图像建模(mask Image Modeling, MIM)成为视觉表征学习的热门借口任务[2,4,21]。MIM旨在从损坏的输入中重建掩码令牌。目前的MIM方法可以根据重建目标分为两类。SimMIM[51]指出随机掩码patch的原始像素值是一个很好的重建目标，一个轻量级的预测头就足以进行预训练。与SimMIM不同的是，MAE[21]只将可见patch作为编码器的输入。在编码器和解码器的中间添加掩码令牌。这种非对称设计大大减少了编码器的计算开销。为了进一步增强编码器的特征提取能力，CAE[7]通过在编码器和解码器中间添加特征对齐模块，将编码器和解码器显式分离。Jean-Baptiste等[1]提出通过合成混合视频重构原始视频来学习表征。代替人工构建重建目标，利用网络生成重建目标也得到了广泛的应用。在这类作品中，图像标记器用于将图像转换为视觉标记。BEiT[4]使用预训练的离散VAE (dVAE)[40,41]作为标记器。然而，最初在dVAE中使用的MSE损失不足以迫使标记器捕获高级语义。PeCo[15]提出将感知相似度损失应用于dVAE的训练，可以驱动标记器生成更好的语义视觉标记，有助于预训练。此外，贝叶斯[4]中的标记器需要离线预训练，这限制了模型的自适应能力。为了解决这个问题，iBOT[57]提出使用在线标记器来生成视觉标记。同时也有研究在分层视觉变压器上使用MAE的工作。UM-MAE[28]提出了一种新的掩蔽策略，用于适应MAE预训练基于金字塔的vit(如PVT [46]， Swin[36])。GreenMIM[25]也将MAE应用于层次结构。它将局部窗口划分为几个大小相等的组，并提出了一种最优分组算法来找到最优组大小。而不是设计特定的掩蔽或分组策略，我们专注于重新安排输入和目标。实证评价表明，MixMAE在不同层次结构下都能获得更好的性能。我们的MixMAE还受到其他基于混合的训练技术的启发[34,52,54]。与那些为监督学习进行数据增强的工作相比，MixMAE可以在没有人工注释的情况下有效地预训练分层变压器。

3. 方法

在本节中，我们将介绍MixMAE，用于通过蒙面图像建模(MIM)学习视觉表示。我们首先简要回顾一下MIM，然后介绍MixMAE如何创建训练输入并执行图像重建，以及分层Transformer架构。最后，我们提出了如何降低借口任务的难度来提高预训练效率。

3.1. A Revisit of Masked Image Modeling  蒙面图像建模的再探讨

继BERT[14]之后，最近的研究[4,21,51]提出了用于学习视觉表征的MIM。给定输入图像x, MIM首先将图像划分为不重叠的图像块x p，按照ViT划分[16]。然后，它对一个随机掩码M进行采样，以掩码的一部分图像补丁，并用一个特殊符号[mask]填充被掩码的地方，即:´x p = x p⊙M + [mask]⊙(1−M)，其中⊙表示元素乘法。掩码图像x p由图像编码器处理以产生潜在表示并利用轻量级解码器(头部)基于潜在表征重构原始图像。重建目标可以选择归一化原始像素[21,51]或视觉标记[4,15]。MIM计算重建图像斑