[论文精读]Masked Autoencoders are scalable Vision Learners

摘要本文证明了掩码自编码器(MAE)是一种可扩展的计算机视觉自监督学习算法。我们的
MAE方法很简单:我们盖住输入图像的随机块并重建缺失的像素。它基于两个核心设计。首先，我们开发了一个非对称编码器-解码器架构，其中一个编码器仅对块的可见子集(没有掩码标记)进行操作，以及一个轻量级解码器，该解码器从潜在表示和掩码标记重建原始图像。其次，我们发现如果用比较高的掩盖比例掩盖输入图像，例如75%，这会产生一个重要的和有意义的自我监督任务。这两种设计的结合使我们能够高效地训练大型模型:我们加速了训练(3倍或更多)并提高了准确性。我们的可扩展方法允许学习泛化良好的高容量模型:例如，在仅使用ImageNet-1K数据的方法中，原始的ViT-Huge模型实现了最佳的准确率(87.8%)。（注意，这里ViT原论文有提到自监督的方法去训练，但是作者尝试后发现效果不好，因此没有展开，还是推荐用有标号的更大的数据集训练，而这里是说用较小的数据集，然后自监督训练可以达到和较大的数据集接近的性能）。在下游任务的迁移性能优于监督预训练，并显示出有希望的扩展行为。
1.简介

深度学习见证了一个模型架构的不断爆炸式增长的能力和容量。在硬件快速发展的帮助下，今天的模型可以很容易地过拟合一百万张图像，并开始需要数亿张通常是公开无法访问的标记图像。这种对数据的需求已经在自然语言处理（NLP）中成功地解决了自监督预训练的问题。这些解决方案基于GPT中的自回归语言建模和BERT 中的掩码自编码，其在概念上很简单：删除部分数据，并学习预测被删除的内容。这些方法现在可以训练包含超过1000亿个参数的广义NLP模型。掩码自编码器的思想，是一种更一般的去噪自编码器的形式，是自然的，也适用于计算机视觉。事实上，相关的视觉领域的研究在BERT之前就有。然而，尽管随着BERT的成功，人们对这一想法产生了极大的兴趣，但在视觉领域的自编码方法的进展仍落后于NLP。我们疑惑：是什么使得掩码自编码器在视觉和语言之间有所不同？我们试图从以下几个角度来回答这个问题：
（i）直到最近，架构都是不同的。在视觉中，卷积网络在过去十年中占主导地位。卷积通常在常规网格上运行，要将掩码标记或位置嵌入这些“指示器”集成到卷积网络中并不简单。然而，随着ViT的引入，这一架构差距已经得到了解决，不应该再构成障碍。
（ii）语言和视觉之间的信息密度是不同的.语言是由人类产生的信号，具有高度的语义性和信息密集性。当训练一个模型只预测每个句子中缺失的一些单词时，这项任务似乎能诱导复杂的语言理解。相反，图像是具有大量空间冗余的自然信号——例如，一个缺失的块可以从相邻的块中恢复，只有很少的对部分、目标和场景的高级理解。为了克服这种差异并鼓励学习有用的特性，我们展示了一个简单的在计算机视觉中表现很好的策略：屏蔽非常多的随机块。这种策略在很大程度上减少了冗余，并创造了一个具有挑战性的自监督任务，需要对低级图像统计之外的整体理解。

