【超分辨率】何凯明新作:Masked Autoencoders Are Scalable Vision Learners

0.前言

本文提出了一种新的自编码器：MAE。即可以说本文提出了一种新的自编码器，也可以说本文在计算机视觉领域提出了一个新的任务，即通过掩盖图像的一部分来利用没有掩盖的部分对图像进行重建。
这个新的编码器和新的任务也就构成了一种新的自编码方案，如下图：

这个东西在标题中被称为可扩展的视觉学习器，对于这个可扩展，我的可能不太准确的理解是本文定义的这个任务和自编码器，形成了一种新的图像处理方案，它不仅仅只是用于对于图像丢失像素的重建这个任务上，还能够用在像图像分类、图像分割等领域中，所以是可扩展的，但具体是如何扩展的，以我现在的认知还无法想象到。但是在原文中配图的解释中，有这样一句话：
After pre-training, the decoder is discarded and the encoder is applied to uncorrupted images to produce representations for recognition tasks.
意思大概是在预训练之后，编码器就可以单独使用在没有被掩盖的图像上来产生用于识别任务的某种编码表示，这也就体现了文章标题所说的可扩展性。

MAE方案有两个核心点：
第一：本文构建了一个非对称的编码器-解码器结构，用一个比较复杂的编码器来编码未被掩盖的图像块，用一个比较轻量级的解码器来重建原始图像，解码是从未被掩盖的图像块和掩盖的标记来解码重建原图像的，未被掩盖的图像块很好理解，掩盖的标记其实就是下图的这些部分 ：

也就是说解码的过程并不是只利用编码器编码的内容，还要利用掩盖图像的标记块，关于这些灰色的方块到底是以什么值输入编码器的，暂时还不知道，感觉可能应该是全0或全1.这个样一个自编码器就是本文的第一个核心点。
第二：本文发现掩盖一个图像的高比例部分（例如75%）会产生一个非常有意义且不平凡的自监督任务（简单来说就是一个非常了不起的自监督任务），也就是掩盖大部分来进行自监督恢复。

这两个核心点，让本文提出的自监督方案可以加速大模型的训练和提高准确率（这也是一种可扩展的体现）。作者举例说利用这个方案使得ViT-Huge模型仅仅使用ImageNet-1k数据就达到了87.89%的准确率。

1.介绍

深度学习发展迅速，在硬件的帮助下，现在训练百万甚至更多的数据的网络也成为可能，但这通常也有可能带来过拟合问题以及带标签的图像获取困难的问题。
在NLP中，这种对数据的需求已经被自监督学习的方法所解决。例如基于自回归的GPT和基于自编码的BERT，它们的思想十分简单，就是移除数据的一部分然后通过模型恢复它们。这种方法现在使得能够训练出参数量达到千亿级别的NLP模型。
其实这种掩码的自编码器在CV上的研究是要先于BERT的，作者在此提出了一个问题：视觉和语言的自编码器的区别究竟在哪里？
作者尝试从下面三个部分来回答这个问题：
(1)结构不同。近十年来在视觉领域在占主导地位的都是卷积神经网络，卷积只是在规则的网格上运行，并没有集成掩码标志和位置编码（mask tokens and positional embeddings, NLP里常用的操作）。这种结构上的不同已经被最近的VIT给解决了。
(2)信息密度不同。语言是人类产生的一种信号，它的信息密度非常大，几乎没有多少冗余，每个单词和句子都含有丰富的语义信息，冗余非常小，因此如果移除一部分信息，恢复出来的信息可能和原来完全不同。而图像则不同，图像的空间冗余很大，图像中的某个像素或部分是可以通过它相邻的像素或部分来恢复重建。为了减小这种信息密度的不同，让视觉的模型学习到更多有用的特征，本文才提出了这样的策略：随机掩盖图像的很大一部分从而减小冗余性，并且这也创造了一个有挑战性的自监督任务，要求对图像有更全面的理解而不仅仅是低级别的图像特征统计。

这里我的理解是：这样的一个任务也反映了图像中确实存在大量的冗余信息，也就是我们在处理图像时，把很多的资源或算力都用在了处理冗余信息，所以本文提出的方案几乎可以扩展到CV领域的各个方面，减少资源和算力的浪费。
(3)语言和图像自编码器的解码器不同。图像自编码器中的解码器主要是恢复像素，相比识别任务的模型，解码器的输出其实是有更少的语义信息。比如超分辨任务，仅仅是恢复分辨率，并没有改变图像的语义。但分类任务其实最终输出的是这个图像的类别，这种类别也可以看做是一种语义信息。而在NLP中，每个单词都有丰富的语义，所以恢复单词的解码器自然也就能够恢复出丰富的语义信息。作者发现，在图像中，解码器的设计非常关键，它决定了学到特征表示的语义级别。

