MAE(掩码自编码器)是可扩展的计算机视觉自监督学习方法

原文：He, Kaiming, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll'ar and Ross B. Girshick. “Masked Autoencoders Are Scalable Vision Learners.” ArXiv abs/2111.06377 (2021).

1. Abstract

本文证明了掩码自编码器(MAE)是一种可扩展的计算机视觉自监督学习方法。我们的MAE方法很简单：我们对输入图像的patches进行随机掩码，然后重建缺失的像素。MAE基于两个核心设计。首先，我们开发了一个非对称的编码器-解码器架构，其中编码器仅对可见的patches子集(没有掩码的tokens)进行操作，同时还有一个轻量级的解码器，可以从潜在表示和掩码tokens重建原始图像。其次，我们发现对输入图像进行高比例的掩码，例如75%，会产生一项非凡的、有意义的自监督任务。将这两种设计结合起来，使我们能够高效地训练大型模型：加快模型训练速度(3倍甚至更多)并提高精度。我们的可扩展方法允许学习具有良好泛化性能的高容量模型：例如，在仅使用ImageNet-1K数据的方法中，vanilla ViT-Huge模型达到了最高精度(87.8%)。MAE在下游任务上的迁移性能优于有监督的预训练，并且显示出有前景的可扩展行为。

2. Method, Experiment & Result

图1. MAE的架构。在预训练期间，对输入图像的patches进行随机掩码，例如掩码率为75%。编码器应用于可见patches这个小子集上。在编码器之后引入掩码tokens，所有已编码的patches和掩码tokens由一个小型解码器处理，该解码器以像素为单位重建原始图像。在预训练之后，解码器被丢弃，编码器被应用于未损坏的图像，为识别任务生成表示。

图2. 在ImageNet验证图像上的示例结果。对于每个三元组，我们显示掩码图像(左)、MAE重建图像(中)和真实图像(右)。掩码率为80%，196个patches只剩下39个。更多示例见附录。

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

图4. 使用掩码率为75%的预训练MAE重建ImageNet验证图像，但应用于掩码率更高的输入图像。预测结果与原始图像似是而非，表明该方法可以泛化。

图5. 掩码率。高掩码率(75%)适用于微调(顶部)和线性探测(底部)。本文所有图的y轴为ImageNet-1K验证精度(%)。

表1. 使用ViT-L/16在ImageNet-1K上进行MAE消融实验。我们报告了微调(ft)和线性探测(lin)的精度(%)。如未指定，默认值为：解码器的深度为8，宽度为512，重建目标为非标准化像素，数据扩增为随机调整大小的裁剪，掩码率为75%，预训练长度为800个epochs。默认设置以灰色标记。

表2. MAE的训练时间(800个epochs)，用TensorFlow在128个TPU-v3内核上进行基准测试。

图6. 掩码采样策略决定了pretext任务的难度，影响了重建质量和表示(表1f)。每个输出都来自一个使用指定掩码策略训练的MAE。左：随机采样(我们的默认设置)。中：删除大随机块的分块采样。右：网格采样，每四个补丁保留一个。图像来自验证集。

图7. 训练schedules。较长的训练schedule会带来显著改善。这里的每一点都是一个完整的训练schedule。使用的模型为ViT-L，默认设置如表1所示。

表3. 与之前在ImageNet-1K上的结果进行比较。预训练数据是ImageNet-1K训练集。所有自监督方法都通过端到端微调进行评估。ViT的型号为B/16、L/16、H/14。每列的最佳值都加了下划线。

图8. MAE预训练vs.有监督预训练，通过在ImageNet-1K上微调进行评估。我们将MAE的结果与IN1K或JFT300M训练的原始ViT结果进行比较。

图9. 在表1的默认设置下，ViT-L关于Transformer块数量的部分微调结果。MAE表示的线性可分性较差，但如果对一个或多个块进行调整，则始终优于MoCo v3。

表4. 使用ViT Mask R-CNN基线模型进行COCO目标检测和分割。所有条目都基于我们的实现。自监督条目使用无标签的IN1K数据。掩码AP与方框AP的趋势相似。

表5. 使用UperNet进行ADE20K语义分割。BEiT结果是使用官方代码复制的。其他条目基于我们的实现。自监督条目使用无标签的IN1K数据。

表6. 将pixels或tokens作为MAE重建目标的对比。△是使用dVAE tokens和标准化pixels之间的区别。这种差异在统计学上不显著。

表7. 预训练设置。

表8. 端到端微调设置。

表9. 线性探测设置。

表10. 从头开始训练ViT(有监督)。

表11. 掩码编码方法的线性探测结果。我们的微调结果如表3所示。

图10. 在ImageNet验证图像上的随机样本。对于每个三元组，我们显示掩码图像(左)、MAE重建图像(中)和真实图像(右)。掩码率为75%。

图11. 使用ImageNet训练的MAE，在COCO验证图像上的随机样本。对于每个三元组，我们显示掩码图像(左)、MAE重建图像(中)和真实图像(右)。掩码率为75%。

3. Conclusion / Discussion

扩展性好的简单算法是深度学习的核心。在自然语言处理中，简单的自监督学习方法可以从指数扩展的模型中获益。在计算机视觉中，尽管在自监督学习方面取得了一定进展，但实用的预训练范式主要是有监督的。在这项研究中，我们在ImageNet和迁移学习上观察到，自动编码器(一种简单的自监督方法，类似于NLP中的技术)能够提供可扩展的好处。视觉中的自监督学习现在可能正走上与NLP类似的轨道。

另一方面，我们注意到，图像和语言是不同性质的信号，必须认真处理这种差异。图像只是光的记录(without a semantic decomposition into the visual analogue of words)。我们不尝试删除对象，而是删除最有可能不形成语义片段的随机patches。同样地，我们的MAE重建像素，这些像素不是语义实体。然而，我们观察到(例如，图4)，MAE能够推断出复杂的、整体的重建，表明它已经学会了许多视觉概念，即语义。我们假设这种行为是通过MAE内部丰富的隐藏表示产生的。我们希望这一观点将启发未来的工作。

更广泛的影响。MAE基于训练数据集的统计信息预测内容，因此将反映这些数据中的偏见，包括具有负面社会影响的偏见。我们的模型可能会生成不存在的内容。基于MAE生成图像时，这些问题需要进一步的研究和考虑。

关注“多模态人工智能”公众号，一起进步！