【翻译】 Masked Autoencoders Are Scalable Vision Learners

论文链接：https://arxiv.org/pdf/2111.06377v3.pdf

摘要

本文表明，掩码自编码器(MAE)是用于计算机视觉的可扩展自监督学习器。我们的MAE方法很简单：我们屏蔽输入图像的随机块并重建丢失的像素。它基于两个核心设计。首先，我们开发了一个非对称编码器-解码器架构，其中一个编码器只对可见的补丁子集进行操作（没有掩码标记），还有一个轻量级的解码器，可以从潜在表示和掩码标记重建原始图像。其次，我们发现屏蔽大部分输入图像（例如 75%）会产生重要且有意义的自我监督任务。结合这两种设计使我们能够高效地训练大型模型：我们加速训练（3倍或更多）并提高准确性。我们的可扩展方法允许学习泛化良好的高容量模型：例如，在仅使用ImageNet-1K数据的方法中，vanillaViT-Huge模型实现了最佳准确率(87.8%)。下游任务中的传输性能优于有监督的预训练，并显示出有希望的扩展行为。

1.介绍

深度学习见证了能力和容量不断增长的架构的爆炸式增长[33, 25, 57]。在硬件快速增长的帮助下，今天的模型很容易过拟合一百万张图像[13]，并开始需要数亿张——通常是公开无法访问的——标记图像[16]。
　　这种对数据的需求已经通过自监督预训练在自然语言处理(NLP)中成功解决。基于GPT[47,48,4]中的自回归语言建模和BERT[14]中的掩码自动编码的解决方案在概念上很简单：它们移除了一部分数据并学习预测移除的内容。这些方法现在可以训练包含超过一千亿个参数的可泛化NLP模型[4]。
　　掩码自动编码器的想法，一种更通用的去噪自动编码器[58]的形式，是自然的，也适用于计算机视觉。事实上，与视觉密切相关的研究[59, 46]早于BERT。然而，尽管随着BERT的成功，人们对这个想法产生了浓厚的兴趣，但视觉中自动编码方法的进展落后于NLP。我们问：是什么使视觉和语言之间的掩码自动编码不同？我们尝试从以下几个角度来回答这个问题：
　　(i)直到最近，架构还是不同的。在视觉方面，卷积网络[34]在过去十年中占主导地位[33]。卷积通常在规则网格上运行，将“指标”（例如掩码标记[14]或位置嵌入[57]）集成到卷积网络中并不容易。然而，随着Vision Transformers(ViT)[16]的引入，这种架构差距已经得到解决，不应再成为障碍。
　　(ii)语言和视觉的信息密度不同。语言是人类生成的具有高度语义和信息密集度的信号。当训练一个模型来预测每个句子的几个遗漏单词时，这个任务似乎会导致复杂的语言理解。相反，图像是具有大量空间冗余的自然信号——例如，可以从相邻的块中恢复丢失的块，而对部分、对象和场景的高级理解很少。为了克服这种差异并鼓励学习有用的特征，我们展示了一个简单的策略在计算机视觉中效果很好：屏蔽了很大一部分的随机patches。这种策略在很大程度上减少了冗余并创建了一项具有挑战性的自我监督任务，需要超越低级图像统计的整体理解。要对我们的重建任务有一个定性的认识，请参见图 2-4。
　　(iii)自动编码器的解码器将潜在表示映射回输入，在重建文本和图像之间扮演不同的角色。在视觉中，解码器重建像素，因此其输出的语义级别低于普通识别任务。这与语言相反，在语言中，解码器预测包含丰富语义信息的缺失词。虽然在BERT中，解码器可能很简单（MLP）[14]，但我们发现对于图像，解码器设计在确定学习到的潜在表示的语义级别方面起着关键作用。