（iii）自编码器的解码器，它将潜在的表示形式映射回输入端，在重建文本和图像之间扮演着不同的角色。在视觉中，解码器重建像素，因此其输出的语义级别低于常见的识别任务。这与语言相反，在语言中，解码器预测包含丰富语义信息的缺失单词。虽然在BERT中，解码器可以是比较简单的（一个MLP），但我们发现，对于图像，解码器的设计在确定学习到的潜在表示的语义水平中起着关键作用。
在此分析的驱动下，我们提出了一种简单、有效的、可扩展的掩码自编码器（MAE）形式，用于视觉表示学习。我们的MAE从输入图像中屏蔽随机块，并重建像素空间中缺失的块。它具有非对称编码解码器的设计。我们的编码器只操作可见的块子集，并且我们的解码器是轻量级的，并与掩码标记一起重构来自潜在表示的输入。在我们的非对称编解码器中，将掩码标记转移到小解码器会大大减少计算。在这种设计下，一个非常高的掩蔽比（例如，75%）可以实现一个双赢的场景：它优化了精度，同时允许编码器只处理一小部分（例如，25%）的块。这可以减少3倍或更多的预训练时间，同时可以减少内存消耗，使我们能够轻松地将MAE扩展到大型模型。
我们的MAE学习了非常高容量的模型，它们可以推广得很好。通过MAE预训练，我们可以在ImageNet-1K上训练像ViT-Large/-Huge 这样的需要大量数据的模型，从而提高泛化性能。使用一个普通的ViT-Huge模型，当我们在ImageNet-1K上进行微调后，我们达到了87.8%的精度。这优于之前所有只使用ImageNet-1K数据的结果。我们还评估了在目标检测，实例分割，和语义分割上的迁移学习的效果。在这些任务中，我们的预训练比有监督的预训练对手获得了更好的结果，更重要的是，我们通过扩大模型观察到显著的收益。这些观察结果与在NLP 的自监督预训练中所看到的观察结果一致，我们希望它们能够使我们的领域能够探索类似的轨迹。

