MAE论文详解

---

前言

 MAE是Facebook团队在2021年11月发布的一篇论文，《Masked Autoencoders Are Scalable Vision Learners》，带掩膜的自编码器是可扩展的视觉学习器，MAE就是Masked AutoEncoders的缩写，一作是何凯明大神。MAE通过在图片上进行掩码实现在CV领域进行大规模的自监督训练，并且仅通过在ImageNet-IK数据集上进行训练实现了最好的准确率，超过了当时的有监督训练方法。接下来对这篇论文进行详细介绍。

一、MAE理论

 在NLP领域，自监督训练获得了很大的成功，例如BERT和GPT等。ViT实现了将Transformer应用在CV领域，但是其仍然是有监督训练。MAE的目的便是使CV领域可以像NLP中BERT那样通过预测掩码（类似完型填空）实现大规模的自监督预训练。因此，MAE可以看作是BERT的CV版本。那么在CV领域使用预测掩码机制的困难是什么呢？

一、模型架构。在NLP领域用的是Transformer架构，而CV领域中用的多是CNN架构。不过ViT的出现使得Transformer可以用在CV领域中从而解决了这个问题。
二、信息冗余。语言文字具有很强的信息密度和语义性，而图片中有着很多的冗余信息，在预测掩码时模型可以根据周围像素进行插值还原，使得模型预测难度降低，难以学到有用的特征表示。
三、解码器。在BERT中，最终的预测任务是一个词，可以看作的分类任务，因此解码器使用一个简单的MLP层就可以。而在CV中，预测的是大量的像素点，因此解码器的结构应该更加复杂。

 上述三个问题中，第一个问题已经被ViT解决。对于第二个问题，MAE通过增大掩码的比例，减少图片中的冗余信息，迫使模型学习到有用的特征表示，在论文中，MAE对于图片中的75%都打上了掩码。对于第三个问题，在CV中预测掩码相当于是进行图片重构，因此不能用简单的MLP实现，MAE使用的一系列Transformer块堆叠而成。

 总结一下就是MAE有两个核心理念：一：在模型结构上使用非对称的编码器-解码器结构来进行自监督训练。二：在训练时对于图片中的大部分区域(75%)都打上Mask。

二、MAE整体框架

 接下来我们介绍一下MAE的整体框架。

整体流程
 一、首先对于输入的图片，将其划分为Patch，然后随机将75%的Patch打上Mask。
 二、将没有被Mask的图片输入到编码器中进行编码。
 三、 将编码器中的输出和之前被Mask的Patch组合起来输入到解码器中进行解码重构。
 四、将重构后的图片与输入的原图计算MSE Loss。
编码器
 MAE的编码器是ViT，输入是没有被打上Mask的Patch（带位置编码）。
解码器
 MAE的解码器是堆叠的Transformer层，输入由两部分组成：1.编码器的输出。2.Mask的tokens。每一个mask token 都是一个共享的、学习的向量，它指示了这里有一个待预测的 token。然后将这两部分组合起来并加入位置编码告诉模型每一个mask tokens所对应的图像中的位置。
 注意：MAE是一个非对称的结构，其编码器是要比解码器要大，这么设计有两个原因：一：编码器的输入是所有Patch中的25%，而解码器的输入是所有Patch，解码器的计算量相比较要大一些，为了减少训练时间，解码器应该设计的小一点，而编码器由于计算量较少，所以模型结构设计大一点增加特征提取能力也可以接受。二：在进行预测时只用到了编码器进行特征提取，而并没有用到解码器，因此在设计时应该重点关注编码器的能力。

重构目标
 MAE预测的是每一个Patch区域中的像素值。解码器的输出元素(tokens)中，每一个元素是一个向量，代表的是一个掩码Patch中的像素信息表示，然后将其通过一个线性层将元素的向量维度映射成与Patch中的像素数量一样，最后将其Reshape成图片，与原图计算MSELoss。在计算Loss时，只计算掩码区域的Loss。

三、MAE简单实现

接下来介绍一下MAE的训练流程实现。
1 输入一张照片，将其划分为Patch。
2 将这些Patchs经过一个全连接层变成tokens。
3 将这些tokens加入位置编码。
4 将tokens随机打乱，将打乱后的前25%tokens输入的解码器中。
5 将编码器的输出与可学习的mask tokens组合在一起并且恢复成原来在图片中的位置顺序，使解码器的输入与目标图像中的Patch相对应，然后对这些tokens加入位置编码。
6 经过解码器。
7 将解码器的输出Reshape成图片。

四、实验

 MAE论文中做了很多实验，有关于模型结构的，模型微调的等等。详细的实验结果大家可以去论文中看一下。
 有一个实验是关于编码器的输入的，如果编码器的输入不包含Mask，与编码器输入中包含Mask时相比其不但速度更快，并且精度更高。通过实验说明了为什么编码器的输入中不包含Mask。

总结

 MAE将自监督的学习方法用到了CV领域中并且取得了很好的效果，在ImageNet-1K的数据集上进行训练效果超越了有监督的方法，并且其训练成本与其他大模型相比不算大。ViT的出现是把Transformer用到了CV领域，而MAE则是把NLP中的大模型预训练方法用到了CV领域，实现了CV领域的大模型训练。