论文阅读笔记：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在视觉和语言之间有所不同？我们试图从以下角度回答这个问题：
- 直到最近，架构还是不同的。在视觉上，卷积网络在过去十年中占主导地位。卷积通常在规则网格上运行，将“指示符”如masked token或positional embedding集成到卷积网络中并不容易。然而，随着ViT的引入，这一体系结构差距已经得到了解决，并且不再是一个障碍。
- 语言和视觉的信息密度是不同的。语言是人类产生的信号，具有高度的语义和信息密度。当训练一个模型来预测每个句子中只有几个遗漏的单词时，这个任务似乎会导致复杂的语言理解。相反，图像是具有高度空间冗余的自然信号，例如，缺失的patches可以从相邻patches中重新覆盖，而对局部、对象和场景的高层次理解很少。为了克服这种差异并鼓励学习有用的特性，我们展示了一种简单的策略在计算机视觉中很有效：随机mask大部分patches。这一策略大大减少了冗余，并创建了一个具有挑战性的自我监督任务，需要对低级图像统计之外的整体理解。
- 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训练集上进行自监督预训练。然后我们再进行有监督训练，用端到端微调或线性评估来评估学习到的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%。