[Transformer]MAE:Masked Autoencoders Are Scalable Vision Learners

From Fackbook AI Research
Author：Kaiming He, Xinlei Chen,Saining Xie, Yanghao Li, Piotr Dollar, Ross Girshick

Paper

Abstract

本文证明了masked autoencoder（MAE）可用于计算机视觉的可扩展自监督学习器。
MAE想法很简单：首先会随机对输入图像的patch施加掩膜，然后重建丢失的像素。MAE基于两个核心设计：
首先本文提出了一个非对称的编码-解码结构，encoder只对未掩盖的patch部分的子集进行操作；然后通过一个轻量级的decoder从隐空间和mask token对图像进行重建。

其次，本文会掩盖大部分输入图像（高达75%），依旧能获得有意义的自监督结果。


结合以上两大设计可以高效训练大型网络，将训练速度提升3x并且提升精度。本文这种可扩展的方法可以学习泛化性良好的高性能模型，仅使用ImageNet-1K数据集训练的ViT-Huge模型就可达到87.8%的精度，并且迁移到下游任务的性能也超过有监督的预训练模型，展现出惊人的可扩展能力。

Section I Introduction

深度学习领域涌现一大批性能、容量不断增长的网络框架，如ResNet,Transformer；得益于硬件算力的提升，如今的网络模型可以轻易消化一百万张图像，甚至需要成百上千万的图像，并向数以亿计的带标签图像进发。
在NLP领域对海量数据的需求可以通过自监督预训练解决，比如GPT和BERT中的masked autoencoding,这些概念实现起来很简单：通过删除部分数据并且学习预测被删除的部分。这样可以用于训练超过千亿参数的泛化NLP模型。




Autoencoding是一种通用的去噪自编码器，也适用于计算机视觉。自从BERT大获成功，人们对其应用于视觉领域也产生了极大的兴趣，但是autoencoding在视觉领域的应用还是落后于NLP。我们不仅有这样一个疑问：是什么造成了这种差距？






本文尝试从以下角度回答这一问题：





（1）网络结构的差异




在计算机视觉领域CNN是主流框架，卷积操作通常在一系列grid上进行，要想将，asked token或者positional embedding集成到CNN并不容易。但是随着ViT的提出，这方面的应用不再存在那么多阻碍。





（2）语言任务和视觉任务中信息密度差异。









语言是人类产生的语义高度集中的交流信号，当训练语言模型预测语句中缺失的单词，需要复杂的语言理解机制。但是图像是自然信号，空间信息高度冗余，当某个patch缺失可以通过相邻的patch进行重建而不需要对其他部分、场景进行理解。




为了克服这种信息密度上差异，鼓励有用特征的学习，本文进行的实验是随机掩盖很大一部分的patch,这种策略很大程度上减少了冗余，变成一个极具挑战的自监督任务，需要对图像形成超越低层信息的全局性理解。可以参考Fig2-4体会一下。





（3）Autoencoder的decoder在语言和视觉任务中的角色是不同的。




在视觉任务中decoder负责将输入的隐空间表示复原回像素级别，其输出的语义水平比一般的识别任务更低；而在语言任务中，decoder需要预测缺失的单词，这是包含丰富语义信息的。虽然BERT中decoder并不十分重要，仅通过一层MLP得到最终输出，但是在数据额任务中decoder至关重要，它决定了能表示的语义水平的高低。





基于以上分析，本文提出一种加单、高效的可扩展的掩膜自编码器（Masked AutoEncoder,MAE）来学习视觉表征。MAE会对输入图像进行随机掩膜，然后重建像素空间中缺失的部分。





MAE是一个非对称的编码器-解码器结构，encoder只对未掩盖的patch进行操作，decoder则负责从隐空间中结合mask token对图像进行重建，参见Fig 1.






