Masked Autoencoders Are Scalable Vision Learners

论文地址：Masked Autoencoders Are Scalable Vision Learners
论文代码：MAE—github

Abstract

  在MoCo（未了解过的同学可以移步：Momentum Contrast）这篇论文中，作者最后这样总结道：“MoCo从ImageNet-1M到Instagram-1B的改进一直很明显，但相对较小，这表明更大规模的数据可能没有得到充分利用。我们希望使用高级的pretext task来改善这一点。也就是说除了简单的instance discrimination任务，例如在语言和视觉中还可以采用MoCo来执行诸如masked auto-encoding(MAE，带掩码的自动编码器)之类的pretext task。”因此就有了这篇MAE论文。
  本文证明了MAE是一种可扩展的计算机视觉自监督学习器。MAE方法很简单：屏蔽输入图像的随机patches并重建缺失的像素。它基于两个核心设计：首先，使用了一个非对称编码器-解码器体系结构，其中一个编码器只在可见的patches子集上运行（没有掩码的tokens），另一个轻量级解码器从潜在表示和掩码tokens重建原始图像。其次，作者发现设置输入图像的高比例（例如75%）会产生一项不平凡且有意义的自我监督任务。将这两种设计结合起来能够高效地训练大型模型：加快训练速度（提高3倍或更多，因为掩码比例如果在75%以上，那么只要计算输入图像的1/4或者更少，自然速度就加快了3倍或以上）并提高精度。该可扩展方法允许学习具有良好通用性的高容量模型：例如，在仅使用ImageNet-1K数据的方法中，普通的ViT-H模型的精度最高（87.8%）。下游任务中的迁移性能优于有监督的预训练，并表现出良好的伸缩行为。

Introduction

  深度学习在过去一段时间内有着爆炸式的增长。在硬件快速增长的帮助下，如今的模型很容易超过100万张图像，并开始需要数亿张通常难以获取的标签图像。自然语言处理（NLP）通过自监督的预训练成功地解决了这种对数据的需求。这些解决方案基于GPT中的autoregressive language（自回归语言）建模和BERT中的masked autoencoding，在概念上很简单：它们删除一部分数据，并学习预测删除的内容。这些方法现在可以训练包含超过1000亿个参数的可推广NLP模型。 最近很多研究将BERT应用到cv上面，但是它在cv上的应用不如NLP。于是作者提出问题并从以下3个方面给出自己的解答：是什么使masked autoencoding在视觉和语言之间有所不同？

• 最近这些年cnn在cv上一直占据着主导地位。卷积通常在常规网格上运行，将masked tokens集成到卷积网络中并不容易。因为在transformer中的mask是一个特定词，它会一直保留下来且和其他的词区别开，但是在cnn中卷积窗口在扫过掩码（把一块图片盖住：如把像素换成特定值等）时无法区分它的边界，因为卷积窗口无法像transformer中的那样把它单独作为一个特定的东西拿出来，导致在后面训练的过程中很难还原出这个掩码的具体信息。然而，随着ViT逐渐应用在cv上，这一问题就得到了解决。
• 语言和视觉的信息密度是不同的。语言是人类产生的、语义和信息高度密集的信号。相反，图像是具有高度空间冗余的自然信号。比如说：一个缺失的patch可以从相邻的patch中恢复（比如插值），只需要很少的高级别的理解部分、对象和场景。为了克服这种差异并鼓励学习有用的特性，我们展示了一种简单的策略并在cv中很有效：给很大一部分随机patch加上掩码。这种策略在很大程度上减少了冗余，并创建了一个具有挑战性的自监督任务，这需要对低级别图像统计之外的整体理解。
• 自动编码器的解码器将潜在表示（可以认为是语义空间上的一个表示）映射回输入，在重建文本和图像之间起着不同的作用。在视觉中，解码器重建像素，因此其输出的语义级别低于普通识别任务。这与语言不同，在语言中解码器预测包含丰富语义信息的缺失单词。虽然在BERT中，解码器可能是微不足道的一个MLP，但我们发现对于图像而言解码器设计在确定学习到的潜在表示的语义水平方面起着关键作用。

  在这种分析的推动下，作者提出了一种简单、有效、可扩展的MAE，用于视觉表征学习。MAE从输入图像中给随机patches加上掩码，并在像素空间中重建缺失的patches。它是一个不对称的编码器—解码器设计。编码器只对可见的patches子集进行操作，解码器是轻量级的，可以根据潜在表示和掩码标记重建输入。在非对称的编码器-解码器中，将masked token转移到小型解码器会大大减少计算量。在这种设计下，非常高的masked ratio（例如75%）可以实现双赢：它优化了精度，同时允许编码器只处理一小部分（例如25%）patches。这可以将总体预训练时间减少3倍或更多，同样可以减少内存消耗，能够轻松地将MAE扩展到大型模型。另外，作者还评估了迁移学习在目标检测、实例分割和语义分割方面的应用。在这些任务中，预训练比受监督的预训练获得了更好的结果且更重要的是，通过扩展模型观察到了显著的收益。这些观察结果与NLP自我监督预训练中的观察结果一致。

