掩蔽自动编码器

论文链接

核心概念

掩蔽自动编码器（MAE）的核心思想就是删除部分数据并让模型学习剩下的数据去预测被删除的数据。

MAE从输入图像中屏蔽随机补丁，并在像素空间中重建丢失的补丁。它具有非对称编码器设计。我们的编码器只对可见的补丁子集（不带掩码令牌）进行操作，我们的解码器是轻量级的，可以从潜在表示和掩码令牌重建输入。在我们的非对称编码器-解码器中，将掩码令牌转移到小型解码器会大大减少计算量。在这种设计下，非常高的掩蔽率（例如75%）可以实现双赢方案：它优化了精度，同时允许编码器仅处理一小部分（例如25%）补丁。这可以将总体预训练时间减少3倍或更多，同样可以减少内存消耗，使我们能够轻松地将MAE扩展到大型模型。

方法原理

掩蔽自动编码器（MAE）是一种简单的自动编码方法，它根据原始信号的部分观测结果重建原始信号。像所有的自动编码器一样，我们的方法有一个编码器将观察到的信号映射到潜在表示，还有一个解码器从潜在表示重建原始信号。与经典的自动编码器不同，我们采用了一种非对称设计，允许编码器仅对部分观察信号（无掩码标记）进行操作，并采用一种轻量级解码器，从潜在表示和掩码标记重建完整信号。说明了下面介绍的想法。

蒙版

按照ViT，我们将图像划分为规则的非重叠面片。然后，我们对一个子集的面片进行采样，并遮罩（即移除）剩余的面片。我们的采样策略很简单：我们在不替换的情况下，按照均匀分布对随机斑块进行采样。我们简单地称之为“随机抽样”。

具有高掩蔽率（即移除的面片的比率）的随机采样在很大程度上消除了冗余，因此创建了一个无法通过从可见相邻面片外推来轻松解决的任务。均匀分布可防止潜在的中心偏移（即，图像中心附近有更多遮罩面片）。最后，高度稀疏的输入为设计高效编码器创造了机会，下面介绍。

MAE编码器

我们的编码器是ViT，但仅适用于可见的、无遮罩的补丁。就像在标准ViT中一样，我们的编码器通过添加位置嵌入的线性投影嵌入面片，然后通过一系列变换块处理结果集。然而，我们的编码器只在整个集合的一小部分（例如25%）上运行。去除遮罩贴片；不使用掩码令牌。这使我们能够用一小部分的计算和内存来训练非常大的编码器。整个集合由轻量级解码器处理，如下所述。

MAE解码器

MAE解码器的输入是由（i）编码可见补丁和（ii）掩码令牌组成的全套令牌。每个掩码标记都是一个共享的学习向量，表示存在要预测的缺失面片。我们向这个完整集合中的所有标记添加位置嵌入；如果没有这一点，掩码令牌将没有关于其在映像中位置的信息。解码器具有另一系列转换器块。

MAE解码器仅在预训练期间用于执行图像重建任务（仅编码器用于生成用于识别的图像表示）。因此，解码器架构可以以独立于编码器设计的方式灵活地设计。我们用非常小的解码器进行实验，比编码器更窄、更浅。例如，与编码器相比，我们的默认解码器每个令牌的计算量为10%。通过这种非对称设计，全套令牌仅由轻量级解码器处理，这大大减少了预训练时间。

重建目标

我们的MAE通过预测每个遮罩面片的像素值来重建输入。解码器输出中的每个元素都是表示面片的像素值向量。解码器的最后一层是一个线性投影，其输出通道的数量等于面片中像素值的数量。解码器的输出被重塑以形成重构图像。我们的损失函数计算像素空间中重建图像和原始图像之间的均方误差（MSE）。我们只计算遮罩面片的损失，类似于BERT。

我们还研究了一种变体，其重建目标是每个遮罩面片的归一化像素值。具体来说，我们计算一个面片中所有像素的平均值和标准偏差，并使用它们来规范化该面片。在我们的实验中，使用归一化像素作为重建目标提高了表示质量。