2.相关工作
Masked language modeling和其自回归对应的方法，如BERT 和GPT，则是NLP中非常成功的预训练方法。这些方法包含了部分输入序列，并训练模型来预测缺失的内容。这些方法已经被证明具有很好的扩展性，并且大量的证据表明，这些预先训练好的表示可以很好地推广到各种下游任务。
Autoencoding是学习表示法的一种经典方法。它有一个编码器，将输入映射到潜在表示，和一个重构输入的解码器。例如，PCA和k-means是自编码器。去噪自动编码器（DAE）是一类自编码器，它会破坏一个输入信号，并学习如何重建原始的、未损坏的信号。一系列的方法可以被认为是在不同的损坏下的广义DAE，例如，掩蔽像素或去除颜色通道。我们的MAE是去噪自编码的一种形式，但在许多方面与经典的DAE不同。
Masked image encoding方法从被掩蔽破坏的图像中学习表示。之前的开创性工作将掩蔽作为DAE中的一种噪声类型。上下文编码器使用卷积网络来绘制大的缺失区域。由于NLP的成功，最近的相关方法都是基于Transformers的。iGPT 处理像素序列并预测未知像素。ViT论文研究了自监督学习中的掩码块预测。最近，BEiT提出了预测分离标记。
self-supervised learning方法对计算机视觉产生了极大的兴趣，通常集中在训练前的任务上。最近，对比学习非常流行，它建模两个或多个视图之间的图像相似性和不相似性（或仅为相似性）。对比方法和相关方法强烈依赖于数据增强。自编码追求一个概念上不同的方向，它表现出我们将呈现的不同行为。
3.方法
我们的掩码自动编码器（MAE）是一种简单的自动编码方法，它可以重建原始信号的部分观察结果。
像所有的自动编码器一样，我们的方法有一个编码器，它将观察到的信号映射到一个潜在的表示，此外还有一个解码器，它将潜在的表示重构为原始信号。与经典的自动编码器不同，我们采用了一种非对称的设计，允许编码器只操作部分的、观察到的信号（没有掩码标记），以及一个轻量级的解码器，从潜在的表示和掩码标记重构完整的信号。
**遮挡。**和ViT 一样，我们将一个图像分割成规则的不重叠的patch。然后，我们对剩下的patch和掩码（即移除）的子集进行采样。我们的采样策略很简单：我们对随机patch进行采样而不进行替换，并遵循均匀分布。我们只是把其称为“随机抽样”。
高掩蔽比（即被删除的块的比例）的随机采样在很大程度上消除了冗余，从而创建了一个不容易通过从可见的相邻块进行外推来解决的任务。均匀分布防止了潜在的中心偏差（即，在图像中心附近有更多的掩蔽块）。最后，高度稀疏的输入为设计一个有效的编码器创造了机会，接下来将介绍。
**MAE encoder.**我们的编码器是一个只应用于可见的，未屏蔽的块的ViT。就像在一个标准的ViT中一样，我们的编码器通过一个添加了位置嵌入的线性投影来嵌入块，然后通过一系列的Transformer块来处理结果集。然而，我们的编码器只操作于整个集合的一个子集（例如，25%）。删除屏蔽块；不使用掩码标记。这允许我们只用一小部分的计算和内存来训练非常大的编码器。完整的集合由一个轻量级解码器处理，下面将进行描述。
MAE decoder. MAE解码器的输入是完整的标记集合，包括编码的可见块和掩码标记。每个掩码标记是一个共享的、可学习的向量，表明存在一个需要预测的遗漏块。我们向这个完整集中的所有标记添加位置嵌入；没有这个，掩码标记就没有关于它们在图像中所处的位置的信息。该解码器由另一系列的Transformer块构成。
MAE解码器仅在预训练的时候用于执行图像重建任务（只有编码器用于产生要进行识别的图像表示）。因此，解码器架构可以以独立于编码器设计的方式灵活设计。我们实验使用非常小的解码器，比编码器更窄、更浅。例如，我们的默认解码器相比编码器对于每个标记的计算量小于10%。通过这种不对称的设计，完整的标记集只由轻量级的解码器来处理，这大大减少了预训练的时间。
**重建目标。**我们的MAE通过预测每个掩蔽块的像素值来重建输入。解码器输出中的每个元素都是一个代表一个块的像素值的向量。解码器的最后一层是一个线性映射，其输出通道的数量等于一个块中的像素值的数量。解码器的输出被重塑从而形成一个重建的图像。我们的损失函数计算在像素空间中重建图像和原始图像之间的均方误差（MSE）。我们只计算了掩蔽块上的损失，类似于BERT 。
我们还研究了一个变体，其重建目标是每个掩蔽块的归一化像素值。具体来说，我们计算一个块中所有像素的平均值和标准差，并使用它们来对这个块进行归一化。在我们的实验中，使用归一化像素作为重建目标，提高了表示质量。
简单的实现我们的MAE预训练可以有效地实现，而且重要的是，不需要任何专门的稀疏操作。首先，我们为每个输入块生成一个标记（通过添加位置嵌入的线性映射）。接下来，我们随机打乱标记列表，并根据掩蔽比率删除列表的最后一部分。这个过程产生一个标记子集，相当于不重复的采样块。在编码之后，我们将向编码的块列表添加一个掩码标记列表，并恢复这个完整的列表（以使所有标记与它们的目标对齐。解码器将应用于这个完整的列表（并添加了位置嵌入）。如上所述，不需要进行稀疏操作。这个简单的实现引入了可以忽略不计的开销，因为打乱和恢复操作很快。

4.ImageNet 实验
我们在ImageNet-1K训练集上进行自监督的预训练。然后，我们进行监督训练，通过(i)端到端微调或（ii）线性探测来评估表示。我们报告了一个单一的224×224的top-1验证准确性。
**基线：ViT-Large。**在我们的消融研究中，我们使用ViT-Large（ViT-L/16）作为骨干。ViT-L非常大（比ResNet-50大一个数量级），并且很容易过拟合。以下是从零开始训练的ViT-L与基于我们的基线MAE进行的微调之后的对比：