　　在此分析的推动下，我们提出了一种简单、有效且可扩展的掩码自编码器(MAE)形式，用于视觉表示学习。我们的MAE从输入图像中屏蔽了随机patches，并在像素空间中重建丢失的patches。它具有非对称编码器解码器设计。我们的编码器只对可见的patches子集（没有掩码标记）进行操作，我们的解码器是轻量级的，可以从潜在表示和掩码标记中重建输入（图 1）。在我们的非对称编码器-解码器中将掩码标记转移到小型解码器会导致计算量大幅减少。在这种设计下，非常高的掩蔽率（例如 75%）可以实现双赢：它优化了准确性，同时允许编码器仅处理一小部分（例如 25%）的patches。这可以将整体预训练时间减少3倍或更多，同样减少内存消耗，使我们能够轻松地将MAE扩展到大型模型。
　　我们的MAE学习了非常高容量的模型，可以很好地泛化。通过MAE预训练，我们可以在ImageNet-1K上训练ViT-Large/-Huge[16]等数据饥渴模型，并提高泛化性能。使用普通的ViT-Huge模型，我们在ImageNet-1K上微调时达到了87.8%的准确率。这优于之前仅使用ImageNet-1K数据的所有结果。我们还评估了对象检测、实例分割和语义分割的迁移学习。在这些任务中，我们的预训练比其监督预训练获得了更好的结果，更重要的是，我们通过扩大模型观察到了显着的收益。这些观察结果与NLP中自监督预训练[14、47、48、4]中的观察结果一致，我们希望它们将使我们的领域能够探索类似的轨迹。

2.相关工作

掩码语言建模及其自回归对应物，例如BERT[14]和 GPT[47, 48, 4]，是非常成功的NLP预训练方法。这些方法保留一部分输入序列并训练模型来预测缺失的内容。这些方法已被证明具有出色的扩展性[4]，并且大量证据表明，这些预训练的表示可以很好地推广到各种下游任务。
自编码是学习表征的经典方法。它有一个将输入映射到潜在表示的编码器和一个重构输入的解码器。例如，PCA和k-means是自动编码器[29]。去噪自编码器(DAE)[58]是一类自编码器，它破坏输入信号并学习重建原始的、未破坏的信号。一系列方法可以被认为是不同损坏下的广义DAE，例如，屏蔽像素[59, 46, 6]或去除颜色通道[70]。我们的MAE是一种去噪自动编码形式，但在许多方面与经典DAE不同。
掩码图像编码方法从被掩码损坏的图像中学习表示。[59]的开创性工作将掩蔽作为DAE中的一种噪声类型。上下文编码器[46]使用卷积网络修复大的缺失区域。受NLP成功的启发，最近的相关方法[6, 16, 2]基于Transformers[57]。iGPT[6]对像素序列进行操作并预测未知像素。ViT论文[16]研究了用于自监督学习的掩码patches预测。最近，BEiT[2]提出预测离散标记[44, 50]。
自监督学习方法对计算机视觉产生了浓厚的兴趣，通常侧重于不同的预训练借口任务[15, 61, 42, 70, 45, 17]。最近，对比学习[3, 22]很流行，例如[62, 43, 23, 7]，它对两个或多个视图之间的图像相似性和不相似性（或仅相似性[21, 8]）进行建模。对比和相关方法强烈依赖于数据增强[7, 21, 8]。自动编码追求概念上不同的方向，它表现出我们将要展示的不同行为。

3.方法

我们的掩码自编码器(MAE)是一种简单的自编码方法，可以在给定部分观察的情况下重建原始信号。与所有自动编码器一样，我们的方法有一个编码器，将观察到的信号映射到潜在表示，以及一个解码器，从潜在表示重建原始信号。与经典的自动编码器不同，我们采用非对称设计，允许编码器仅对部分观察信号（无掩码标记）进行操作，并采用轻量级解码器从潜在表示和掩码标记重建完整信号。图1说明了这个想法，接下来介绍。