简单的实现

我们的MAE预培训可以有效实施，而且重要的是，不需要任何专门的稀疏操作。首先，我们为每个输入面片生成一个标记（通过添加位置嵌入的线性投影）。接下来，我们随机洗牌令牌列表，并根据掩蔽率删除列表的最后一部分。此过程为编码器生成一小部分令牌，相当于在不替换的情况下采样补丁。编码后，我们将掩码标记列表附加到已编码补丁列表中，并取消填充此完整列表（反转随机洗牌操作），以将所有标记与其目标对齐。解码器应用于这个完整列表（添加了位置嵌入）。如前所述，不需要稀疏操作。这个简单的实现引入了可忽略不计的开销，因为洗牌和取消洗牌操作速度很快。

实验结果

ImageNet实验

掩蔽比

下图显示了掩蔽率的影响。最佳比率出人意料地高。75%的比率对线性探测和微调都有好处。这种行为与典型掩蔽率为15%的BERT形成对比。我们的掩蔽率也远高于计算机视觉相关工作中的掩蔽率（20%至50%）

该模型推断缺失的补丁会产生不同但似乎合理的输出（图4）。它可以理解物体和场景的格式塔，而这不能简单地通过延伸线条或纹理来完成。我们假设这种类似推理的行为与学习有用的表征有关。

下图还显示了线性探测和微调结果遵循不同的趋势。对于线性探测，精度随着掩蔽率的增加而稳定增加，直到达到最低点：精度差距达到∼ 20%（54.6%对73.5%）。对于微调，结果对比率的敏感度较低，并且大范围的掩蔽比率（40–80%）工作良好。下图中的所有微调结果都优于从头开始的培训（82.5%）。

解码器设计

我们的MAE解码器可以灵活设计，如表1a和1b所示。

表1a改变了解码器深度（变压器块的数量）。足够深的解码器对于线性探测非常重要。这可以用像素重建任务和识别任务之间的差距来解释，即自动编码器中的最后几层更专门用于重建，但与识别无关。一个合理的深度解码器可以解释重建的特殊性，使潜在的表示处于更抽象的层次。这种设计可以使线性探测提高8%（表1a，“lin”）。但是，如果使用微调，编码器的最后一层可以进行调整以适应识别任务。解码器深度对改善微调的影响较小（表1a，“ft”）。有趣的是，我们使用单块解码器的MAE可以在微调（84.8%）时表现出色。请注意，单个转换器块是将信息从可见令牌传播到掩码令牌的最低要求。这种小型解码器可以进一步加快训练速度。

在表1b中，我们研究了解码器宽度（信道数）。我们默认使用512-d，它在微调和线性探测下表现良好。更窄的解码器也可以很好地进行微调。

总的来说，我们默认的MAE解码器是轻量级的。它有8个块，宽度为512-d（表1中为灰色）。与ViT-L（24个块，1024-d）相比，每个令牌只有9%的触发器。因此，虽然解码器处理所有令牌，但它仍然是整个计算的一小部分。

面具令牌

我们MAE的一个重要设计是跳过编码器中的掩码令牌，然后将其应用到轻量级解码器中。表1c研究了该设计。

如果编码器使用掩码标记，则性能更差：在线性探测中，其精度下降14%。在这种情况下，预训练和部署之间存在差距：此编码器在预训练的输入中有很大一部分掩码令牌，而未损坏的图像中不存在这些令牌。此间隙可能会降低部署的准确性。通过从编码器中移除掩码标记，我们约束编码器始终看到真实的面片，从而提高精度。

此外，通过跳过编码器中的掩码令牌，我们大大减少了训练计算。在表1c中，我们将整体训练失败次数减少了3.3次。在我们的实现中，这导致了2.8倍的挂钟加速（见表2）。对于较小的解码器（1块）、较大的编码器（ViT-H）或两者兼而有之，墙上的时钟加速比甚至更大（3.5–4.1×）。请注意，当掩蔽率为75%时，加速比可以大于4倍，部分原因是自我注意复杂性是二次的。此外，内存大大减少，这可以使训练更大的模型或通过大批量训练加快速度。时间和内存效率使我们的MAE有利于训练非常大的模型。

