【论文和代码阅读】Masked Autoencoders Are Scalable Learners (MAE)

写在最前面

先贴一下MAE的论文链接

https://arxiv.org/pdf/2111.06377.pdfhttps://arxiv.org/pdf/2111.06377.pdf紧随其后的是代码复现

https://github.com/facebookresearch/maehttps://github.com/facebookresearch/mae

MAE主要工作：在海量无标注的图像上，通过图像重建任务来训练模型，预训练好的模型通过简单的微调即可适用于下游图像任务（如图像分类）中。

---

1 引言

     近年来Transformer[1]架构在深度学习（Deep Learning，DL）领域如日中天，GPT[2]、BERT[3]等以Transformer为内核的模型在自然语言处理（Natural Language Processing，NLP）领域取得了显著的效果。然而，在计算机视觉（Computer Vision，CV）领域，卷积神经网络（Convolutional Neural Networks，CNN）的架构仍然占据主导地位。尽管ViT[4]将Transformer引入CV并解决了这个架构差异，但是仍然存在信息密度和解码器的差异：CV中图像是一种原始数据，存在较大的空间冗余，相反NLP中语言拥有较高的语义信息；信息密度的差异进一步造成了解码器（decoder）的差异，CV的解码器企图实现较低语义的像素级的重建，而NLP的decoder用以预测较高语义的单词。为了解决这一问题，何凯明等人提出了Masked Autoencoders（MAE[5]）：MAE随机遮掩输入图像的大部分像素块（patch），根据剩余的像素块对图像进行重建，并取得了出色的成绩。

       图1展示了MAE在ImageNet验证集上的重建结果。对于每个三元组，左边的图像是被遮挡的，中间的图像是MAE重建的，右边的图像是实际的。其中掩蔽率为80%，即在196个patch中只剩下39个对模型可见。

图1 MAE效果图

2 架构

      图2展示了MAE的基本框架。对于一张输入的图片，假定其尺寸为80x80，首先将其切分成为25份patch，其中每份patch的尺寸为16x16。接着随机遮掩大部分的patch，将未遮蔽的patch输入到编码器（encoder）进行编码得到token，将可见的token和被遮蔽的token输入到解码器中进行图像重建。

       值得一提的是，虽然MAE的编码器与ViT完全一致，但是由于其在编码前随机遮蔽了大部分的patch，因此模型效率相较于ViT得到了大幅提升。特别地，蓝色虚线框选的decoder部分仅用于预训练（pre-training），目的是通过重建海量图像使模型拥有丰富的先验知识。在具体执行下游任务时，可以将该decoder更换为和实际使用场景适配的decoder。此外，MAE中decoder的transformer模块数小于encoder，单个token的计算量仅为encoder的10%，进一步加速了预训练环节。

图2 MAE架构图

       由于MAE的encoder部分和ViT完全一致，本文通过图3进一步介绍ViT的encoder部分（蓝色虚线框选范围）。对于输入图像，ViT首先将其切分为若干patch，接着将patch输入到线性投射层进行patch embedding操作，为了使patch间彼此知道其所在位置和周围邻居，再对patch进行position embedding操作。由于Transformer的自注意力机制，token间的信息是彼此流通的，同时参考BERT的操作，ViT引入CLS字符用以学习全局图像信息，并将CLS的最后输出用以进行图像分类。不同于MAE，ViT仅有编码器而缺少解码器。

图3 ViT架构图

3 实验

       本节主要展示MAE论文中的部分图片及相关分析。

       图4是MAE原文展示的图片，显示了MAE在部分图像上的重建结果。可以看到MAE在大部分图像上都拥有较好的重建结果。特别地，在红色虚线框选且标号为1的三元图片组中，由于可见patch未能给到小熊的关键信息，因此在重建结果中缺失了旁边的小熊；在红色虚线框选且标号为2的三元图片组中，MAE未能重建出骑马者的小腿部分；在红色虚线框选且标号为3的三元图片组中，MAE为两幅实际缺失头部的卡通人物重建了头部，且表现地较为自然；在红色虚线框选且标号为4的三元图片组中，可见的信息其实完全无法被正常人理解，但是MAE竟能较好地还原；在红色虚线框选且标号为5的三元图片组中，MAE仅根据少量的身体信息，重建出了一只狗，但实际图像是一头牛。

图4 MAE重建图像展示

       为了研究遮蔽率对重建图像性能的影响，何恺明等人通过改变遮蔽率进行了相关实验，结果如图5所示。可以发现，当遮蔽率逐渐增长，图像重建的结果也会越来越偏离实际图像。然而，当遮蔽率达到95%的时候，事实上，汽车、陶器和蘑菇的可见部分已经无法被常人理解了，但是MAE仍能将其较好地重建出来，令人难以置信。