Fig 1是MAE的网络框架，预训练时会随机掩盖很大一部分patch(比如75%);然后encoder对未掩盖的patch进行处理；decoder输入的是所有未掩盖的patch和mask token，用于重建输入图像。预训练后，decoder被丢弃然后仅使用encoder对未处理的图像进行处理，学习到的视觉表征用于识别任务。

这种设计掩盖了很高比例的原始图像，这样做是一种双赢：encoder每次只处理一小部分patch，可以将整体预训练时间提升3倍，减少内存需求，同时又可以提升精度，可以将MAE扩展到其他大型模型。

MAE学习到的模型具有很好的泛化性，m模型容量也很大，在预训练阶段可以在ImageNet-1k上训练一些需要大量数据的模型，比如ViT-Large/Huge用于提升泛化性能。比如原始的ViT-Huge微调后的精度可达87.8%,超过了之前所有仅使用ImageNet-1K的模型。


本文还迁移到了目标检测、实例分割和语义分割任务中，本文的MAE均超过了对应的监督预训练结果；更重要的是本文还观察到通过扩展模型可以进一步提升MAE的性能，这些是呀结果与NLP中自监督预训练观察到的结果是一致的。

Section II Related Work

Masked language modeling


BERT,GPT都是NLP领域用于掩码语言建模非常成功的预训练模型，它们会掩盖输入语句很大一部分来训练模型对缺失部分进行预测并且具有很好的可扩展性，大量实验证明这些预训练模型学习到的特征表达可以很好的适用于下游任务。



Autoencoding




自编码是学习特征表达非常经典的一类方法，通过一个编码器将输入映射到隐空间然后通过解码器重建输入，比如PCA,K-MEANS都是自编码器；DAE则是另一类自编码器，输入的是破坏后的输入信号，用于重建未损坏的原始输入。掩膜像素、移除颜色通道都可以看做是DAE的一种，本文的MAE也是DAE的一种形式，但是在许多方面与经典的DAE又存在不同。




Maked image encoding





主要通过掩膜学习图像特征表达，将，masking也看做是DAE噪声的一种，Context Encoder则借助CNN来学习缺失的区域；受Transformer在NLP中的应用，ViT开始自监督学习masked patch，BEit则提出预测离散的token等。





Self-supervised learning





自监督学习在计算机领域中十分重要，主要通过预训练学习不同任务中的pretext，最近对比学习十分受关注，主要从建模图像相似性和不相似性方面进行建模，但是这种方法强烈依赖于数据增强；autoencoding追求的则是概念方向上的不同。

Section III Method

本文的方法遵循基本的自编码器结构，包含encoder和decoder,encoder负责将输入信号映射到隐空间，decoder负责将隐空间的信号重建会原始信号。但是与经典自编码器不同的地方在于本文采用一种非对称的设计，encoder只对未掩盖的patch进行操作，然后使用一个轻量级的decoder根据隐空间的表示和掩码标记重建原始信号。

Masking


本文遵循ViT的流程，首先将图像划分为不重叠patch，然后对path进行随机采样对剩余的部分施加掩膜，采样策略十分简单：基于均匀分布进行随机采样。由于mask的比例很高因此随机采样就很大程度上消除了冗余，这样也无法借助临近patch通过外推进行重建；而均匀分布的随机采样也可以避免出现中心偏置，即图像中间有更多的mask；最后这种高度稀疏的输入有助于设计高效的自编码器。


MAE Encoder



本文encoder使用ViT，仅对未掩膜的patch进行操作。具体设计遵循原始的ViT，将patch经过linear projection并进行位置嵌入；送入Transformer block进行处理。


注意输入的patch是所有patch的子集（25%左右），mask patch会被移除，也不使用mask token。这样仅使用一小部分算力和内存就可以训练较大的encoder。



MAE Decoder




Decoder的输入则是所有token的全集，即已编码的可见patch+mask token。每一个mask token是一个共享的可学习的向量，用来表征该位置是否需要预测一个缺失的patch.




