FAIR 何恺明、Piotr、Ross等新作，MAE才是YYDS！仅用ImageNet1K，Top-1准确率87.8%！

关注公众号，发现CV技术之美

介绍一篇今天新出的论文。

▊ 写在前面

在本文中，作者证明了masked autoencoders（MAE） 是一种可扩展的计算机视觉自监督学习方法。本文的MAE方法很简单：mask输入图像的随机patch，并重建丢失的像素 。它基于两个核心设计的。

首先，作者开发了一种非对称的编码器-解码器体系结构，其中的编码器只对可见的patch子集（没有mask的token）进行操作，同时还有一个轻量级的解码器，该解码器从潜在表示和mask token中重建原始图像。

其次，作者发现mask高比例的输入图像patch（例如75%）会变成一个不错且有意义的自监督任务。将这两种设计结合起来，能够高效地训练大型模型：本文的模型加快训练速度（3倍或更多）并提高精度。

本文的可扩展方法使得学习具有良好通用性的高容量模型：例如，在仅使用ImageNet-1K数据时，ViT-Huge模型的Top-1精确度为87.8%。 此外，下游任务中的迁移性能优于有监督的预训练。

▊ 1. 论文和代码地址

Masked Autoencoders Are Scalable Vision Learners

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

代码地址：未开源

▊ 2. Motivation

深度学习见证了学习能力和容量不断增长的网络结构的爆炸式增长。在硬件快速发展的帮助下，如今的模型很容易拟合超过百万张图像，并开始使用需要数亿张通常无法公开获取的标签图像。

自然语言处理（NLP）通过自监督预训练成功地解决了这种对数据的偏好。基于GPT中的自回归语言建模和BERT中的masked autoencoding的解决方案在概念上很简单：它们删除一部分数据并学习预测删除的内容。这些方法使得现在可以训练包含超过1000亿个参数的NLP模型。

masked autoencoding是一种更为通用的去噪自动编码器（denoising autoencoders），它的思想很自然，也适用于计算机视觉。然而，尽管随着BERT的成功，人们对这一想法产生了极大的兴趣，但视觉中自动编码方法的进展却落后于NLP。作者就提出了一个问题：是什么让masked autoencoding在视觉和语言之间有所不同？ 

作者试图从以下角度回答这个问题：

1）网络结构是不同的 。在视觉上，卷积网络在过去十年中占主导地位。卷积通常在规则网格上运行，并且将诸如mask tokens或位置嵌入等“指示符”集成到卷积网络中并不容易。然而，随着视觉Transformer（ViT）的引入，这一结构差距已经得到解决，不应再构成障碍。

2）语言和视觉的信息密度不同 。语言是人类产生的信号，具有高度的语义和信息密度。当训练一个模型来预测每个句子中只有几个遗漏的单词时，这项任务似乎可以诱导复杂的语言理解。

相反，图像是具有高度空间冗余的自然信号，例如，缺失的patch可以从相邻的patch中恢复。为了克服这一差异并鼓励学习有用的特征，作者展示了一种在计算机视觉中很有效的简单策略：mask大部分随机patch 。这一策略大大减少了冗余，并创建了一个具有挑战性的自监督任务，需要对低级图像统计之外的整体进行理解（如上图所示）。

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

在这一分析的推动下，作者提出了一种简单、有效、可扩展的masked autoencoder（MAE） ，用于视觉表征学习。MAE从输入图像中随机的mask部分patch，并在像素空间中重建丢失的patch。它是一个简单非对称的编码器-解码器设计。

如上图所示，编码器只在可见的patch子集上运行（不带mask tokens），而解码器是轻量级的，并且从潜在表示和mask tokens中重建输入。在非对称编码器-解码器中，小型解码器会大大减少计算量。

在这种设计下，非常高的masking ratio （例如75%）可以实现双赢方案：它优化了精度，同时允许编码器仅处理一小部分（例如25%）patch 。这可以将总体预训练时间减少3倍或更多，同样可以减少显存消耗，使MAE能够轻松地扩展到大模型。

本文的MAE学习非常高容量的模型，这些模型具有很好的通用性。通过MAE预训练，可以在ImageNet-1K上训练像ViT Large/-Large这样的data-hungry模型，从而提高泛化性能。使用vanilla ViT-Huge模型，在ImageNet-1K上进行微调时，实现了87.8%的Top-1精度。

这比以前所有只使用ImageNet-1K数据的结果都要好。作者还评估了迁移到目标检测、实例分割和语义分割的结果。在这些任务中，MAE的预训练比其监督的预训练取得了更好的结果，更重要的是，本文的方法通过扩展模型观察到了显著的性能收益。

▊ 3. 方法

本文的masked autoencoder（MAE）是一种简单的自动编码方法，它根据原始信号的部分观测结果重建原始信号。像所有的自动编码器一样，本文的方法有一个编码器将观察到的信号映射到潜在表示，再用一个解码器从潜在表示重建原始信号。