图5 MAE不同遮掩率下的重建结果

4 代码

       为进一步探究MAE论文的真实性及其有效性，本文找到了其开源代码并进行了相关分析和实验。

       在代码demo中，可以通过导入在海量图片上pre-training的模型，对输入的图片进行重建。进一步观察源码，发现作者并没有如其论文描述的，将可见的patch输入encoder，相反，作者将图像的所有patch都输入了encoder。

       如图6，7，8所示：在run_one_image()函数中，代码将完整的图片传入到模型中；在模型的forword()函数中，代码又将完整的图片传给了forward_encoder()函数；在forward_encoder()函数中，代码对完整的图片执行patch embedding操作。显而易见，对完整图片执行patch embedding操作是极其不合理且浪费资源的，因为按原文所说，模型可见的仅为未遮蔽的patch，那么出于学术清白和代码效率考虑，就没有必要对所有patch进行编码。

图6 run_one_image()函数部分代码

图7 forword()函数部分代码

图8 forward_encoder()函数部分代码

       开源demo代码中采用的是一张松鼠图片，为了排除预训练模型对该图片细节已知的可能性，本文从网络上挑选了一张小狗图片，遮蔽85%的patch后，重建结果如图9所示。其中，左1表示原始图片，左2表示遮盖后模型可见的图像，右2表示模型重建后的图像，右1表示模型重建后的图像叠加模型可见的图像。

图9 小狗图片重建

       可以观察到，在仅看见小狗右眼的情况下，MAE模型能够较好的还原出其左眼。特别地，在仅有的3个patch出现一丝白色像素的情况下，MAE就能准确重建出小狗头部的白色竖线毛发。这更让人怀疑其作弊的可能性，MAE模型是否真的对整张图片已知？

       进一步阅读源码发现，给定遮蔽率，对于同一张图片，其遮盖的patch位置是绝对固定的，不会发生变化。于是为了验证MAE模型的真实性，本文将遮蔽后的图片进行保存，并在遮蔽率不变的情况下将其作为输入，这样就能在保持遮盖图像不变的情况下，排除模型提前已知图像的可能性。

      如图10所示，将遮蔽后的图像作为输入，MAE模型仍能实现相同的重建效果。可见transfor-mer结构完全移植到CV上完全可行，MAE利用BERT的思路对patch进行遮掩，能实现非常好的效果。

图10 遮盖后的小狗图片重建

5 思考

       在本节中，本文还使用其他图像测试了代码的重建效果。

       图11展示了围棋被遮掩重建的结果，在重建的结果中其实可以通过黑白的比例来判断某一个棋格上是黑子还是白子。进一步地，如果在预训练阶段使用上万组高手实战棋局，每一组都是上百幅有序图片组成棋局，其中当前棋局作为模型输入，下一手棋局作为模型重建的目标输出，是否有可能训练出一个另类的阿尔法狗？

图11 围棋图像重建

       图12展示了西红柿辣椒图像的重建结果，可以看出MAE根据零散的红色、圆形等语义信息，较好地重建了西红柿图像，但是在重建青色辣椒时，误将上下两只辣椒分割成了左右两只辣椒，但也在情理之中。

图12 西红柿辣椒图像重建

       图13展示了MAE对棋盘边缘的重建结果，可以看到，在仅知道右下角一个patch的情况下，MAE仍能较好地重建出原始图像，并且并未对已知的白色patch做过分的解读。

图13 棋盘边角图像重建

       最后，我认为MAE是一项很新颖的技术，它不应该仅仅用于传统图像识别、图像分类、图像分割上。何恺明等人最初的目的可能仅仅是通过图像重建任务，使模型可以实现自监督训练，从而在CV领域引入预训练环节，通过海量图像让模型进行学习。但在我看来，图像重建工作完全可以在其他领域大展身手，除了前文提到的围棋图像阿尔法狗，MAE还可以用于台风气旋图的预测、画家灵感的启迪、犯罪分子的头像还原等。

参 考 文 献

[1] Llion Jones, Aidan N Gomez, Lukasz Kaiser, Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 2017: 5998-6008.

[2] Alec Radford, Karthik Narasimhan, Tim Salimans, IlyaSutskever. Improving language understanding by generative pre-training. 2018.

[3] Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2019: 4171-4186.

[4] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby. An image is worth 16x16 words: transformers for image recognition at scale. 9th International Conference on Learning Representations, 2021.

[5] Kaiming He, Xinlei Chen, SainingXie, Yanghao Li, Piotr Dollar, Ross Girshick. Masked autoencoders are scalable vision learner. CoRR, 2021, abs/2111.06377.

最后附上一张精简版的海报