本文给所有token都添加了位置嵌入，不然mask token将不包含它们在图像中的位置信息。Decoder依旧包含一系列Transformer blocks.
注意MAE Decoder只在预训练时期使用，Encoder最后用于产生识别任务的特征表述，因此可以以更加灵活的方式来设计decoder的结构，比如本文尝试了非常小的decoder以及更浅、更窄的decoder结构。比如本文默认decoder每个token的计算量仅为encoder的1/10；这种非对称设计显著减少了预训练时间。





Reconstruction target






重建目标通过预测每一个mask patch的像素值来重建输入，decoder的每一个输出是代表一个patch的一组像素vector，decoder的最后一层是一层线性映射，输出通道数等于patch中的像素数量，因此decoder的输出可以重建回一个patch.
损失函数计算的是重建图像与原始图像的MAE损失，并且仅计算mask patch上的损失，这一点类似BERT。






本文还研究了使用归一化图像像素的效果，可以提高重建效果。







Simple Implementation







MAE预训练十分便捷无需特殊的稀疏操作。








第一步就是对每一个patch经过线性映射+位置嵌入转换成token;









第二步就是随机shuffle并移除最后一部分token,移除比例根据实验设定；









第三步将tokens送入encoder获得编码后的encoded patches；










第四步将encoded patched还原回对应位置，并在缺失部分补上mask patches；










第五步将对齐后的token送入decoder，decoder输出每一个patch对应的像素集合；











第六步计算预测的图像与原始图像之间的MSE loss；












第七步是取训练好的encoder部分作为下游任务的basemodel，在下游任务中进行微调。

Section IV ImageNet Experiments

在ImageNet-1K上预训练，然后在下游任务进行微调。backbone是ViT-Large，比ResNet-50还大因此容易过拟合。需要从头训练+正则化训练效果比较好，微调的话精度很大程度取决于预训练的效果。




Part 1Main Preperties

Table 1展示了消融实验的结果。


Masking ratio



Fig 5展示的是使用线性探测和微调两种不同掩膜率的性能对比，，使用线性探测性能会一直上升，使用微调则在较大的范围内均保持较好的精度（掩膜率40%-80%），但是所有微调的性能都由于线性探测。可以看到无论是线性探测还是微调，掩膜率高达75%时效果最好；这一点与BERT不太一样，BERT最佳的掩膜率是15%；本文也比其他相近视觉任务有区别，它们一般在20%-15%。



从Fig 4可以看出在比较高的掩膜率（75%）会产生不同的但依旧合理的重建输出表示模型有很好的泛化能力，重建的图像不能仅仅通过简单的扩展线条或纹理来获得，本文认为这种是学习到有用的特征表述后的一种类似推理的行为。




Decoder Design

Decoder depth
编码器深度的影响，主要就是Transformer block的数量，decoder深度对线性探测的影响比较大，因为decoder越深，图像重建使用的特征越抽象，与重建越相关，与识别任务可能越不相关。通过加大深度获得更抽象的特征表述可以在线性探测上获得8%的性能提升；但是因为微调还会在识别任务上进行重新学习，因此decoder 深度的影响没有那么大。

很有趣的一点是即使使用single-block在微调上的性能也很优秀，这可以进一步加速预训练。


Decoder width



通道的数目影响的是Decoder width，默认使用512-d，在线性探测和微调上性能都比较出色。
最终decoder包含8个block，通道数是512，比ViT-L每个token的计算量小得多，非常节约算力和内存。

Mask token



本文还探究了encoder是否需要使用mask token.如果encoder使用mask token性能会更差，线性探测中性能会下降14%。主要是因为预训练和最终实时时输入的差异太大，预训练输入的是mask后的结果，应用的时候使用的是原始未mask的结果，二者存在很大的差异。


不使用mask token可以确保encoder看到的都是真实的patch，这样有助于提升精度。
不使用mask token还可以进一步减少计算量，节省下来的算力和内存可以用来训练更大的模型。
Reconstruction target