重建目标

我们比较了表1d中的不同重建目标。到目前为止，我们的结果是基于没有（每个面片）标准化的像素。使用标准化像素可以提高精度。这种逐块标准化增强了局部对比度。在另一个变体中，我们在面片空间中执行PCA，并使用最大的PCA系数（此处为96）作为目标。这样做会降低准确性。两个实验都表明高频分量在我们的方法中是有用的。

我们还比较了预测代币的MAE变体，即BEiT中使用的目标。特别是对于这种变体，我们使用DALLE预先训练的dVAE作为标记器，如下。这里，MAE解码器使用交叉熵损失预测令牌索引。与未标准化像素相比，这种标记化将微调精度提高了0.4%，但与标准化像素相比没有优势。它还降低了线性探测精度。在§5中，我们进一步表明，标记化在迁移学习中是不必要的。

我们基于像素的MAE比标记化简单得多。dVAE标记器还需要一个预训练阶段，这可能取决于额外的数据（250M图像）。dVAE编码器是一个大型卷积网络（ViT-L的40%触发器），增加了不寻常的开销。使用像素不会遇到这些问题。

数据扩充

表1e研究了数据扩充对我们MAE预培训的影响。

我们的MAE使用仅裁剪的增强效果很好，无论是固定大小还是随机大小（都具有随机水平翻转）。添加颜色抖动会降低结果，因此我们不会在其他实验中使用它。

令人惊讶的是，即使不使用数据扩充（只使用中心裁剪，不使用翻转），我们的MAE也能正常工作。这一特性与对比学习及相关方法截然不同，后者严重依赖于数据扩充。有人观察到，对于BYOL和SimCLR，仅使用裁剪增强将准确度分别降低13%和28%。此外，没有证据表明对比学习可以在没有增强的情况下工作：图像的两个视图是相同的，很容易满足一个简单的解决方案。

在MAE中，数据增强的作用主要通过随机掩蔽（下一步烧蚀）来实现。每个迭代的掩码都不同，因此它们生成新的训练样本，而不考虑数据的增加。借口任务由于掩蔽而变得困难，并且需要较少的强化以使训练正规化。

掩模采样策略

在表1f中，我们比较了不同的掩模取样策略，如图6所示。

我们的MAE采用分块掩蔽，在50%的比率下工作良好，但在75%的比率下降解。这项任务比随机抽样更难，因为观察到更高的训练损失。重建也更加模糊。

我们还研究了网格采样，它定期保持每四个补丁中的一个。这是一项更简单的任务，训练损失也更低。重建过程更为尖锐。然而，表示质量较低。

简单随机抽样最适合我们的MAE。它允许更高的掩蔽比，这提供了更大的加速优势，同时也享受良好的准确性。

培训时间表

到目前为止，我们的烧蚀是基于800个历元的预训练。图7显示了培训计划长度的影响。随着训练时间的延长，准确度稳步提高。事实上，我们甚至在1600个时代也没有观察到线性探测精度的饱和。这种行为不同于对比学习方法，例如，MoCo v3[9]在ViT-L的300个历元处饱和。请注意，MAE编码器在每个历元中仅看到25%的补丁，而在对比学习中，编码器在每个历元中看到200%（两个裁剪）甚至更多（多裁剪）补丁。

实验结果对比

与自我监督方法的比较

在表3中，我们比较了自监督ViT模型的微调结果。对于ViT-B，所有方法的性能都很好。对于ViT-L，方法之间的差距更大，这表明大型模型面临的挑战是减少过度拟合。