我们注意到，从头开始训练有监督的ViT-L并不重要，只需要一个具有强正则化的好配方。即便如此，我们的MAE预训练贡献了很大的进步。这里的微调在前50个epoch（从头开始200个），这意味着微调精度很大程度上依赖于预训练。
4.1.主要属性
我们使用表1中的默认设置来进行MAE消融实验。我们观察到了几个有趣的特性。
**掩蔽率。**图5显示了掩蔽率的影响。最佳比率要惊人地高。75%的比值有利于线性探测和微调。这种行为与BERT 相反，后者的典型掩蔽比为15%。我们的掩蔽比也远高于计算机视觉中的相关作品[6,16,2]（20%-50%）。
该模型以产生不同的、但看似可信的输出推断出缺失的块。它可以理解物体和场景的格式塔，而不是简单地通过扩展线条或纹理来完成。我们假设这种类似推理的行为与学习有用的表征有关。
图5还显示了线性探测和微调结果遵循不同的趋势。对于线性探测，准确率随着掩蔽率的增加而稳步增加，直到最佳点：精度差距高达大约20%（54.6%vs73.5%）对于微调，准确度对掩蔽率的敏感性较低，而且在很大范围的掩蔽率下（40-80%）工作得很好。图5中的所有微调结果都优于从头开始的训练（82.5%）。
**解码器设计。**我们的MAE解码器可以灵活设计，如表1a和表1b所示。表1a改变了解码器的深度（Transformer块的数量）。一个足够深的解码器对于线性探测是很重要的。这可以用像素重建任务和识别任务之间的差距来解释：自动编码器的最后几层更擅长用于重建，但与识别的相关性较小。一个合理的深度解码器可以解释重构的专门化，使得潜在的表示在一个更抽象的层次上。这种设计在线性探测方面可以提高高达8%（表1a，“lin”）。然而，如果使用微调，编码器的最后一层可以调整以适应识别任务。解码器深度对微调的提升影响较小（表1a，“ft”）。
有趣的是，我们使用单块解码器的MAE可以通过微调达到强大的性能（84.8%）。请注意，单个Transformer块是将信息从可见标记传播到掩码标记的最小要求。这样小的一个解码器可以进一步加快训练速度。
在表1b中，我们研究了解码器的宽度（通道数）。我们默认使用512-d，它在微调和线性探测下表现良好。一个较窄的解码器也可以很好地进行微调。
总的来说，我们的默认MAE解码器是轻量级的。它有8个块，宽度为512-d（表1中的灰色）。每个标记只需要9%的FLOPs相比ViT-L（24块，1024-d）。因此，虽然解码器处理所有的标记，但它仍然是整个计算的一小部分。

掩码标记我们的MAE的一个重要设计是在编码器中跳过掩码标记[M]，然后将其应用到轻量级解码器中。表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标记化器需要一个额外的预训练阶段，这可能依赖于额外的数据。dVAE编码器是一个大型的卷积网络（vit-L的40%），并增加了大量的开销。使用像素的话并不会遇到这些问题。
数据增强
表1e研究了数据增强对我们的MAE预训练的影响。我们的MAE使用只进行裁剪的增强，效果很好，无论是固定大小还是随机大小（都有随机的水平翻转）。添加颜色抖动会降低结果，所以我们不在其他实验中使用它。
令人惊讶的是，即使没有使用数据增强（只有中心裁剪，没有翻转），我们的MAE也表现得很好。这一特性与严重依赖于数据增强的对比学习和相关方法[62,23,7,21]有显著的不同。通过[21]观察到，仅使用裁剪增强可在BYOL上降低13%的精度，在SimCLR上降低28%的精度。此外，没有证据表明对比学习可以在不需要增强的情况下工作：一个图像的两个视图是相同的，并且可以很容易地满足一个平凡解。
在MAE中，数据增强的作用主要通过随机掩蔽来实现。每次迭代的掩模都是不同的，所以无论数据增强量如何，它们都会生成新的训练样本。掩蔽使代理任务变得困难，因此只需要较少的增强来规范训练。
掩码抽样策略。
在表1f中，我们比较了不同的掩模采样策略，如图6所示。[2]中提出的块掩蔽策略倾向于删除大的块（图6中间）。我们的块级掩蔽的MAE在50%的比例下工作得相当有效，但在75%的比例下降。这个任务比随机抽样更难，因为观察到更高的训练损失。重建的效果也更加模糊。
我们还研究了网格采样，它定期保留每四个补丁中的一个（图6右）。这是一项更容易完成的任务，而且训练损失也更低。重建更加尖锐。但是，表示的质量较低。
简单的随机抽样最适合我们的MAE。它允许一个更高的掩蔽比，这提供了一个更大的加速效益，同时也有良好的精度。
训练进度表
到目前为止，我们的消融实验是基于800个epoch的预训练。图7显示了训练epoch的影响。随着更长时间的训练，准确性会稳步提高。事实上，即使在1600个时代，我们也没有观察到线性探测精度的饱和。这种行为不同于对比学习方法，例如，MoCo v3 [9]在ViT-L的300个时代达到饱和。请注意，MAE编码器每个epoch只看到25%的补丁，而在对比学习中，编码器每个epoch看到200%（两次复制）甚至更多块。