本文还对比了不同的重建目标的效果，分别是：直接对像素进行重建；对patch内像素归一化后进行重建；对patch像素PCA后取主成分进行重建。效果最好的是直接对像素重建和对patch内归一化后进行重建，表示使用高频分量最有效。
本文还是用了dVAE对token进行预测，效果也不是很好。


Data augmentation


MAE在仅使用cropping的数据增强效果也很好，如果再加入color-jitter范围会下降。令人惊奇的一点是即使不使用数据增强MAE效果也很好，这与对比学习截然相反。


主要是因为在MAE中mask随机施加掩膜的操作就是一种数据增强的操作了，因为每次迭代的掩膜都是不同的，都会生成新的训练样本，这就是很好的数据增强了。



Mask sampling strategy




Fig 6展示的是不同掩膜采样方式的对比结果。有block-wise masking，grid-wise sampling，ramdom sampling。结果显示随机采样效果最好。







Training schedule





Fig 7展示的是训练时长对性能的影响，可以看到随着训练时间的正常准确性稳步提升。实际是1600epochs本文也没有看到精度饱和；这不像对比学习，在300epoch就开始饱和；并且MAE每次只能看到25%的原始patch，而对比学习每次看到的是200%甚至更多的原始patch。

Part 2 Comparison Results

Comparisons with self-supervised method

Table 3展示的是预训练自监督的对比结果。实验结果显示MAE既有很好的可扩展性，模型越大性能稳步提升；并且backbone是原始的ViT，我们认为更先进的网络性能会进一步提升。
MAE即使预训练1600epochs耗费的时间也比其他模型要少。

Comparisons with supervised pre-training


Table 1中表明线性探测与微调之间没有什么关联，线性探测早期比较流行，但是它忽略了非线性特征的作用；本文还探索了部分微调的策略即只微调后面几层，前面层冻住。


Fig 9展示了对比结果，可以看到只微调一个Transformer block就能将精度从73.5%提升至81.0%；如果只微调一个block的一半也能获得79.1%的精度。比线性探测好得多。




Section V Transfer Learning

本文还前移到了目标检测和语义分割任务，使用的是Table 3中预训练的模型。

目标检测结果参见Table 4。
backbone依旧是ViT但结合了FPN进去用于目标检测，可以看到MAE取得了最优的AP指标，比其他监督预训练的结果都要好；在ViT-L上优于监督预训练4个点，而MAE 更简单更快速。


语义分割的结果展示在Table5，可以看到ViT-L的结果比对应的监督学习模型提升了3.7%。




Pixels vs Tokens





Table 6展示的是重建目标选择pixel还是token更有效。尽管使用dVAE编码后的token比使用未归一化的像素效果更好，但是与归一化后的额像素性能相当，因此本文认为MAE使用tokenization用处不大。

Section VI Discussion and Conclusion

扩展性良好的简洁算法是深度算法的核心内容，NLP中一些简单的自监督学习方法可以指数级别的缩放模型；在CV领域预训练主要基于监督学习。本文采取类似NLP中自监督学习的方式在ImageNet上和迁移学习中应用Antuencoder可以灵活的进行扩展，CV中的自监督学习也许能采取和NLP中类似的探索方式。

另一方面，本文注意到图像和语言两种信号的不同，这种差异必须谨慎处理。图像仅仅记录了光信息没有分解成语义部分，而本文的随机移除也尽可能不形成图像中的语义段；而MAE复原的对象也是像素，并不是具有语义的实体。
但实验结果显示MAE依旧能推断出一些复杂的、整体的结构，这也说明了它学习到了视觉概念。我们认为这是通过MAE中丰富的隐层表示学到的，希望这一发现能为未来的工作提供思路。

Broader impacts

本文提出的方法会基于学习到的训练集统计数据来进行预测，因此会包含数据中的偏差、具有副作用的样本等；本文的模型可能会生成不存在的内容，因此当基于本工作生成图像时，上述这些问题有待进一步的探索和考虑。