导读

凯明出品，必属精品。没有花里胡哨的修饰，MAE就是那么简单的强大，即结构简单但可扩展性能强大。MAE通过设计一个非对称的编码解码器，在预训练阶段，通过高比例的掩码原图，将可见部分输入到编码器中；解码器是一个轻量级的transformer stack，融合编码器编码后的表征和mask token，实现了归一化像素的目标重构；在推测识别阶段，去掉解码器，输入完整的图片到decoder中实现识别任务。实验表明，在ImageNet-1K上，MAE实现了自监督方法里最好的性能。

摘要

本文表明，掩码自编码器(Masked Autoendoers, MAE)是可用于计算机视觉的可扩展自监督学习器。MAE手段很简单：我们将输入图像分块并随机掩码，重建丢失的像素。它基于两个核心设计：1、我们开发了一个非对称的编码-解码架构，其中编码器只在可见的图像子块上进行操作，以及一个轻量级的解码器，它从潜在表征和掩码令牌重构原始图像；2、我们发现高比例的图像遮挡，例如75%，会产生一个具有重要意义的自监督任务。结合这两种设计使我们能够有效并高效地训练大模型：我们提升训练速度三倍之多并提升精度。我们可扩展方法允许学习具有良好泛化的大容量模型：例如，一个原始的ViT-Huge模型只使用ImageNet-1K就能获取最好的87.8%的精度。在下游任务中的迁移性能超过有监督的预训练并呈现出有希望的可扩展行为。

MAE结构：在预训练阶段，图像块的很大一部分比例（例如75%）被遮挡。编码器应用在可见的这部分图像块。掩码令牌在编码器之后被引入进来，与被编码的子块一起被一个小型的解码器处理，这个解码器用于重构原图的像素。预训练后，这个解码器被丢弃，编码器被应用于未损坏的图像去实现识别任务。

1.介绍

深度学习见证了容量和能力不断增长的架构的爆炸式发展，如从AlexNet，ResNet再到Transformer。借助硬件的快速发展，今天模型能够在一百万张图像上快速过拟合，并且开始需要成百上千万张不可实现的标注样本。

对于数据的渴望，已经在自然语言理解(natural language processing, NLP)中通过自注意的预训练成功解决。这些方案，基于GPT中的自回归语言模型和BERT中的掩码自编码，在概念上很简单：他们去掉数据的一部分，并去学习去掉的内容。这些方法现在能够训练具有百万参数的可泛化的NLP模型。

MAE的思路，一种更通用的去噪自编码器，同样自然地适用于计算机视觉。实际上，视觉中的相关研究工作比BERT先出现。然而，尽管随着BERT的成功对这一想法产生巨大的兴趣，视觉中MAE的进展却落后于NLP。**我们不禁会问：什么让MAE在视觉和语言之间有所不同？**我们尝试从以下几个方面来回答这个问题：

直到最近，两者的架构还是不同的。在视觉中，卷积网络在过去的十年里呈统治地位。卷积通常在规则网格上操作，而将掩码令牌或位置编码等“指示器”应用到卷积网中并不直接。但是，这个架构的鸿沟已通过引入ViT得到解决，并不再成为障碍。
信息密度在图像和语言是不一样的。语言是人类产生的信号，是高语义和密集信息的。当训练一个模型去预测每个句子中少量缺失的单词时，这个问题会引发复杂的语言理解。反而，图像是自然的信号，带有严重的空间冗余，例如，一个缺失块能够从相邻像素块中复原，而对局部、物体和场景的高层理解很少。为了克服这个差别和鼓励学习有用的特征，我们展示了一个简单的策略在CV中工作良好：遮挡很大部分的随机块。这个策略很大程度上降低冗余，和创造了一个有挑战的自监督任务，它需要整体的理解，超越低级的图像统计。
这个自编码器的解码器将隐含的表征映射回输入，在重建文本和图像中扮演着不同的角色。在视觉中，解码器重建像素，因此它的输出比通常的识别任务是更低的语义级别。这与语言截然相反，在语言中解码器预测丢失的单词包含有丰富的语义信息。虽然在BERT中，解码器可能是很简单的一个MLP，我们发现对于图像，解码器的设计在决定学习到的隐含表征的语义级别中，扮演着一个核心角色。

