Masked Autoencoders Are Scalable Vision Learners

Abstract

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

1.Introduction

深度学习见证了能力和容量不断增长的架构的爆炸式增长[29,25,52]。在硬件快速发展的帮助下，今天的模型可以很容易地过拟合一百万张图像[13]，并开始需要数亿张通常是公开无法访问的标记图像[16]。

这种对数据的需求已经在自然语言处理(NLP)中通过自我监督预训练成功地解决了。基于GPT[42,43,4]中的自回归语言建模和BERT[14]中的掩码自动编码的解决方案在概念上很简单:它们删除一部分数据并学习预测被删除的内容。这些方法现在可以训练包含超过一千亿个参数的可泛化NLP模型[4]。

掩码自编码器是一种更通用的去噪自编码器[53]，它的想法很自然，也适用于计算机视觉。事实上，在BERT之前就有与视觉密切相关的研究[54,41]。然而，尽管随着BERT的成功，人们对这一想法产生了浓厚的兴趣，但视觉领域的自动编码方法的进展落后于NLP。我们的问题是:是什么让掩码自动编码在视觉和语言之间有所不同?我们试图从以下角度回答这个问题:

(i)直到最近，架构都是不同的。在视觉领域，卷积网络[30]在过去十年中占据主导地位[29]。卷积通常在规则网格上运行，将“指标”(如掩码token[14]或位置嵌入[52])集成到卷积网络中并不简单。然而，这种架构上的差距已经通过视觉转换器(ViT)[16]的引入得到了解决，并且不再是一个障碍。

(ii)语言和视觉的信息密度不同。语言是人类产生的具有高度语义和信息密集的信号。当训练一个模型来预测每句话中只缺几个单词时，这项任务似乎可以诱导复杂的语言理解。相反，图像是具有大量空间冗余的自然信号。一个缺失的patch可以从邻近的patch中恢复，而对局部、对象和场景几乎没有高层次的理解。
为了克服这种差异并鼓励学习有用的特征，我们展示了一个简单的策略在计算机视觉中很有效:掩盖非常高比例的随机patch。这种策略在很大程度上减少了冗余，并创建了一个具有挑战性的自我监督任务，该任务需要超越低级图像统计的整体理解。要对我们的重建任务有一个定性的认识，请参见图2 - 4。



(iii)自动编码器的解码器，将潜在表示映射回输入，在重建文本和图像之间起着不同的作用。在视觉上，解码器重建像素，因此其输出的语义水平低于普通识别任务。这与语言相反，在语言中，解码器预测包含丰富语义信息的缺失单词。虽然在BERT中，解码器可能是微不足道的(一个MLP)[14]，但我们发现，对于图像，解码器的设计在决定学习到的潜在表征的语义水平方面起着关键作用。

在此分析的推动下，我们提出了一种简单、有效、可扩展的掩码自编码器(MAE)用于视觉表示学习。我们的MAE从输入图像中屏蔽随机补丁，并在像素空间中重建缺失的补丁。它具有非对称编解码器设计。我们的编码器仅对补丁的可见子集(没有掩码令牌)进行操作，我们的解码器是轻量级的，可以从潜在表示和掩码令牌一起重建输入(图1)。将掩码令牌转移到我们的非对称编码器-解码器中的小型解码器中，可以大大减少计算量。在这种设计下，非常高的掩码比率(例如，75%)可以实现双赢:它优化了精度，同时允许编码器只处理一小部分(例如，25%)的patch。这可以将总体预训练时间减少3倍或更多，同样减少内存消耗，使我们能够轻松地将MAE扩展到大型模型。

我们的MAE学习非常高容量的模型，泛化得很好。通过MAE预训练，我们可以在ImageNet-1K上训练像vit- Large /-Huge[16]这样的数据密集型模型，提高泛化性能。使用普通的ViT-Huge模型，在ImageNet-1K上进行微调时，我们达到了87.8%的准确率。这优于以前仅使用ImageNet-1K数据的所有结果。我们还评估了对象检测、实例分割和语义分割的迁移学习。在这些任务中，我们的预训练比有监督的预训练取得了更好的结果，更重要的是，我们通过缩放模型观察到显著的收益。这些观察结果与NLP中自我监督预训练中的观察结果一致[14,42,43,4]，我们希望它们将使我们的领域能够探索类似的轨迹。

2.Related Work

掩模语言建模及其自回归模型，如BERT[14]和GPT[42,43,4]，是NLP中非常成功的预训练方法。这些方法保留输入序列的一部分，并训练模型来预测缺失的内容。这些方法已被证明具有出色的扩展性[4]，大量证据表明，这些预训练的表征可以很好地推广到各种下游任务。