掩蔽。遵循ViT[16]，我们将图像划分为规则的非重叠patches。然后我们对patches的子集进行采样并屏蔽（即删除）剩余的patches。我们的采样策略很简单：我们按照均匀分布对随机补丁进行采样，无需替换。我们简单地将其称为“随机抽样”。
　　具有高掩蔽率（即移除patches的比率）的随机采样在很大程度上消除了冗余，从而创建了一项无法通过从可见相邻patches外推来轻松解决的任务（见图 2-4）。均匀分布可防止潜在的中心偏差（即，图像中心附近有更多的掩蔽patches）。最后，高度稀疏的输入为设计高效编码器创造了机会，接下来介绍。
MAE编码器。我们的编码器是ViT[16]，但仅应用于可见的、未屏蔽的patches。就像在标准ViT中一样，我们的编码器通过添加位置嵌入的线性投影嵌入补丁，然后通过一系列Transformer块处理结果集。然而，我们的编码器只对整个集合的一小部分（例如25%）进行操作。被屏蔽的patches被移除；不使用掩码token。这使我们能够仅使用一小部分计算和内存来训练非常大的编码器。全套由轻量级解码器处理，如下所述。
MAE解码器。MAE解码器的输入是完整的token集，包括(i)编码的可见patches和(ii)掩码token。参见图1。每个掩码标记[14]是一个共享的、学习的向量，表示存在要预测的缺失patches。我们为这个完整集合中的所有标记添加了位置嵌入；如果没有这个，掩码token将没有关于它们在图像中的位置的信息。解码器有另一系列的Transformer模块。
　　MAE解码器仅在预训练期间用于执行图像重建任务（仅使用编码器生成用于识别的图像表示）。因此，可以以独立于编码器设计的方式灵活地设计解码器架构。我们试验了非常小的解码器，比编码器更窄更浅。例如，与编码器相比，我们的默认解码器每个token的计算量小于 10%。通过这种非对称设计，全套token仅由轻量级解码器处理，这显着减少了预训练时间。
重建目标。我们的MAE通过预测每个掩码patches的像素值来重建输入。解码器输出中的每个元素都是一个表示patch的像素值向量。解码器的最后一层是线性投影，其输出通道的数量等于patch中像素值的数量。解码器的输出被重新整形以形成重建的图像。我们的损失函数计算像素空间中重建图像和原始图像之间的均方误差(MSE)。我们仅在掩码patches上计算损失，类似于BERT[14]。
　　我们还研究了一个变体，其重建目标是每个掩码patch的归一化像素值。具体来说，我们计算一个patch中所有像素的均值和标准差，并使用它们来规范化这个patch。在我们的实验中，使用归一化像素作为重建目标可以提高表示质量。
简单的实现。我们的MAE预训练可以高效实施，重要的是，不需要任何专门的稀疏操作。首先，我们为每个输入patch生成一个token（通过线性投影和添加的位置嵌入）。接下来，我们根据掩蔽率随机打乱tokens列表并删除列表的最后一部分。此过程为编码器生成一小部分tokens，相当于采样patches而无需替换。编码后，我们将掩码tokens列表附加到编码patches列表中，并反打乱这个完整列表（反转随机洗牌操作）以将所有tokens与其目标对齐。解码器应用于此完整列表（添加了位置嵌入）。如前所述，不需要稀疏操作。这个简单的实现引入的开销可以忽略不计，因为shuffle和unshuffling操作很快。

4.Imagenet实验

我们在ImageNet-1K(IN1K)[13]训练集上进行自我监督的预训练。然后我们进行有监督的训练，以通过(i)端到端微调或(ii)线性探测来评估表示。我们报告了单个224×224裁剪的top-1验证精度。详情见附录 A.1。
Baseline：ViT-Large。我们在消融研究中使用ViT-Large(ViT-L/16)[16]作为主干。ViT-L非常大（比ResNet-50[25]大一个数量级）并且容易过拟合。以下是从头开始训练的ViT-L与从我们的基线MAE微调的比较：