在这些分析的推动下，我们展示了一个简单的有效的、可扩展的MAE形式，用于视觉表征学习。我们MAE在输入图像上遮挡随机的图像块，并在像素空间重构遮挡掉的像素块。它具有非对称的编码-解码设计。我们编码器仅仅在可见的图像块上（不带有掩码令牌）操作，并且我们的解码器是轻量级的，它能够从隐藏的表征和掩码令牌中重构输入。在这种非对称的编码解码器中将掩码令牌转移到小型的解码器中，能够大大减少运算量。在这种设计下，一个很高的掩码比例，如75%可以实现双赢局面：它优化了精度，并使得编码器处理较少部分的像素块。这能够将整体预训练时间减少三倍之多，同时减少内存消耗，使我们能够轻松扩招MAE到大模型上。

我们的MAE学习了泛化好、高容量的模型。通过MAE预训练，我们能够在ImageNet-1K上训练如ViT-Large/Huge等数据渴望型的模型，并提高了泛化能力。使用朴素的ViT-Huge模型，我们能够在ImageNet-1K上微调并取得了87.8%的精度。这个性能超过当前只使用ImageNet-1K的所有结果。我们同样评估了迁移学习在目标检测，实力分割和语义分割。在这些人物上，我们预训练取得了超过它使用有监督的预训练的同类。更重要的是，通过扩大模型能够取得显著的提升。这些观察与NLP中自监督的预训练的结果相一致，我们希望他们能够带领我们的领域去探索类似的轨迹。

2.相关工作

Masked language modeling. 掩码语言建模和它的自回归同类，如BERT和GPT在NLP中用于预训练非常成功。这些方法保留输入序列的一部分，并训练模型去预测缺失的内容。这些方法展示了很好的扩展性，并且大量证据表明这些预训练的表征能够很好泛化到不少下游任务。

Autoencoding. 自编码是一种经典的学习表征的方法。它有一个编码器能够将输入映射到一个隐含的表征，和一个解码器能够重构输入。例如，PCA和K-means都是自编码器。去噪自编码器（Denoising autoencoders，DAE）是一种自编码器，它破坏输入信号并学习重建原始的未破坏的信号。很多方法可以看成一个广义的DAE，在不同破坏程度下，例如掩码像素，或去掉颜色通道等。我们的MAE是DAE的一种形式，但与经典的DAE存在很多不同。

Masked image encoding. 掩盖图像编码方法从被掩盖破坏的图像中去学习表征。先驱工作（来自Bengio组）将掩码当做DAE中的一种噪音类型。Context encoder使用卷积网络来修复大量缺失的区域。受NLP中成功方法的启发，相关工作基于Transformer去完成的。iGPT在像素序列上运行，去预测未知的像素。ViT论文通过自监督学习去学习遮挡块的预测。最近，BEiT提出预测离散令牌。

Self-supervised Learning. 自监督学习方法在计算机视觉中呈现了极大的兴趣，通过关注预训练的预文本任务。最近，对比学习很流行，很多论文尝试建模两个或多个视图之间的相似性或非相似性。对比和相关方法强烈依赖数据增强。自编码追求一个概念上不同的方法，它呈现出不同的行为，就是我们接下来将要展示的。

方法

MAE是一个简单的自编码器方法，能够重构原始的信号当给出它的部分观测。就像所有的自编码器，我们的方法有一个编码器能够将观测到的信号映射到一个隐藏的表征，和一个解码器能够从这个隐藏的表征中重构原始的信号。不同于典型的自编码器，我们采用一个非对称的设计，如图1所示。

Masking.

跟随ViT，我们将一幅图片分割成规则的非重叠的子块。然后我们采用一部分子块并遮挡剩下的图像块。我们采样的策略非常直接：我们随机采样图像块而没有采用替代，遵循一个均匀分布。我们简单地将其称为随机抽样。
带有高比例遮挡的随机抽样大大地消除冗余，因此创造了一个任务，它不能够简单从可见的相邻块外推出来。这个均匀分布防止潜在的重心偏差，即图像重心具有更多的遮挡块。最后，高度稀疏的输入为设计高效编码器创造了机会，如下。

MAE encoder.

我们的编码器是一个ViT但是仅用于可见的未遮挡的部分。正如一个标准的ViT，我们的编码器通过带有位置嵌入的线性投影来嵌入图像块，然后一系类的Transformer来处理结果集。但是我们的编码器只在很小的一部分子集，例如25%的全集上操作。遮挡的子块被移除；不使用掩码标记。这允许我们去训练一个很大的编码器，，只使用一小部分计算和内存资源。全集被一个轻量的解码器来处理，如下。