4.2.与以前的结果进行比较
**与自监督方法的比较。**在表3中，我们比较了自监督ViT模型的微调结果。对于ViT-B，所有的方法的性能都很接近。对于ViT-L，方法之间的差距更大，这表明更大的模型面临的挑战是减少过拟合。
我们的MAE可以很容易地扩大规模，并在更大的模型上显示出稳步的提高。我们使用ViT-H（224大小）获得了86.9%的精度。通过对448大小的微调，我们实现了87.8%的精度，仅使用IN1K数据。在所有仅使用IN1K数据的方法中，之前最好的精度是87.1%（512大小）的[67]，基于先进的网络。在竞争激烈的IN1K基准测试（没有外部数据）中，我们以显著的优势提高了最先进的水平。我们的研究结果是基于普通的ViT，我们期望先进的网络会表现得更好。
与BEiT [2]相比，我们的MAE更准确，同时更简单、更快。我们的方法重建了像素，与BEiT预测标记相比： BEiT报告了在使用ViT-b重建像素时，[2]退化了1.8%。我们不需要dVAE预训练。此外，我们的MAE比BEiT快得多（每个epoch3.5倍），原因如表1c所示。

表3中的MAE模型进行了1600个epoch的预训练，以获得更好的精度（图7）。即便如此，在相同的硬件上训练时，我们的预训练总时间比其他方法少。例如，在128个TPU-v3内核上训练ViT-L，我们MAE的1600个epoch的训练时间为31小时，而MoCov3的300个epoch[9]的训练时间为36小时。
与有监督的预训练之间的比较在最初的ViT论文[16]中，ViT-L在IN1K中训练时退化。我们实施的监督训练（见A.2）效果更好，但准确性达到饱和。见图8。
我们的MAE预训练，只使用IN1K，可以更好地概括：对于更高容量的模型，从头开始训练的增益更大。它遵循了类似于[16]的JFT-300M监督预训练的趋势。这种比较表明，我们的MAE可以帮助扩大模型的大小。
4.3.部分微调
表1显示，线性探测和微调的结果在很大程度上是不相关的。在过去的几年中，线性探测一直是一种流行的方法；然而，它忽略强大但非线性特征——这确实是深度学习的一种优势。作为中间立场，我们研究了一个部分微调方法：微调最后几层，同时冻结其他层。该方法也被用于早期的工作中，如[65,70,42]。
图9显示了结果。值得注意的是，仅微调一个Transformer块就可以将精度从73.5%显著提高到81.0%。此外，如果我们只微调最后一个块（即它的MLP子块的“一半”），我们可以得到79.1%，比线性探测好得多。这个变体本质上是对MLP头进行微调。微调几个块（例如，4或6）可以实现接近完全微调的精度。
在图9中，我们还与MoCo v3 [9]进行了比较，这是一种与ViT-L结果相结合的对比方法。MoCo v3具有较高的线性探测精度；然而，它所有的部分微调结果都比MAE差。当调优4个块时，差距为2.6%。虽然MAE表示的线性可分离性较少，但它们具有更强的非线性特征，当非线性头部被调优时表现良好。