我们注意到从头开始训练有监督的ViT-L是很重要的，并且需要一个具有强正则化的良好配方（82.5%，参见附录 A.2）。尽管如此，我们的MAE预训练还是有很大的改进。此处微调仅针对50个epoch（而从头开始200个），这意味着微调精度在很大程度上取决于预训练。

4.1. 主要属性

我们使用表1中的默认设置烧蚀我们的MAE（见标题）。观察到几个有趣的特性。

掩蔽比。图5显示了掩蔽比的影响。最佳比率高得惊人。75%的比率对线性探测和微调都有利。这种行为与BERT[14]形成对比，后者的典型掩蔽率为15%。我们的掩蔽率也远高于计算机视觉相关工作[6, 16, 2]中的掩蔽率（20%到50%）。
　　该模型推断缺失的patches以产生不同但合理的输出（图4）。它使物体和场景的格式塔变得有意义，这不能简单地通过延伸线条或纹理来完成。我们假设这种类似推理的行为与学习有用的表征有关。
　　图5还显示线性探测和微调结果遵循不同的趋势。对于线性探测，准确度随着掩蔽比稳步增加，直到最佳点：准确度差距高达20%（54.6%对73.5%）。对于微调，结果对比率不那么敏感，并且很宽范围的掩蔽比率(40-80%)效果很好。图5中的所有微调结果都优于从头开始训练(82.5%)。
解码器设计。我们的MAE解码器可以灵活设计，如表1a和1b所示。
　　表1a改变了解码器深度（Transformer块的数量）。足够深的解码器对于线性探测很重要。这可以通过像素重建任务和识别任务之间的差距来解释：自动编码器中的最后几层更专门用于重建，但与识别的相关性较低。一个合理深度的解码器可以解释重建专业化，将潜在表示留在更抽象的层次上。这种设计可使线性探测提高8%（表1a，“lin”）。但是，如果使用微调，则可以调整编码器的最后几层以适应识别任务。解码器深度对改进微调的影响较小（表1a，‘ft’）。
　　有趣的是，我们的带有单块解码器的MAE可以通过微调(84.8%)表现得很好。请注意，单个Transformer块是将信息从可见token传播到掩码token的最低要求。这么小的解码器可以进一步加快训练速度。
　　在表1b中，我们研究了解码器宽度（通道数）。我们默认使用512-d，它在微调和线性探测下表现良好。更窄的解码器也适用于微调。
　　总的来说，我们默认的MAE解码器是轻量级的。它有8个块，宽度为512-d（表1中的灰色）。与ViT-L（24个区块，1024-d）相比，每个token只有9%的FLOP。因此，虽然解码器处理所有tokens，但它仍然是整体计算的一小部分。
Mask token。我们MAE的一个重要设计是跳过编码器中的掩码token[M]，稍后将其应用到轻量级解码器中。表1c研究了这种设计。
　　如果编码器使用掩码token，它的性能会更差：在线性探测中其精度下降14%。在这种情况下，预训练和部署之间存在差距：该编码器在预训练的输入中具有很大一部分掩码token，这在未损坏的图像中不存在。这种差距可能会降低部署的准确性。通过从编码器中移除掩码token，我们限制编码器始终看到真实的patches，从而提高准确性。
　　此外，通过跳过编码器中的掩码token，我们大大减少了训练计算。在表1c中，我们将整体训练FLOP减少了3.3倍。这导致我们的实现中2.8倍挂钟加速（见表2）。对于较小的解码器（1块）、较大的编码器(ViT-H)或两者，挂钟加速甚至更大（3.5-4.1倍）。请注意，对于75%的掩蔽率，加速可以大于4倍，部分原因是自注意力复杂度是二次的。此外，内存大大减少，可以训练更大的模型或通过大批量训练加快速度。时间和内存效率使我们的MAE有利于训练非常大的模型。