我们的MAE可以很容易地扩展，并且与更大的型号相比已经显示出稳定的改进。我们使用ViT-H（224号）获得86.9%的准确率。通过使用448大小进行微调，我们仅使用1K数据就实现了87.8%的精度。基于高级网络，在所有仅使用1K数据的方法中，先前的最佳精度为87.1%（512大小）。我们在极具竞争力的基准IN1K（无外部数据）中以不平凡的优势超越了最先进的技术。我们的结果基于vanilla ViT，我们预计高级网络将表现得更好。

与BEiT相比，我们的MAE更精确，同时更简单、更快。与预测标记的BEiT相比，我们的方法重建像素：当使用ViT-B重建像素时，BEiT报告了1.8%的退化。我们不需要dVAE预训练。此外，我们的MAE比BEiT快得多（每个历元3.5倍），原因如表1c所示。

表3中的MAE模型经过1600个时代的预训练，以获得更高的精度（图7）。即使如此，如果在相同的硬件中进行培训，我们的总预培训时间也比所有其他方法要短。例如，对于ViT-L，我们的MAE的训练时间对于1600个历元是31小时，而MoCo v3的训练时间对于300个历元是36小时，使用相同的128个TPU-v3核。

与有监督的预培训进行比较

在最初的ViT论文中，当在1K中训练时，ViT-L会退化。参见图8。我们改进的监督配方对于从头开始的训练效果更好，但精确度已经饱和。

我们的MAE预培训仅使用1K，可以更好地推广：对于容量更大的模型，从头开始培训的收益更大。这一趋势与中的JFT-300M监督预训类似。这一比较表明，我们的MAE可以帮助扩大模型尺寸。

迁移学习实验

我们评估了迁移学习在COCO上的目标检测和分割以及在ADE20K上的语义分割。我们使用表3中预先训练的模型。

目标检测和分割

我们在COCO上对Mask R-CNN进行端到端的微调。ViT主干适合与FPN一起使用。我们将此对象检测系统应用于表4中的所有条目。我们报告用于对象检测的box AP和用于实例分割的mask AP。

与有监督的预培训相比，我们的MAE在所有配置下都表现得更好（表4）。对于较小的ViT-B，我们的MAE比有监督的预训练高2.4分（50.3比47.9，AP框）。更重要的是，随着ViT-L的增大，我们的MAE预训练比有监督的预训练高出4.0分（53.3比49.3）。

基于像素的MAE优于或等同于基于令牌的BEIT，而MAE则更简单和更快。MAE和BEiT均优于MOCO V3，MOCO V3与监督前培训一致。

语义分割

我们在ADE20K上的实验按照[2]中的代码使用SuperNet。表5显示，我们的MAE显著改善了ViT-L的转移结果，比有监督的训练前对应者（53.6比49.9）好3.7分。基于像素的MAE优于基于令牌的BEiT。这些观察结果与COCO中的一致。

结论

扩展性好的简单算法是深度学习的核心。在NLP中，简单的自监督学习方法可以从指数缩放模型中获益。在计算机视觉中，尽管在自我监督学习方面取得了进展，但实际的预培训模式仍主要受到监督。在这项研究中，我们在ImageNet和转移学习中观察到，自动编码器——一种简单的自我监督方法，类似于NLP中的技术——提供了可扩展的优势。视觉中的自我监督学习现在可能正走上与NLP类似的轨道。

另一方面，我们注意到，图像和语言是不同性质的信号，必须认真处理这种差异。图像只是光照记录，没有语义分解成文字的视觉模拟。我们不尝试删除对象，而是删除最有可能不形成语义段的随机补丁。同样，我们的MAE重建像素，而像素不是语义实体。然而，我们观察到（如图4），我们的MAE推断出复杂的整体重建，表明它已经学习了许多视觉概念，即语义。我们假设这种行为是通过MAE内部丰富的隐藏表征发生的。我们希望这一观点将启发今后的工作。

更广泛的影响：提出的方法基于培训数据集的学习统计预测内容，从而反映这些数据中的偏见，包括具有负面社会影响的偏见。模型可能会生成不存在的内容。这些问题需要进一步的研究和考虑时，在这项工作的基础上生成图像。