这些观察结果表明，线性可分性并不是评价表示质量的唯一指标。人们还观察到（例如，[8]），线性探测与迁移学习性能并没有很好的相关性，例如，对于目标检测。据我们所知，线性评价在NLP中并不经常用于预训练的基准测试。
5.迁移学习实验
我们使用表3中的预训练模型来评估下游任务中的迁移学习。

目标检测和分割

我们在COCO [37]上端到端的微调Mask R-CNN [24]。FPN适用于ViT主干的使用。我们对表4中的所有条目都应用了这种方法。我们展示了用于目标检测的方框AP和用于实例分割的掩码AP。
与有监督的预训练相比，我们的MAE在所有配置下都表现得更好（表4）。对于较小的ViT-B，我们的MAE比有监督的预训练高出2.4分（50.3 vs. 47.9，APbox）。更重要的是，对于较大的ViT-L，我们的MAE预训练比有监督的预训练多出4.0分（53.3 vs. 49.3）。
基于像素的MAE比基于标记的BEiT更好或相当，而MAE更简单、更快。MAE和BEiT都优于MoCov3，MoCov3与有监督的预训练相当。
语义分割
我们使用UperNet [63]在ADE20K [72]上进行了实验（见A.4）。表5显示，我们的预训练比监督预训练显著提高了结果，例如，ViT-L提高了3.7分。我们基于像素的MAE也优于基于标记的BEiT。这些观察结果与COCO的观察结果一致。
分类任务
表6研究了iNaturalists[56]和Places [71]任务上的迁移学习（见A.5）。在iNat上，我们的方法显示出很强的尺度行为：随着更大的模型，精度显著提高。我们的结果远远超过了之前的最佳结果。在Places上，我们的MAE优于之前的最佳结果[19,40]，而其是通过对数十亿张图像的预训练获得的。
像素与标记
表7对比了像素与标记作为MAE重建目标。虽然使用dVAE标记比使用非标准化像素更好，但它在统计上与在我们测试的所有案例中使用标准化像素相似。这再次表明，标记化对我们的MAE并不是必要的。

6.讨论和结论
可扩展性好的简单算法是深度学习的核心。在自然语言处理领域，简单的自监督学习方法可以从指数级缩放模型中获益。在计算机视觉中，尽管在自监督学习方面取得了进展，但实际的预训练范式主要还是有监督的。在本研究中，我们在ImageNet上和迁移学习中观察到，自编码器(一种简单的自监督方法，类似于NLP中的技术)提供了可扩展的好处。视觉领域的自监督学习现在可能走上了与NLP类似的轨道。
另一方面，我们注意到图像和语言是不同性质的信号，必须仔细处理这种差异。图像仅仅是记录下来的光，没有将语义分解为文字的视觉模拟。相比于试图删除对象，我们删除随机块，最有可能不形成一个语义段。同样，我们的MAE重建像素，像素不是语义实体。然而，我们观察到我们的MAE推断出复杂的、整体的重建，这表明它已经学习了许多视觉概念，即语义。我们假设这种行为是通过MAE内部丰富的隐藏表示发生的。我们希望这一观点将启发未来的工作。
**更广泛的影响。**所提出的方法基于训练数据集的学习统计来预测内容，因此将反映这些数据中的偏见，包括具有负面社会影响的偏见。模型可能生成不存在的内容。当建立在这项工作所产生的图像上时，这些问题需要进一步的研究和考虑。