重建目标。我们在表1d中比较了不同的重建目标。到目前为止，我们的结果基于没有（每个patch）归一化的像素。使用归一化像素可提高准确性。这种逐块归一化在局部增强了对比度。在另一个变体中，我们在patch空间中执行PCA并使用最大的PCA系数（此处为96）作为目标。这样做会降低准确性。两个实验都表明高频分量在我们的方法中很有用。
　　我们还比较了预测token的MAE变体，即BEiT[2]中使用的目标。特别是对于这个变体，我们使用DALLE预训练的dVAE[50]作为标记器，遵循[2]。这里MAE解码器使用交叉熵损失预测token索引。这种token化与非标准化像素相比将微调精度提高了0.4%，但与标准化像素相比没有优势。它还降低了线性探测精度。在第5节中，我们进一步表明在转移学习中不需要token化。
　　我们基于像素的MAE比token化简单得多。dVAE分词器需要一个额外的预训练阶段，这可能取决于额外的数据（250M图像[50]）。dVAE编码器是一个大型卷积网络（ViT-L的40% FLOPs）并增加了非平凡的开销。使用像素不会遇到这些问题。
数据增强。表1e研究了数据增强对我们的MAE预训练的影响。
　　我们的MAE使用仅裁剪增强效果很好，无论是固定大小还是随机大小（两者都具有随机水平翻转）。添加颜色抖动会降低结果，因此我们不会在其他实验中使用它。
令人惊讶的是，即使没有使用数据增强（只有中心裁剪，没有翻转），我们的MAE也表现得很好。这一特性与对比学习和相关方法[62, 23, 7, 21]截然不同，后者严重依赖于数据增强。据观察[21]，对于BYOL[21]和SimCLR[7]，使用仅裁剪增强分别将准确度降低了13%和28%。此外，没有证据表明对比学习可以在没有增强的情况下工作：图像的两个视图是相同的，可以很容易地满足一个简单的解决方案。
　　在MAE中，数据增强的作用主要是通过随机掩码来实现的（下一个消融）。每次迭代的掩码都不同，因此无论数据增强如何，它们都会生成新的训练样本。掩蔽使借口任务变得困难，并且需要较少的增强来规范训练。

掩码采样策略。在表1f中，我们比较了不同的掩码采样策略，如图6所示。
　　[2]中提出的逐块屏蔽策略倾向于移除大块（图6中）。我们的带块屏蔽的MAE在50%的比率下工作得相当好，但在75%的比率下性能下降。这项任务比随机抽样更难，因为观察到更高的训练损失。重建也更模糊。
　　我们还研究了网格采样，它定期保留每四个补丁中的一个（图6右）。这是一项更简单的任务，并且具有更低的训练损失。重建更加清晰。但是，表示质量较低。
　　简单随机抽样最适合我们的MAE。它允许更高的掩蔽率，这提供了更大的加速优势，同时还享有良好的准确性。
训练计划。到目前为止，我们的消融是基于800 epoch的预训练。图7显示了训练计划长度的影响。随着训练时间的延长，准确率稳步提高。事实上，即使在1600个epoch，我们也没有观察到线性探测精度的饱和。这种行为与对比学习方法不同，例如，MoCo v3[9]在ViT-L的300个epoch时饱和。请注意，MAE编码器每个epoch只能看到25%的patches，而在对比学习中，编码器每个epoch看到200%（two-crop）甚至更多（multi-crop）patches。

4.2. 与之前结果的比较