自动编码是学习表征的经典方法。**它有一个将输入映射到潜在表示的编码器和一个重建输入的解码器。**例如，PCA和k-means是自编码器[26]。去噪自编码器(DAE)[53]是一类破坏输入信号并学习重建原始未损坏信号的自编码器。一系列方法可以被认为是不同破坏下的广义DAE，例如，屏蔽像素[54,41,6]或去除颜色通道[64]。我们的MAE是一种去噪的自动编码形式，但与经典的DAE有许多不同之处。

掩码图像编码方法从被掩码的图像中学习表示。[54]的开创性工作将掩蔽作为DAE中的一种噪声类型。上下文编码器[41]使用卷积网络绘制大型缺失区域。受NLP成功的影响，最近的相关方法[6,16,2]都是基于Transformers[52]。iGPT[6]对像素序列进行操作，并预测未知像素。ViT论文[16]研究了自监督学习的掩模patch预测。最近，BEiT[2]提出预测离散token[39,45]。

自监督学习方法在计算机视觉领域引起了极大的兴趣，通常关注于预训练的不同借口任务[15,55,37,64,40,17]。最近，对比学习[3,22]很受欢迎，例如[56,38,23,7]，它对两个或多个视图之间的图像相似性和不相似性(或仅相似性[21,8])进行建模。对比和相关方法强烈依赖于数据增强[7,21,8]。自动编码追求一个概念上不同的方向，它表现出不同的行为，正如我们将要介绍的。

3.Approach

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

掩蔽。根据ViT[16]，我们将图像划分为规则的不重叠的小块。然后我们对patch的子集进行采样，并掩码(即删除)剩余的patch。我们的采样策略很简单:我们对随机斑块进行采样，不进行替换，遵循均匀分布。我们简单地称之为“随机抽样”。

具有高掩码率(即去除斑块的比例)的随机采样在很大程度上消除了冗余，从而创建了一个不能通过从可见的邻近斑块外推轻松解决的任务(见图2 - 4)。均匀分布防止了潜在的中心偏差(即在图像中心附近有更多的掩蔽斑块)。最后，高度稀疏的输入为设计高效的编码器创造了机会，下面将介绍。

MAE编码器。我们的编码器是一个ViT[16]，但只应用于可见的，未遮挡的补丁。就像在标准ViT中一样，我们的编码器通过添加位置嵌入的线性投影嵌入补丁，然后通过一系列Transformer块处理结果集。然而，我们的编码器只在完整集合的一小部分(例如，25%)上运行。遮罩补丁被移除;不使用掩码令牌。这允许我们只用一小部分的计算和内存训练非常大的编码器。完整的集合由一个轻量级解码器处理，下面将介绍。

MAE解码器。MAE解码器的输入是一整套token，由(i)编码的可见补丁和(ii)掩码token组成。参见图1。每个掩码token[14]是一个共享的学习向量，表示存在待预测的缺失patch。我们为这个完整集合中的所有token添加位置嵌入;如果没有这个，掩码token将没有关于它们在图像中的位置的信息。解码器有另一系列的Transformer块。

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

重建目标。我们的MAE通过预测每个被屏蔽patch的像素值来重建输入。解码器输出中的每个元素都是代表一个patch的像素值向量。解码器的最后一层是一个线性投影，其输出通道的数量等于一个patch中的像素值的数量。对解码器的输出进行重构以形成重构图像。我们的损失函数在像素空间中计算重建图像和原始图像之间的均方误差(MSE)。我们只计算掩码patch上的损失，类似于BERT[14]。我们还研究了一种变体，其重建目标是每个被屏蔽patch的归一化像素值。具体来说，我们计算一个patch中所有像素的均值和标准差，并使用它们对该patch进行归一化。在我们的实验中，使用归一化像素作为重建目标提高了表示质量。

简单的实现。我们的MAE预训练可以有效地实现，重要的是，它不需要任何专门的稀疏操作。首先，我们为每个输入补丁生成一个标记(通过添加位置嵌入的线性投影)。接下来，我们随机打乱令牌列表，并根据屏蔽比率删除列表的最后一部分。这个过程为编码器生成一小部分标记，相当于采样补丁而不进行替换。编码后，我们将一个掩码令牌列表添加到编码补丁列表中，并取消这个完整列表(反转随机洗牌操作)，以使所有令牌与其目标对齐。解码器应用于这个完整的列表(添加了位置嵌入)。如前所述，不需要稀疏操作。这个简单的实现引入的开销可以忽略不计，因为变换和解变换操作非常快。