MAE decoder.

MAE解码器的输入是整个标记，包含可见块被编码后的隐表征和遮挡的标记。每一个遮挡标记是一个共享的可学习的向量，能够指示一个可预测的消失子块的存在。我们给所有标记添加了位置嵌入；没有这个，遮挡标记将没有他们的位置信息。这个解码器具有其他的一系列Transformer结构。
这个MAE解码器仅仅用在预训练阶段去执行图像重建任务，只有编码器被用于识别而产生图像表征。因此，这个解码器架构可以灵活设计，以一种独立于编码器设计的方式。我们实验采用非常小的解码器，比编码器更窄更浅。例如，我们默认的解码器在每个标记上的运算量不到解码器上的10%。带有这种非对称的设计，整个标记集合只被一个轻量的编码器所处理，这在很大程度上降低预训练时间。

Reconstruction Target.

我们的MAE通过预测每个遮挡标记的像素值来重构输入。在解码器输出的每个元素都是一个向量，表示一个子块的像素值。解码器的最后一层是一个线性投影，它的输出通道数与子块的像素个数相等。这个解码器的输出被重塑成一幅图像。我们的损失函数计算重构和原始图像在像素空间的均方误差，即MSE。我们仅在遮挡的部分计算误差，类似于BERT。
我们同样学习了一个变量，它的重构目标是每一个遮挡块的归一化的像素值。特别地，我们计算一个遮挡块中所有像素的平均值与标准偏差，来归一化这个块。在我们的实验中，使用归一化的像素值作为重构目标，能够提升表征质量。

Simple implementation.

我们的MAE预训练操作是可以高效地实现的，重要的是，它不需要任何特别的稀疏操作。首先，针对每个输入块，我们产生了一个对应的令牌标记，可以通过线性投影和添加位置编码信息。然后，我们随机打乱这个标记序列，并根据遮挡比例去掉列表的最后一部分。这个操作为编码器产生了一小部分令牌子集，相当于在不替换的情况下对子块采样。在编码后，我们想编码后的块中添加遮挡令牌序列，并且恢复整个列表的序列，让令牌与他们的目标块对应起来（随机打乱的逆操作）。解码器被应用在这个恢复的带有位置信息的全序列上。如上所述，不需要任何的稀疏操作。这个简单的实现引入了可忽略的开销，因为混淆和反混淆操作是快速的。

ImageNet Experiments.

我们在ImageNet-1K上进行自监督的预训练。然后我们采用有监督的训练去评估表征，通过两种方式：端到端的微调，线性探测。我们汇报了在224×224图像上的Top1的准确率。具体细节如下：

BaseLine：ViT-Large. 我们采用ViT-L作为消融试验的骨架模型。ViT-L非常大，几乎比ResNet-50大一个数量级，并容易过拟合。以下表格是是否采用MAE的性能对比：

策略	Top 1 Acc(%)
原始的ViT	76.5
重新配置菜谱，从头训练ViT	82.5
使用MAE，微调训练好的ViT	84.9

我们发现从头开始训练一个有监督的ViT-L并不简单（nontrivial），我们为ViT-L设计了一套强正则化的参数配置，能够取得82.5%的精度。尽管如此，我们的MAE预训练进一步提升到84.9%的精度。这里，微调只采用了50个epoch，而从头开始训练需要200个epoch，表明微调的精度高度依赖预训练。

主要属性分析.

炼丹的精华，只管采用这个配方即可。

