Masked Autoencoders Are Scalable Vision Learners——论文学习

论文地址：https://arxiv.org/pdf/2111.06377.pdf

官方源码：https://github.com/facebookresearch/mae

一、主要内容

        本文证明了掩码自编码器(MAE)是一种可扩展的计算机视觉自监督学习算法。本文的MAE方法很简单:屏蔽输入图像的随机补丁并重建缺失的像素。它基于两个核心设计。首先，本文开发了一个非对称编码器-解码器架构，其中一个编码器仅对patch的可见子集(没有掩码令牌)进行操作，以及一个轻量级解码器，该解码器从潜在表示和掩码令牌重建原始图像。其次，本文发现掩盖输入图像的高比例，例如75%，产生了一个重要的和有意义的自我监督任务。

        这两种设计的结合使本文能够高效地训练大型模型:本文加速了训练(3倍或更多)并提高了准确性。本文的可扩展方法允许学习泛化良好的大容量模型:例如，在仅使用ImageNet-1K数据的方法中，vanilla ViT-Huge模型达到了最好的准确率(87.8%)。下游任务的迁移性能优于监督预训练，并显示出有希望的缩放行为。

二、提出背景

        基于GPT中的自回归语言建模和BERT中的掩码自动编码的解决方案在概念上很简单:它们删除部分数据并学习预测被删除的内容。然而，尽管随着BERT的成功，人们对这一想法产生了浓厚的兴趣，但视觉领域的自动编码方法的进展落后于NLP。本文的问题是:是什么让掩码自动编码在视觉和语言之间有所不同?本文试图从以下几个角度来回答这个问题:

（1）在视觉领域，卷积网络在过去十年中占据主导地位。卷积通常在规则网格上运行，将“指标”(如掩码令牌或位置嵌入)集成到卷积网络中并不简单。然而，随着视觉变压器(ViT)的引入，这个架构上的差距已经得到了解决，不再是一个障碍。

（2）不同语言之间的信息密度不同。语言是人类产生的信号，是高度语义化和信息密集的。训练时一个预测每句话只缺几个词的模型，这项任务似乎诱发了复杂的语言不足。相反，图像是自然的信号重空间冗余，例如：一个缺失的补丁可以被重新修复，从邻近的补丁覆盖，对部件，对象和场景的高级理解很少。为了克服这种差异并鼓励学习有用的特征，本文展示了一个简单的策略在计算机视觉中很有效:掩盖非常高比例的随机补丁。这种策略在很大程度上减少了冗余，并创建了一个具有挑战性的自我监督任务，该任务需要超越低级图像统计的整体理解。

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

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

图 1 本文的MAE架构

三、相关工作

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

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

        遮罩图像编码方法从被遮罩损坏的图像中学习表示。受NLP成功的影响，最近的相关方法都是基于Transformers。iGPT对像素序列进行操作并预测未知像素。ViT论文研究了自监督学习的掩膜补丁预测。最近，BEiT提出了预测离散令牌的方法。

四、模型详细介绍

图 2 ImageNet验证图像的示例结果

       Masking。根据ViT，本文将图像划分为规则的不重叠的小块。然后对patch的子集进行采样，并掩码(即删除)剩余的patch。本文的采样策略很简单：对随机斑块进行采样，不进行替换，遵循均匀分布。简单地称之为“随机抽样”。如图 3 所示，“随机抽样”的效果最好。

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

图 3 掩码采样策略，影响重构质量。左:随机抽样(默认值)。Middle:块抽样[2]，删除大的随机块。右图:网格采样，每四个补丁保留一个。图像来自验证集。

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

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

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

        重建目标。本文的MAE通过预测每个被屏蔽补丁的像素值来重建输入。解码器输出中的每个元素都是代表一个patch的像素值向量。解码器的最后一层是一个线性投影，其输出通道的数量等于一个patch中的像素值的数量。对解码器的输出进行重构以形成重构图像。本文的损失函数在像素空间中计算重建图像和原始图像之间的均方误差(MSE)。我们只计算屏蔽补丁上的损失，类似于BERT。

        本文还研究了一种变体，其重建目标是每个被屏蔽补丁的归一化像素值。具体来说，本文计算一个patch中所有像素的均值和标准差，并使用它们对该patch进行归一化。在本文的实验中，使用归一化像素作为重建目标提高了表示质量。

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