与经典的自编码器不同，作者采用的设计允许编码器仅对部分观察信号（无mask tokens）进行操作，并采用一个轻量级解码器，从潜在表示和mask token中重建完整信号 （如上图所示）。

Masking

按照ViT，作者将图像划分为规则的非重叠patch。然后对一个子集的patch进行采样，并移除剩余的patch。作者的采样策略很简单：在不替换的情况下，按照均匀分布对随机patch进行采样，称之为“随机抽样”（“random sampling”）。

使用高masking ratio（即移除的patch的比率）的随机采样在很大程度上消除了冗余，因此创建了一个无法通过从可见的相邻patch进行外推能够轻松解决的任务（如上图）。均匀分布可防止潜在的中心偏移（即，图像中心附近有更多mask patch）。最后，高度稀疏的输入为设计高效编码器创造了可能。

MAE encoder

本文的编码器是ViT，但仅用于可见的、无mask的patch。就像在标准ViT中一样，本文的编码器通过线性投影获得token，并加入位置编码，然后通过一系列Transformer块处理。不同的是，本文的编码器只需要在整个集合的一小部分（例如25%）上运行 。这使作者能够用一小部分的计算和显存来训练非常大的编码器。

MAE decoder

MAE解码器的输入是由编码的可见token和mask token组成的全套token。每个mask token都是一个共享的学习向量，表示要预测的缺失patch。作者向这个完整集合中的所有token添加位置嵌入；如果不加入这一点，mask token将没有关于其的位置信息。

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

Reconstruction target

本文的MAE通过预测每个mask patch的像素值来重建输入。解码器输出中的每个元素都是表示patch的像素值向量。解码器的最后一层是一个线性投影，其输出通道的数量等于patch中像素值的数量。解码器的输出被reshape以形成重构图像。计算像素空间中重建图像和原始图像之间的损失函数为均方误差（MSE） 。并只计算mask patch上的损失，类似于BERT。

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

Simple implementation

MAE预训练可以有效实现，而且重要的是，不需要任何专门的稀疏操作。首先，为每个输入patch生成一个token（通过带添加位置嵌入的线性投影）。接下来，随机shuffle这个token列表，并根据masking ratio移除列表的最后一部分。此过程为编码器生成一小部分token，相当于在不替换的情况下采样了token。

编码后，作者将mask token列表附加到已编码patch列表中，并unshuffle整个列表，以将所有token与其目标对齐。解码器应用于这个完整列表（添加了位置嵌入）。如前所述，不需要稀疏操作。这个简单的实现引入了可忽略不计的计算开销，因为shuffle和unshuffle操作速度很快。

▊ 4.实验

4.1. ImageNet Experiments

上表展示了baseline（ViT-L/16）和本文MAE训练的实验对比。

4.1.1. Main Properties

默认设置下，MAE的消融实验结果如上表所示。

Masking ratio

上图显示了mask  ratio的影响。

Decoder design

上表展示了不同decoder深度和宽度的实验结果。

Mask token

上表展示了Encoder编码时是否用mask token的实验结果。

Reconstruction target

上表展示了不同重建目标的实验结果。

Data augmentation

上表展示了不同数据增强的实验结果。

Mask sampling strategy

上图展示了不同采样策略的图示。

上表展示了不同采样策略的结果。

Training schedule

上表展示了训练过程中的log。

4.1.2. Comparisons with Previous Results

Comparisons with self-supervised methods

上表展示了和SOTA方法的对比。

Comparisons with supervised pre-training

上图展示了和监督预训练方法的结果对比。

4.1.3. Partial Fine-tuning

上表展示了ViT-L部分微调的结果。

4.2. Transfer Learning Experiments

Object detection and segmentation

上表展示了COCO object detection 和segmentation结果。

Semantic segmentation

上表展示了ADE20K semantic segmentation的实验结果。

Pixelsvs. tokens

上表展示了Pixels和Tokens作为重建目标的实验结果。

▊ 5. 总结

扩展性好的简单算法是深度学习的核心。在NLP中，简单的自监督学习方法可以指数级别的增益模型。在计算机视觉中，尽管在自监督学习方面取得了进展，但实际的预训练模式仍主要受到监督。在这项研究中，作者在ImageNet和迁移学习中观察到，自动编码器（autoencoder）提供了非常强的优势。视觉中的自监督学习现在可能正走上与NLP类似的轨道。

另一方面，作者注意到，图像和语言是不同性质的信号，必须仔细处理这种差异。图像仅仅是记录下来的光，没有语义分解成文字的视觉模拟。因此，作者没有尝试删除对象，而是删除最有可能不构成语义段的随机patch。同样，MAE重建了像素，而像素不是语义实体。

▊ 作者简介

研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。

知乎/公众号：FightingCV

END

欢迎加入「计算机视觉」交流群👇备注：CV