掩码自编码器是可扩展的视觉学习者

paper题目：Masked Autoencoders Are Scalable Vision Learners

paper是FAIR发表在arxiv 2021的工作

paper地址：​​链接​​

Abstract

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

1. Introduction

深度学习已经见证了能力和容量持续增长的架构的爆炸式增长。在硬件快速发展的帮助下，今天的模型可以很容易地过度拟合一百万张图像，并开始要求数以亿计的–通常是公众无法获得的–标记的图像。

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

掩码自编码器的想法，是一种更通用的去噪自编码器的形式，在计算机视觉中也很自然且适用。事实上，与视觉密切相关的研究早于BERT。然而，尽管随着BERT的成功对这一想法产生了浓厚的兴趣，但视觉自编码方法的进展却落后于NLP。我们问：是什么让掩码自动编码在视觉和语言之间有所不同？我们试图从以下几个方面来回答这个问题：

(i) 直到最近，架构还是不同的。在视觉方面，卷积网络 在过去十年中占据主导地位。卷积通常在规则网格上运行，将诸如掩码标记或位置嵌入之类的“指标”集成到卷积网络中并不简单。然而，这一架构差距已通过引入视觉Transformer(ViT)得到解决，不应再成为障碍。

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

图 2. ImageNet 验证图像的示例结果。对于每个三元组，我们展示了掩码图像（左）、MAE重建（中）和ground truth（右）。掩蔽率为80%，196 个切片中只剩下39个。更多示例在附录中。 由于在可见切片上没有计算损失，可见切片上的模型输出质量较差。人们可以简单地用可见的切片覆盖输出以提高视觉质量。作者故意选择不这样做，可以更全面地展示该方法的行为。

图 3. COCO 验证图像的示例结果，使用在 ImageNet 上训练的 MAE（与图 2 中的模型权重相同）。观察最右边两个例子的重建，虽然与基本事实不同，但在语义上是合理的。

图 4. 使用掩蔽率 75% 预训练但应用于具有更高掩蔽率的输入的 MAE 重建 ImageNet 验证图像。预测与原始图像似乎有很大不同，表明该方法可以泛化。

(iii) 自编码器的解码器将潜在表示映射回输入，在重建文本和图像之间扮演不同的角色。在视觉中，解码器重建像素，因此其输出的语义级别低于常见的识别任务。这与语言形成对比，在语言中，解码器预测包含丰富语义信息的缺失词。虽然在 BERT 中，解码器可能很简单（一个 MLP），但作者发现对于图像，解码器设计在确定学习的潜在表示的语义级别方面起着关键作用。

在此分析的推动下，作者提出了一种简单、有效且可扩展的掩码自编码器 (MAE) 形式，用于视觉表示学习。MAE从输入图像中屏蔽随机切片，并在像素空间中重建缺失的切片。它具有不对称的编码器解码器设计。编码器只对可见的切片子集（没有掩码标记）进行操作，解码器是轻量级的，可以从潜在表示中重构输入以及掩码标记（图 1）。在非对称编码器-解码器中将掩码tokens转移到小型解码器会大大减少计算量。在这种设计下，非常高的掩蔽率（例如，75%）可以实现双赢：它优化了准确性，同时允许编码器只处理一小部分（例如，25%）的切片。这可以将整体预训练时间减少 3 倍或更多，同时减少内存消耗，能够轻松地将 MAE 扩展到大型模型。

图 1. MAE架构。在预训练期间，图像切片的大量随机子集（例如 75%）被屏蔽掉。编码器应用于可见切片的小子集。在编码器之后引入掩码标记，完整的编码切片集和掩码标记由一个小型解码器处理，该解码器以像素为单位重建原始图像。预训练后，解码器被丢弃，编码器应用于未损坏的图像（完整的切片集）以进行识别任务。

MAE学习了泛化能力非常高的模型。通过MAE预训练，可以在 ImageNet-1K 上训练数据密集型模型，如 ViT-Large/-Huge，并提高泛化性能。使用 vanilla ViT-Huge 模型，在 ImageNet-1K 上微调时达到 87.8% 的准确率。这优于之前仅使用 ImageNet-1K 数据的所有结果。作者还评估了目标检测、实例分割和语义分割的迁移学习。在这些任务中，作者的预训练比有监督的预训练取得了更好的结果，更重要的是，通过扩大模型观察到了显著的收益。这些观察结果与NLP中自监督预训练中的观察结果一致，希望它们能让CV领域探索类似的轨迹。

2. Related Work

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

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

