论文阅读笔记：Masked Autoencoders Are Scalable Vision Learners

---

arXiv:2111.06377v1
Author: Kaiming He; Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollar, Ross Girshick

摘要

本文证明了，在CV领域中，Masked Autoencoders（MAE）是一种scalable的自监督学习器。
MAE方法很简单：我们随机mask掉输入图像的patches并重建这部分丢失的像素。
它基于两个核心设计。首先，我们开发了一种非对称的encoder-decoder结构，其中，encoder仅在可见的（即没有被mask的）patches子集上运行，同时还有一个轻量级decoder，用于从潜在表示和masked tokens重建原始图像。其次，我们发现高比例mask的输入图像（如75%）会产生一项非平凡的且有意义的自监督任务。这两种设计的结合使我们能够高效地训练大型模型：我们加快训练速度（3倍或更多）并提高准确性。我们的可扩展方法允许学习具有良好通用性的高容量模型：例如，在仅使用ImageNet-1K数据的方法中，vanilla ViT巨型模型的精确度最高（87.8%）。下游任务中的迁移性能优于有监督的预训练，并有良好的scaling。

介绍

基于GPT的自回归语言建模和BERT中的Masked Autoencoders的解决方案在概念上很简单：它们删除部分数据并学习预测删除的内容。这些方法现在可以训练包含超过1000亿个参数的可推广NLP模型。

Masked Autoencoders是一种更一般的去噪方式，这种想法很自然，也适用于计算机视觉。事实上，与视觉密切相关的研究[49,39]先于BERT。然而，尽管随着BERT的成功，人们对这一想法产生了极大的兴趣，但视觉中自动编码方法的进展却落后于NLP。我们问：是什么让Masked Autoencoders在视觉和语言之间有所不同？我们试图从以下角度回答这个问题：

1. 直到最近，架构还是不同的。在视觉上，卷积网络在过去十年中占主导地位。卷积通常在规则网格上运行，将“指示符”如masked token或positional embedding集成到卷积网络中并不容易。然而，随着ViT的引入，这一体系结构差距已经得到了解决，并且不再是一个障碍。
2. 语言和视觉的信息密度是不同的。语言是人类产生的信号，具有高度的语义和信息密度。当训练一个模型来预测每个句子中只有几个遗漏的单词时，这个任务似乎会导致复杂的语言理解。相反，图像是具有高度空间冗余的自然信号，例如，缺失的patches可以从相邻patches中重新覆盖，而对局部、对象和场景的高层次理解很少。为了克服这种差异并鼓励学习有用的特性，我们展示了一种简单的策略在计算机视觉中很有效：随机mask大部分patches。这一策略大大减少了冗余，并创建了一个具有挑战性的自我监督任务，需要对低级图像统计之外的整体理解。
   
3. AE的decoder将潜在表示映射回输入，在重建文本和图像之间起着不同的作用。在视觉中，解码器重建像素，因此其输出的语义级别低于普通识别任务。这与语言相反，在语言中，解码器预测包含丰富语义信息的缺失单词。虽然在BERT中，decoder可能是微不足道的（仅一个MLP），但我们发现，对于图像，decoder的设计在确定学习到的潜在表示的语义水平方面起着关键作用。

在这三点分析的推动下，我们提出了一种简单、有效、可扩展的MAE，用于视觉表征学习。我们的MAE从输入图像中随机mask tokens，并在像素空间中重建这些tokens。它有一个非对称的编解码器设计。我们的编码器只在可见的patches子集上运行，我们的解码器是轻量级的，并将潜在表示的输入与masked tokens一起重建（图1）。在我们的非对称编码器-解码器中，将masked tokens转移到小型解码器会大大减少计算量。在这种设计下，非常高的mask率（如75%）可以实现双赢：它优化了精度，同时降低了编码器的量级，使其仅处理一小部分（即25%）补丁。这可以将总体预训练时间减少3倍或更多，同样可以减少内存消耗，使我们能够轻松地将MAE扩展到大型模型。

实现

我们的MAE是一种简单的自动编码方法，它根据原始信号的少部分观测值重建原始信号。像所有的自动编码器一样，我们的方法有一个编码器将观察到的信号映射到潜在表示，还有一个解码器从潜在表示重构原始信号。与经典的自动编码器不同，我们采用了一种非对称设计，允许编码器仅对部分观察信号（无掩码标记）进行操作，并采用一种轻量级解码器，从潜在表示和掩码标记重新构造完整信号。图1说明了下面介绍的想法。

MASKING

按照ViT[16]，我们将图像划分为规则的非重叠的patches。然后，我们对一个子集的patches进行采样，并MASK（即移除）剩余的面片。我们的采样策略很简单：我们在不替换的情况下，按照均匀分布对随机patches进行采样。我们简单地称之为“随机抽样”。