Approach

  带掩码的自动编码器（MAE）是一种简单的自动编码方法，根据原始信号的部分观察结果重建原始信号。与所有自动编码器一样，本文方法也有一个编码器，将观察到的信号映射到潜在表示，还有一个解码器，从潜在表示中重建原始信号。与经典的自动编码器不同，作者采用了一种非对称设计，允许编码器仅对部分观察到的信号（不带掩码的tokens）进行操作，并采用一个轻量级解码器，从潜在表示和掩码tokens中重建完整信号。

Masking

  作者遵循ViT将图像分割为规则的非重叠的patches，然后对一部分补丁进行采样并给剩余的patches加上掩码（移除patches）。本文的采样策略很简单：在不替换的情况下遵循均匀分布随机采样，简单地称之为“随机抽样”。比例高masking的随机采样（即移除的patches的比率）在很大程度上消除了冗余，因此产生了一个无法通过从可见的相邻patches中推测来轻松解决的任务。均匀分布可防止潜在的中心偏移（即图像中心附近有更多masking的patches）。最后，高度稀疏的输入为设计高效编码器创造了机会。

MAE encoding

  本文的编码器是ViT，但仅适用于可见的、未加掩码的patches。就像在标准ViT中一样，编码器通过把位置编码信息的线性投影嵌入patches，然后通过一系列transformer块处理结果集。然而，编码器只在完整集合的一小部分（例如25%）上运行。去除了带掩码的patches，使用了不带掩码的tokens。这使我们能够仅仅只用一小部分计算量和内存来训练非常大的编码器。最后整个集合由轻量级解码器处理。

MAE decoder

  MAE解码器的输入是由经过编码器编码后的不带掩码的patches和带掩码的tokens组成的集合。每个带掩码的token都是一个共享的学习向量，表示存在的要预测的patches。并向这个完整集合中的所有tokens添加位置嵌入（如果没有这一点，带掩码的tokens将没有关于其在图像中位置的信息）。该解码器也是一系列transformer块。
论文的总体框架如下图：

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

Reconstruction target

  MAE通过预测每个带掩码的patch的像素值来重建输入。解码器输出中的每个元素都是代表一个patch的像素值向量。解码器的最后一层是一个linear projection，其输出通道的数量等于一个patch中像素值的数量。解码器的输出被reshape以形成重建图像。损失函数计算像素空间中重建图像和原始图像之间的均方根误差（mean squared error，MSE：先平方、再平均、然后开方，也称为标准差）。类似于BERT，作者只计算带掩码的patches的损失。作者还研究了一种变体，其重建目标是每个带掩码的patches的归一化像素值（均值为0，方差为1）。具体来说，就是计算一个patch中所有像素的平均值和标准偏差，并使用它们来规范化这个patch。在实验中，使用归一化像素作为重建目标提高了表示质量。

Simple implementation

  首先，作者为每个输入patch生成一个标记（通过添加位置编码信息的线性投影）。接下来随机打乱tokens，并根据masking ratio移除tokens列表的最后一部分。这个过程为编码器生成一小部分tokens，相当于在不替换的情况下采样patches。编码后，我们将带掩码的tokens附加到已编码的patches列表中，并反转随机打乱操作（还原到原来的顺序），以将所有标记与其目标对齐。解码器应用于这个完整列表（添加了位置嵌入）。如前所述，不需要稀疏操作。由于打乱和取消随机打乱操作的速度很快，这个简单的实现带来的开销可以忽略不计。

Experiments

消融实验

  作者在ImageNet-1K训练集上进行自监督的预训练（只使用图片，不使用标签）。然后进行监督训练，用①端到端微调（可以更改模型所有可学习的参数）或②linear probing（只能更改最后一层线性输出层）来评估表征并报告了在中心裁剪的224×224尺度大小的图片上的top-1验证精度。
使用ViT-L作为主干。ViT-L非常大（比ResNet-50大一个数量级），并且容易过度拟合。以下从左到右分别是从零开始训练的ViT-L、在ViT-L之上加了一些强正则化等措施与从基线MAE微调的ViT-L之间的比较：

以下从六个方面来分别做实验来探测影响：

两种方式随着masking的比例大小效果变化：

训练时间：

作者还给出了其它一些实验的对比，比如如何随机采样（block、grid、random）、训练的epoch数以及MAE与有监督预训练方法对比等。

Transfer Learning Experiments

作者 在COCO上评估了目标检测和分割中的迁移学习，并在ADE20K上评估了语义分割。

最后，看一些实验重建的总体效果图：