与自监督方法的比较。在表3中，我们比较了自监督ViT模型的微调结果。对于ViT-B，所有方法的表现都很接近。对于ViT-L，方法之间的差距更大，这表明更大模型的挑战是减少过度拟合。
　　我们的MAE可以轻松扩展，并且从更大的模型中已经显示出稳定的改进。我们使用ViT-H（224 尺寸）获得了86.9%的准确率。通过使用448尺寸进行微调，我们仅使用IN1K数据实现了87.8%的准确率。在所有仅使用IN1K数据的方法中，以前的最佳准确率是87.1%（512大小）[67]，基于高级网络。在竞争激烈的IN1K基准（无外部数据）中，我们以显着的优势改进了最先进的技术。我们的结果基于vanilla ViT，我们预计高级网络的性能会更好。
　　与BEiT[2]相比，我们的MAE更准确，同时更简单、更快。与预测token的BEiT相比，我们的方法重建像素：BEiT报告在使用ViT-B.2重建像素时降低了1.8%[2]。我们不需要dVAE预训练。此外，由于表1c中研究的原因，我们的MAE比BEiT快得多（每个epoch 3.5倍）。
　　表3中的MAE模型经过1600个epoch的预训练以获得更好的准确性（图7）。即便如此，在相同硬件上训练时，我们的总预训练时间仍少于其他方法。例如，在128个TPU-v3内核上训练ViT-L，我们的MAE的训练时间为31小时，1600个epochs，而MoCo v3的训练时间为36小时，300个epochs[9]。
与监督预训练的比较。在最初的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]。ViT 主干适用于FPN[36]（见A.3）。我们将此方法应用于表4中的所有条目。我们报告用于对象检测的框AP 和用于实例分割的掩码AP。

　　与监督预训练相比，我们的MAE在所有配置下表现更好（表4）。使用较小的ViT-B，我们的MAE比监督预训练高2.4个点（50.3对47.9，APbox）。更重要的是，使用更大的ViT-L，我们的MAE预训练比监督预训练高4.0个点（53.3对49.3）。
　　基于像素的MAE优于或与基于token的BEiT 相当，而MAE更简单、更快。MAE和BEiT 都优于MoCo v3，并且MoCo v3与监督预训练相当。
语义分割。我们使用UperNet[63]（见A.4）在ADE20K[72]上进行实验。表5显示我们的预训练显着改善了监督预训练的结果，例如ViT-L提高了3.7个点。我们基于像素的MAE也优于基于token的BEiT。这些观察结果与COCO中的观察结果一致。

分类任务。表6研究了iNaturalists[56]和Places[71]任务的迁移学习（见 A.5）。在iNat上，我们的方法显示出强大的缩放行为：更大的模型可以显着提高准确性。我们的结果大大超过了之前的最佳结果。在Places上，我们的MAE优于之前通过对数十亿张图像进行预训练获得的最佳结果[19, 40]。

像素和token。表7比较了作为MAE重建目标的像素与令牌。虽然使用dVAE token比使用非标准化像素更好，但在我们测试的所有情况下，它与使用标准化像素在统计上相似。这再次表明我们的MAE不需要token化。

6.讨论与结论

可扩展的简单算法是深度学习的核心。在NLP中，简单的自监督学习方法（例如[47, 14, 48, 4]）可以从指数缩放模型中获益。在计算机视觉中，尽管自监督学习取得了进展，但实际的预训练范式主要受到监督（例如[33, 51, 25, 16]）。在这项研究中，我们在ImageNet和迁移学习中观察到自动编码器——一种类似于NLP技术的简单自我监督方法——提供了可扩展的好处。视觉中的自监督学习现在可能走上与NLP类似的轨迹。
　　另一方面，我们注意到图像和语言是不同性质的信号，必须谨慎处理这种差异。图像只是记录的光，没有将语义分解为文字的视觉类比。我们没有尝试删除对象，而是删除最有可能不形成语义段的随机补丁。同样，我们的MAE重建像素，它们不是语义实体。尽管如此，我们观察到（例如，图 4）我们的MAE推断出复杂的整体重建，表明它已经学习了许多视觉概念，即语义。我们假设这种行为是通过MAE内部丰富的隐藏表示发生的。我们希望这个观点能激发未来的工作。