具有高掩蔽率的随机采样在很大程度上消除了冗余，因此产生了一个无法通过从可见相邻面片外推来轻松解决的任务。均匀分布可防止潜在的中心偏移（即，图像中心附近被mask更多）。最后，高度稀疏的输入为设计高效编码器创造了机会。

MAE编码器

我们的编码器使用ViT，但仅适用于可见的、没被mask的patches。就像在标准ViT中一样，我们的编码器通过添加positional embeddings的Linear Projection来得到embedded patches，然后通过一系列transformer block处理结果集。然而，我们的编码器只在整个集合的一小部分上运行。去除masked patches；不使用mask tokens。(也就是说直接去除而非填充零值)，这使我们能够训练非常大的编码器，只需要计算和内存的一部分。之后，整个集合由轻量级解码器处理。

MAE解码器

MAE解码器的输入是由：（i）经过encoder编码的可见patches；（ii）masked tokens，如图1。每个masked tokens都是一个共享的学习向量，表示要预测的缺失patches。我们向这个完整集合中的所有tokens添加positional embedding。解码器具有另一系列Transformer模块。

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

重构目标
我们的MAE通过预测每个masked patches的像素值来重构输入值。解码器输出中的每个元素都是表示patches的像素值向量。解码器的最后一层是一个linear projection，其输出通道的数量等于每个patch中包含的像素数量。解码器的输出形成重构图像。损失函数计算像素空间中重建图像和原始图像之间的均方误差（MSE）：只计算被mask的patches的损失，类似于BERT。

我们还研究了一种变体，其重建目标是每个masked patches的归一化像素值。具体来说，我们计算每个patches中所有像素的平均值和标准偏差，并使用它们来规范化patches。在我们的实验中，使用归一化像素作为重构目标提高了表示质量。

简单的实现

我们的MAE预训练可以被高效实现，并不需要任何专门的稀疏化操作。首先，我们为每个输入patch生成一个对应token（即添一个positional embedding的Linear projection，正如ViT所做的那样）。接下来，我们随机shuffle这些token的列表，并根据dropout rate删除列表的最后那部分。此过程为encoder生成一小部分（未被mask）的token，而无需更换采样的patches。编码后，我们将masked token列表附加到已编码的patch列表中，并unshuffle此完整列表，以将所有token与其目标对齐。解码器应用于这个完整的、添加了poisitional embedding的列表。如前所述，这不需要稀疏操作，仅引入了可忽略不计的开销，因为shuffle和unshuffle的操作速度很快。

在 ImageNet 上的简单测试

我们在ImageNet-1K训练集上进行自监督预训练。然后我们再进行有监督训练，用端到端微调或线性探测（linear probing 以下简称lin）来评估学习到的representations。

Baseline: ViT-Large

我们拿ViT-Large/16作为backbone进行消融实验。ViT-L非常大（比ResNet-50大一个数量级），并且容易过拟合。以下是从零开始训练的ViT-L与从作为baseline的MAE进行fine-tune得到的ViT-L之间的比较：

我们注意到，从无到有地（scatch）训练有监督的ViT-L需要一个强正则化（76.5%→82.5%）。即使如此，我们的MAE预训练也有很大的改进。在这里，微调使用了50个epoch（而从头开始需要200个），这意味着这么高的准确率在很大程度上取决于MAE预训练。

主要结构

我们使用表1中的设置进行消融实验。观察到一些有趣的性质。

Masking Ratio

图5显示了Masking Ratio的影响。最佳的dropout rate出人意料地高：75%对线性评估和fine-tune都有好处。这与BERT相反，BERT中的经典的dropout rate为15%。我们的Masking Ratio也远高于计算机视觉相关工作中的20%至50%。

该模型推断缺失的patches会产生不同但似乎合理的输出。它可以理解物体和场景的格式塔(gestalt)，而这不能简单地通过延伸线条或纹理来完成。我们假设这种类似推理的行为与学习有用的表征有关。

图5还表明线性探测和微调结果遵循不同的趋势。对于线性探测，精度随着掩蔽率的增加而稳定增加，直到达到最低点：精度差距达到∼20%（54.6%对73.5%）。对于微调，结果对比率的敏感度较低，并且大范围的掩蔽率（40–80%）工作良好。图5中的所有微调结果都优于从头开始的训练（82.5%）。

Decoder design

我们的MAE解码器可以灵活设计，如表1的a-b所示。

表1a说明足够深的解码器非常重要。这可以通过像素重建任务和识别任务之间的差距来解释：AE中的最后几层更专门用于重构，而与识别无关。一个合理深度的解码器可以导致重构的特殊性，使潜在的表示处于更抽象的层次。这种设计可以使线性探测提高8%（表1a，“lin”）。但是，如果使用微调，编码器的最后一层可以进行调整以适应识别任务。解码器深度对改善微调的影响较小（表1a，“ft”）。