Masking Ratio.
掩码比例为什么固定在75%：75%的高比例掩码对于模型微调或者线性探测都工作良好。
Decoder设计
Decoder设计比较灵活，深度上，比如堆叠几个Transformer块；宽度上，输出多少个通道，通过对比不同数据的结果，数据为王强者胜，没有为什么。
Mask Token
输入到Decoder是否需要mask token呢？通过对比发现，跳过mask token而直接把可见的子块输入到编码器中能取得更好的结果。
不难发现，输入到decoder的输入不带有mask token能取得很大程度的性能提升和节省运算量。两全其美。
Reconstruction Target & Data augmentation.
Mask sampling strategy.
遮挡采样的策略也在此处进行了对比。
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-pjHRxHga-1671717112476)(https://www.liangtengyu.com:9998/images/image_e9edc4fc.png)]
遮挡采样策略决定预文本任务难度，影响重构质量和表征。对比发现，随机采样的重构效果最好，accuracy结果也最好，具体看文中数据。

ViT 架构.

我们采用标准的ViT结构，它具有一些列堆叠起来的Transformer块，每一个Transformer块包含多头自注意块和一个MLP块，都使用LN，即LayerNorm。我们的encoder终止于LN层。由于MAE的decoder和encoder具有不同的宽度，我们在编码器后采用一个线性映射层来匹配它。MAE向编码器和解码器中都添加了位置信息，即Sine-Cosine版本，没有使用相对位置或者层缩放，即Layer scaling（这个某些transformer中被使用）。
我们从编码器的输出中提取特征，用于微调和线性推测。由于ViT含有一个类标签去适应这个设计，在我们的MAE预训练中，我们添加了一个辅助的虚假的令牌到编码器输入中。这个令牌将被当成类标记，用于训练微调或线性推测的分类器。我们的MAE不使用这些令牌，而采用平均池化也同样工作良好。

Result on ImageNet-1K.

Comparison with self-supervised methods. 我们对比自监督ViT模型的微调结果。从以下表格中可以看出，对于ViT-B，所有方法性能几乎相近；对于ViT-L，不同方法之间的差距比较大，表明对于更大模型的挑战是减低过拟合。

预训练数据都是来自于ImageNet-1K训练集，除了BEiT中的tokenizer在250M DALLE上预训练。所有自监督方法都是通过端到端微调来评估的。我们的MAE重构归一化的像素值，并且预训练了1600个epoch。

通过以上表格，我们MAE可以轻松扩展，并在更大模型中展示稳定的提升。我们使用ViT-H获取了86.95%的精度（在224尺度上）。通过更大尺寸如448，我们获得了87.8%的精度，仅仅使用ImageNet-1K数据集。当前所有方法中，只使用ImageNet-1K，获得最好的精度是87.1%而且还是在512尺寸下。
在ImageNet-1k这个高度竞争的数据集上，我们比当前最先进的方法还提升了不小的差距。我们的结果结果就是在朴素的ViT上进行的，如果采用更先进的网络，将会提升更多。

和BEiT对比，我们MAE更加准确，同时更简单更快。我们方法是重构像素，对比BEiT是重构令牌tokens：如果BEiT使用ViT-B重构像素，将会有1.8%性能下降。我们不需要dVAE预训练，同时比BEiT快3.5倍之多。
MAE模型预训练了1600个epoch达到了更好的效果。即使如此，我们整体的预训练时间也比其他模型要少，在相同的硬件条件下。例如，当训练ViT-L时，我们采用128个TPU-v3核处理，1600个epoch需要31个小时；而MoCo v3进行300个epoch就需要36h。

当你好奇，凯明大神怎么那么6的时候，不妨看下凯大神的利器：128个TPU v3核；而256个它就能在7分钟内训练完ResNet50。这还怎么玩？这还怎么玩？这还怎么玩？

Comparison with supervised methods. 在原始的ViT文章中，ViT-L性能下降当在ImageNet-1K上训练时。我们设计有监督的参数菜谱能够工作更好，但精度饱和。我们的MAE预训练，仅仅使用ImageNet-1K能够更好泛化：从头开始训练的收益，对更大容量的模型而更高。表明我们的MAE能够帮助扩展模型大小。

Transfer Learning Experiments

我们评估迁移学习在下游任务中使用预训练好的模型。
Object detection and segmentation. 我们在COCO上以端到端的方式微调了Mask R-CNN。ViT骨架被用来和FPN使用。

讨论与总结

深度学习的核心是算法简单并扩展性好。在NLP中，简单的自监督学习方法可以从指数级的缩放模型中获益；在计算机视觉中，尽管在自监督学习中取得了进展，但实际预训练范式还是有监督主导的。在这个方法中，我们在ImageNet和迁移学习任务中看到，一个自编码器以一个简单的自学习方法，类似于NLP中的技术，能够提供可扩展的收益。视觉中的自我监督学习，可能走上了与NLP相似的轨迹。

另外，我们注意到图像和语言是不同本质的信号，这个差异必须要谨慎对待。图像只是记录下来的光线，没有语义分解到单词的视觉类似物。没有尝试去删除物品，反而随机删掉图像块，他们最可能不回形成语义段。同样，我们的MAE重构像素，同样也不是语义实体。然后，我们观察MAE推断出复杂的整体的重构，表明它学习到大量视觉概念，即语义。我们假设这样行为发生是通过MAE内部丰富的隐藏的表征的方式。我们希望这个观点能够激发未来的工作。