掩码图像编码方法从被掩蔽破坏的图像中学习表示。 [59] 的开创性工作将掩蔽作为 DAE 中的一种噪声类型。上下文编码器 [46] 使用卷积网络修复大的缺失区域。受 NLP 成功的启发，最近的相关方法 [6, 16, 2] 基于 Transformers [57]。 iGPT [6] 对像素序列进行操作并预测未知像素。 ViT 论文 [16] 研究了用于自监督学习的掩蔽切片预测。最近，BEiT [2] 提出预测离散tokens [44, 50]。

自监督学习方法对计算机视觉有极大的吸引力，通常专注于预训练的不同pre-text任务。最近，对比学习 [3, 22] 很流行，例如 [62, 43, 23, 7]，它模拟两个或多个视角之间的图像相似性和不相似性（或仅相似性 [21, 8]）。对比和相关方法强烈依赖于数据增强。自编码追求一个概念上不同的方向，它表现出我们将要呈现的不同行为。

3. Approach

掩码自编码器 (MAE) 是一种简单的自编码方法，可根据部分观察重建原始信号。像所有自编码器一样，本文的方法有一个将观察到的信号映射到潜在表示的编码器，以及一个从潜在表示重建原始信号的解码器。与经典的自动编码器不同，作者采用非对称设计，允许编码器仅对部分观察到的信号（没有掩码tokens）进行操作，并采用轻量级解码器从潜在表示和掩码tokens重建完整信号。图 1 说明了这个想法，接下来介绍。

掩码。在ViT之后，将图像划分为规则的非重叠切片。然后对一个切片子集进行采样并屏蔽（即删除）剩余的切片。抽样策略很简单：按照均匀分布对随机切片进行抽样而不进行替换。简单地将其称为“随机抽样”。

具有高掩蔽率（即，已删除切片的比率）的随机采样在很大程度上消除了冗余，从而创建了一项无法通过从可见的相邻切片外推来轻松解决的任务（参见图 2-4）。均匀分布可防止潜在的中心偏差（即图像中心附近有更多的掩码切片）。最后，高度稀疏的输入为设计高效编码器创造了机会，接下来介绍。

MAE编码器。本文的编码器是ViT，但仅应用于可见的、未屏蔽的切片。就像在标准 ViT 中一样，本文的编码器通过线性投影嵌入切片并添加位置嵌入，然后通过一系列 Transformer block处理结果集。然而，本文的编码器只对整个集合的一小部分（例如 25%）进行操作。被屏蔽的切片被移除；不使用掩码标记。这能够只用一小部分计算和内存来训练非常大的编码器。整个集合由轻量级解码器处理，如下所述。

MAE解码器。 MAE解码器的输入是由 (i) 编码可见切片和 (ii) 掩码tokens组成的完整tokens集。参见图 1。每个掩码标记是一个共享的学习向量，表示存在要预测的缺失切片。将位置嵌入添加到这个完整集合中的所有标记；没有这个，掩码标记将没有关于它们在图像中的位置的信息。解码器有另一个系列的 Transformer 块。

MAE 解码器仅在预训练期间用于执行图像重建任务（仅编码器用于生成用于识别的图像表示）。因此，解码器架构可以以独立于编码器设计的方式灵活设计。作者尝试了非常小的解码器，比编码器更窄更浅。例如，作者的默认解码器与编码器相比，每个tokens的计算量 <10%。采用这种非对称设计，整个集合tokens仅由轻量级解码器处理，大大减少了预训练时间。

重建target。MAE通过预测每个掩码切片的像素值来重建输入。解码器输出中的每个元素都是代表切片的像素值向量。解码器的最后一层是线性投影，其输出通道数等于切片中像素值的数量。解码器的输出被重新整形以形成重建图像。损失函数计算像素空间中重建图像和原始图像之间的均方误差 (MSE)。只计算掩码切片上的损失，类似于BERT。

作者还研究了一种变体，其重建target是每个掩码切片的归一化像素值。具体来说，计算一个切片中所有像素的平均值和标准偏差，并使用它们来规范化这个切片。在本文的实验中，使用归一化像素作为重建target可以提高表示质量。

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

参考文献

[2] Hangbo Bao, Li Dong, and Furu Wei. BEiT: BERT pre-training of image transformers. arXiv:2106.08254, 2021. Accessed in June 2021.

[6] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In ICML, 2020.

[12] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In CVPR Workshops, 2020.

[44] Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In NeurIPS, 2017.

[46] Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. Context encoders: Feature learning by inpainting. In CVPR, 2016.

[50] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea V oss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In ICML, 2021.

[57] Ashish V aswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, 2017.

[59] Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Y oshua Bengio, Pierre-Antoine Manzagol, and Léon Bottou. Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion. JMLR, 2010.