有趣的是，我们使用单块Transformer构成的解码器依然可以在微调（84.8%）时表现出色。这种小型解码器可以进一步加快训练速度。

在表1b中，我们研究了解码器宽度（信道数）。我们默认使用512-d，它在微调和线性探测中表现良好。

总的来说，我们默认的MAE解码器是轻量级的。它有8个块，宽度为512-d（表1中为灰色）。与ViT-L（24个块，1024-d）相比，处理每个token只需要9%。

Mask token

我们MAE的一个重要设计是跳过编码器中的Mask tokens，表1中用[M]表示，然后将其应用到轻量级解码器中。表1c研究了该设计。

如果编码器使用了mask token，则性能更差：在线性探测中，其精度下降14%。在这种情况下，预训练和部署之间存在差距：此编码器在预训练的输入中有很大一部分是mask token，但这在未损坏的图像中不存在。这可能会降低部署的准确性。通过移除mask token，我们约束编码器始终看到真实的patches，从而提高精度。

此外，通过跳过编码器中的掩码令牌，我们大大减少了训练所需计算。此外，内存大大减少，这可以使训练更大的模型或通过大批量训练加快速度。时间和内存效率使我们的MAE有利于训练非常大的模型。

重构目标

我们比较了表1d中的不同重构目标。到目前为止，我们使用未对每个patch的像素进行normalization作为默认，但是使用标准化后的像素可以提高精度。这种对每个patch进行normalization的方法增强了局部对比度。在另一个变体中，我们在面片空间中执行PCA，并使用最大的PCA系数（此处为96）作为目标，但这样做会降低acc。两个实验都表明高频分量在我们的方法中是有用的。在§5中，我们进一步证明了标记化在迁移学习中是不必要的。

我们基于像素的MAE比标记化简单得多。dVAE标记器还需要一个预训练阶段，这可能取决于额外的数据。dVAE编码器是一个大型卷积网络，增加了大量开销。而使用pixels不会遇到这些问题。

data augmentation

间表1-e，没有必要说明，略。值得注意的是就算不做增强， MAE依然可以有效工作，只是lin指标比较低

Mask sampling

比较了3种：block-wise，grid-wise和random三种采样策略，简单的随机mask效果最好，这是因为只有random才能更好的在非常高的masking ratio下工作。见下图和表1-f

训练过程

到目前为止，我们的消融是基于800-epoch的预训练。
随着训练时间的延长，准确度稳步提高。**事实上，我们甚至在1600个epoch也没有观察到线性探测精度的饱和。**这种行为不同于对比学习的方法。请注意，MAE编码器在每个epoch中仅看到25%的patch，而在对比学习中，编码器在每个epoch中看到200%甚至更多的patch。

4.2 与之前的自监督和有监督学习的结果对比

略，反正就是有图有真相地吹MAE有多好

4.3 部分微调

表1显示，线性探测和微调结果在很大程度上是不相关的。在过去的几年里，线性探测已经成为一种流行的方式；然而，它错过了追求强而非线性特征的机会，而这正是深度学习的优势。我们研究了部分微调方式：微调最后七层，同时冻结其他层。

图9显示了这种改变带来的结果。值得注意的是，仅微调一个变压器块可将精度从73.5%显著提高到81.0%。此外，如果我们只微调最后一个块（即其MLP子块）的“一半”，我们可以得到79.1%，比线性探测要好得多。这种变体本质上是对MLP head进行微调。微调几个块（例如，4或6）可以获得相当高的精度，与冻结主干相比，这仍然是一个微不足道的调整。

在图9中，我们还比较了MoCo v3，这是一种与ViT-L结果对比的方法。仅在不进行任何微调时，它比我们的MAE具有更高的线性探测精度。然而，它的所有部分微调结果都比我们的差。调整4个block时，差距为2.6%。这些结果表明，MAE表示的线性可分离性较差，但它们具有较强的非线性特征，并且在微调非线性的head时表现良好。

这些观察结果表明，线性可分性不是评价表征质量的唯一指标。还观察到线性探测与迁移学习性能（例如对象检测）没有很好的相关性。

迁移学习实验

我们在COCO上评估了目标检测和分割中的迁移学习，并在ADE-20K上评估了语义分割。

目标检测和分割。我们在COCO上对Mask R-CNN[23]进行端到端的微调。ViT主干适合与FPN[31]一起使用（见附录A.3）。我们将此对象检测系统应用于表4中的所有条目。我们报告用于对象检测的box AP和用于实例分割的mask AP。

与有监督的预培训相比，我们的MAE在所有配置下都表现得更好。对于较小的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与监督前培训一致。