经过上面的分析，作者就提出了MAE。MAE的编码器只需要处理图像的一小部分，这可以大大减小大模型预训练的时间，并且还能在减小内存消耗的同时提高准确率。

2.相关工作

掩码语言模型和对应的自回归模型包括BERT和GPT已经是比较成功的方法。

自编码是一种用于学习特征表示的比较经典的方法。PCA和K-means是自编码器，去噪编码器（DAE）也是一种经典的编码器，它用于重建图像的像素。本文的MAE也是一种DAE，但是在许多方面都与DAE有所不同。

掩码图像编码是从被掩盖的图像中学习特征表示。DAE是从被噪声掩盖的图像中学习特征表示。近来还有一些视觉上受到NLP启发的工作：iGPT、ViT、BEiT等等。
自监督

自监督学习在CV中也十分重要。比较学习和自编码都算自监督学习的领域，但是其实是两个不同的方向。

3.方法

这个图，再贴到这里。

随机掩盖：本文和ViT一样，也是将图像分成了许多不重叠的图像块，然后随机的掩盖一些块（或者说是随机的移除一些块）。本文采用均匀分布的方式随机移除图像块。
移除高比例的图像块可以很好地减少图像的冗余，均匀分布保证了大部分被遮挡的图像块都靠近中心。
MAE编码器：本文的编码器实际上就是一个ViT（ViT指的是vision transformer结构，是在论文《An image is worth 16x16 words: Transformers for image recognition at scale.》中提出的，可参考论文原文或我的博文《【图像Transformer论文理解】AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE》），但是本文却也和ViT不同，因为本文只输入没有被掩盖的一些图像块，这些图像块仅仅占一小部分，这就使得我们只需要消耗很少的计算资源和内存资源去训练一个非常大的网络。
MAE解码器：解码器的输入包括：（i）编码的可见图像块（ii）掩码标记（就是之前图中的那些灰色的方块）。编码的可见图像块就是编码器的输出。而掩码标记究竟是个啥东西，目前在这里还理解的不是很清楚，原文如下：
Each mask token is a shared, learned vector that indicates the presence of a missing patch to be predicted.
看起来这些灰色的方块是参数共享的，并且是可以学习的。
输入编码器之前，还要给所有的块增加位置编码，这样所有的图像块才能知道他们在图像中的位置信息。
MAE的解码器只是在预训练的时候才被使用，在进行图像识别时，仅仅使用编码器编码出的图像表示就可以了。所以解码器的设计比较简单，往往取决于编码器的复杂程度。本文设计的解码器的计算复杂度基本要小于解码器的10%，这样就可以在预训练时大大减少训练时间。
重建目标：编码器的最后一层采用线性映射，原文说映射出的向量的通道数等于一个块中的像素数，这里这个通道数是啥意思我没太理解，不过我感觉应该是经过最后一层的线性变换，每个掩码块会最终输出一个对应的1维向量，这个向量的长度就是这个块中的像素总数，然后reshape一下就恢复成了图像块。训练时作者采用MSE损失，仅仅计算掩码块的损失（重建的掩码块和之前的掩码块）。
这里还有个小点就是本文经过实验发现，对图像块进行标准化会提高恢复像素的质量。
简单实现：本文的方案实现起来也比较简单，首先通过线性映射和位置编码给每个块都生成一个token，然后随机shuffle这些token形成一个序列，按照比例来移除序列中最后一部分的token，然后就是进行编码和解码的操作了。解码前还有进行unshuffle，根据位置编码重新按照原来的顺序排序。

4.ImageNet实验

作者先是使用MAE来进行预训练，然后用预训练好的编码器来分别进行端到端的微调和线性映射两种方式测试验证集top1的准确率。（端到端的微调用的其实是ViT-Large，而线性映射是指一些类似SVM的线性分类器）
baseline采用的是ViT-Large，VIT-Large是要比ResNet-50更大的网络，很容易过拟合。
这里作者比较了训练ViT-Large预训练和没预训练得到的精度（验证集top1的准确率），如下：

76.5%是ViT原文的准确率(没有经过预训练)，而82.5%是作者复现的准确率，84.9%是增加了MAE预训练的准确率。经过预训练的网络只需经过50epoches的微调就能达到84.9%的准确率，然而没经过预训练则需要训练200epoch。
至于作者复现的比原文好的结果，作者在附加材料中给出了训练的参数。

4.1 主要属性

作者做了许多消融实验。
掩码比例：
这个实验表明75%的掩码比例对于微调和线性分类来说效果最好，这个比例是比较高的，远远高于一些近期的工作。同时这个实验也表明了对于线性分类和端到端的微调来说，掩码率的影响不同，线性分类对于掩码率的变化更敏感。

上面这个实验表明高于75%的掩码率会让恢复的图像和原图像有些不同，但是似乎也合理，作者假设产生这种合理性的原因和网络学习有用的表示有关。
解码器设计：解码器可以灵活设计，如下表：

上表（a）中，blocks代表解码器中的的transformer块的数量，这个表明深的解码器能够提高线性分类器的准确率，作者解释说因为更深的解码器会使得编码器学习到更加抽象的表示，而这种抽象的表示对线性分类更有帮助，对端到端微调的结构影响很有限。
同时，只有1个block的解码器能让ft的准确率达到84.8%，原文提到：因为1个Transforme块是从视觉token到掩码token信息传播的最小要求，这句话是原文说的，我不太理解，主要不能太理解这个token的含义……
上表（b）中，dim表示编码器的通道数（具体含义应该和Transformer结构有关），通道数512时效果是最好的。
通过上面的实验，作者最终确定的MAE解码器默认的参数时8个块和512维，它的浮点数运算量只有ViT-L的9%。
Mask token：先说明一点，原文中所提到的token，我都不是能够很好地理解，因为token是NLP里的一个概念，在NLP处理语言时我还能勉强理解，但是在视觉中总感觉理解的不是很透彻。这里我们可以简单地把Mask token看做是下图中的灰色方块。
本文的一个重要设计就是在编码时，没有用到这些灰色的方块，而在解码时却使用了。如果编码时也用到的话，效果会变差并且浮点数运算会增加，如下表：

字母的具体含义不是很能理解，不过第一行代表的是编码时也用到灰色方块，这导致ft的性能变差，lin的性能甚至变差了14%，浮点数运算还增加了三倍多。
重建目标：作者对重建不同的目标也做了实验，在这之前的那三个实验都是没有对像素进行标准化的，所以作者分别实验的重建目标有：标准化后的图像块（不含mask token和含有mask token的）的像素，图像块的PCA，以及图像块的dVAE token，最终的实验结果如下：

第一行是重建标准化之后的像素（编码时没用到灰色方块，只编码显式的图像块），第二行是编码时用到了灰色方块的结果，这表明标准化对结果有很大的提升，作者分析可能是标准化能够提高图像块的对比度。最终确定把像素作为恢复的目标，当然恢复的要是含有mask token的目标。图中灰色标注的区域为默认的设置，至于为什么不用第二行作为默认设置，是因为第二行的设置用到灰色方块和显式的图像块一起编码，虽然精度挺高，但是浮点数运算却多了三倍多（前面说过的那个表里面有体现），计算复杂度太高了。
数据扩充：数据扩充的实验结果如下：

从这个实验结果能看出即使没有数据扩充，MAE也能保持比较高的性能。这是因为在MAE的整个训练过程中，每一代都是重新随机的进行MASK，这其实也是一种数据扩充，尽管没有对原始数据进行扩充，但是每一代都产生了新的训练样本。
Mask采样策略：对于mask的不同采样策略，作者也做了实验，结果如下：


最终确定随机采样的策略。
训练长度：之前的消融实验的训练长度都是800epoch，对于不同的训练长度作者也做了不同的实验，如下：

更长的训练能够获得更高的性能，本文之前的消融实验默认是训练800个epoch。

4.2 与之前结果的比较

与其它自监督方法的比较：

上面的实验室训练1600epoches的结果，虽然训练长度比较大，但是时间还是比较短的，使用相同的TPU，MAE训练1600epoches需要31小时，但是MoCo v3训练300epoches就需要36小时。
与有监督预训练的比较：

4.3 部分微调

作者微调的操作采用部分微调的方法：冻结其它部分仅仅微调最后几层。
微调不同的层数作者也做了不同的实验：

注意这个图前面最小微调的尺度不是1个block，而是0.5个block，也就是仅仅微调一个Transformer block中的MLP的sub-block部分，就能够达到79.1%的精度。

5.迁移学习的实验

目标检测与实例分割：

语义分割：
在迁移学习的预训练中，对于重建不同的目标，作者又做了不同的实验：

6.讨论和总结

在本文中，作者用一种简单的类似于NLP的处理方式来对图像进行自监督学习，在ImageNet和迁移学习上做了很多实验都证明这种处理方式很有效。这也许可以说明，在视觉上的自监督学习现在可能会走上一条与NLP类似的道路。
但是，图像和语言的差别还是挺大的，图像仅仅对光进行记录，而语言则是记录下来了许多语义信息，这种区别是需要被重点关注的。本文MAE正是关注到了这种区别，所以没有直接重建图像语义，而是重建图像的像素。而且MAE推理除了非常复杂和全面的重建像素，可能是因为它学习到了比较多的视觉语义。作者假定这种学习语义的行为发生在MAE中的隐藏层表示中，作者希望这种假定能够启发未来的相关研究工作。