一文解读Masked Autoencoder(MAE)

前言

  论文链接：https://arxiv.org/pdf/2111.06377.pdf
  跟李沐学AI：https://www.bilibili.com/video/BV1sq4y1q77t?spm_id_from=333.999.0.0
  如果说Vision Transformer是Transformer在CV领域的拓展，那么Masked Autoencoder就是BERT在CV领域的拓展。MAE使用类似于BERT的掩码机制，从图片中随机抹去一些像素，并让模型通过已知像素去构建未知像素，从而迫使模型学习图像中的特征。实验证明：MAE具有很好的像素重构能力。接下来，就开始论文的讲解吧！（注：对Vision Transformer感兴趣的小伙伴可以移步到上一篇博客：https://blog.csdn.net/abbcdc/article/details/123639566?spm=1001.2014.3001.5501）

Masked Autoencoder

Architecture

  由于模型的思想基于Vision Transformer，所以这里先回顾一下：Vision Transformer（ViT）打破了CNN在计算机视觉领域的统治地位，仅使用一个标准的Transformer Encoder，并在大规模数据集预训练的情况下，就能达到和CNN一样甚至更好的效果。具体来说：ViT模型将一张图像分割成一个个小块，每个小块作为一个单位（类似于NLP中句子的一个个单词）然后将这些小块按位置排列成一个序列，送入Transformer Encoder中获取图像信息。假设原始图像像素为224×224，块的大小为16×16，序列的长度就为224²/16²=196，相比于将像素作为单位时的序列长度224²=50176减少了几百倍，而这时的序列长度就在Transformer的承载范围之内了。
  在ViT论文中的最后一段，作者说明他们尝试了使用自监督方式预训练模型，即样本和标签均来自同一物体，但效果不尽人意，因此他们认为仍然应该使用有监督+大规模数据集的预训练方式使模型效果更佳。而Masked Autoencoder提出了新的策略，证明了当使用自监督+不那么大规模的数据集对ViT模型进行预训练，可以达到较优的效果。MAE的思想如下：

将图像分割成一个个小块后，随机对其中的一些小块进行遮蔽，然后将没有遮蔽的像素块按顺序排列，送入Transformer Encoder中，获得特征向量；之后将遮蔽的像素块（只包含位置信息）按原来的位置插入特征向量中，再放入Decoder中，由Decoder重构像素信息，生成原始图片。注意到，图中的Encoder比Decoder大一些，是因为Encoder是一个标准的Transformer编码器，计算量比较大，而Decoder结构就比较简单了。当用于不同的下游任务时，Decoder可以替换成任意结构，而Encoder作为特征提取器需要保持不变。
  总结：Masked Autoencoder使用了掩码机制，利用编码器将像素信息映射为语义空间中的特征向量，而使用解码器重构原始空间中的像素。MAE使用的是非对称的Encoder-Decoder架构，即编码器只能看到未被遮蔽的部分像素块信息，以节省计算开销，而解码器解码的是所有像素块的特征信息。

Approach

  上面是对模型整体架构的介绍，接下来详细讲解各个部分。

1. Masking
     MAE使用了ViT的方法，将图片分割成一个个小块，然后在这些小块中随机、均匀地选取一部分保留，剩下的全部遮蔽。作者强调了要遮蔽大量的像素块（约75%），从而减少像素块之间的冗余信息，使整个任务更具有挑战性，迫使模型去学习图像的全局特征而非局部特征，从而获得更优的图像重构能力。
2. MAE Encoder
     Encoder即一个标准的ViT模型中的Transformer Encoder，架构如下：
   
   不同的是，这里的输入不是全部像素块而只包含未被遮蔽的25%的像素块，节省了计算开销。
3. MAE Decoder
     由于解码器需要重构那些被遮蔽的像素，因此输入包含两部分的信息：未被遮蔽的像素块通过编码器生成的特征向量和被遮蔽的像素块信息。被遮蔽的像素块信息全部通过一个相同的可被学习的特征向量表示。 在这里，Decoder实质上还是一个包含Transformer block的架构，因此输入向量包含了位置编码（同Encoder，Transformer无法学习位置信息，因此需要加入位置信息表明某个像素块在原始图像中的位置）。该解码器只在预训练阶段使用，在迁移到下游任务时，解码器可以根据用户的实际需求替换成任意架构。
4. Reconstruction target
     解码器的最后一层是一个线性层。假设patch的大小是16×16，则线性层的输出维度就为256，然后再reshape为16×16，则得到了重构后的像素块。损失函数是MSE函数，即原始像素和重构后的像素相减再求平方和。只需要在被遮蔽的像素块上计算损失，因为未被遮蔽的像素块信息已经作为输入被编码器和解码器知晓。
5. Simple implementation
     作者简单说明了整个实现流程。将输入序列进行Linear Projection并加上位置编码形成一个个token后，进行shuffle操作，然后选取前25%的token作为Encoder的输入，便完成了随机采样过程。将Encoder的输出和表示被遮蔽的像素块的向量拼接，进行unshuffle操作，即还原原始像素块的排列顺序，然后再加入位置编码，送入Decoder中，完成像素重构。

ImageNet Experiments

Main Properties

  作者对比了使用自监督预训练的MAE和使用有监督预训练的ViT的性能（MAE和ViT的关系：MAE实际上是加上了类似于BERT的掩码机制用于自监督预训练的ViT）。对于MAE，作者将模型在ImageNet-1k数据集上进行自监督预训练，然后使用有监督训练对模型进行fine-tune，对于ViT，作者使用的是ViT-L/16在ImageNet-1k数据集上进行有监督训练。对比结果如下：

在ViT-L/16上使用有监督训练，准确率是76.5%，加上了正则化约束后，准确率提升到了82.5%。而对MAE使用自监督预训练，仅使用50个epoch进行fine-tune，效果就提升到了84.9%并保持稳定，说明了MAE的有效性。

  接下来看一下论文中的主要实验结果，即上面6个表格。第一个表格表示MAE的性能和decoder的Transformer block的数量的关系，ft即end-to-end fine-tuning，表示fine-tune时改变MAE每个层的参数，而lin即linear probing，即在fine-tune时将编码器的参数冻结，直接使用编码器提取的特征，只改变解码器（线性层）的参数。结果表明，当block数量为8时，ft和lin上的效果均达到最佳。第二个表格表示MAE性能和decoder的宽度即每个token的长度的关系，结果表明，当每个token的长度为512时，ft和lin上的效果最优。第三个表格表示编码器的输入只包含未被遮蔽的像素块和包含所有像素块两种情况下MAE的性能，结果表明当输入只包含未被遮蔽的像素块时，模型精度更高且计算量大大减少，节省了计算开销。第四个表格表示预训练阶段解码器执行不同任务对MAE的性能的影响，第一行表示重构像素并且不对像素块做归一化处理，第二行表示重构像素并对像素块做归一化处理，第三行表示进行PCA降维，第四行表示对每一个像素块进行一次映射到token的操作，结果表示重构像素并进行归一化这种方式的效果最佳。第五个表格表示不同的数据增强技术对MAE性能的影响，结果表明对图像进行随机大小的裁剪这种方式带来的性能最佳。
最后一个表格表示输入图像的mask的方式对MAE性能的影响，共有random、block、grid三种方式，它们的效果可以参考上图；ratio表示遮蔽率；结果表明，当使用random的遮蔽方式且遮蔽75%的像素块时，ft和lin上的效果最佳。
  下图表示mask的遮蔽率对MAE性能的影响，结果比较直观，这里就不展开说了（fine-tuning即end-to-end的fine-tune，linear probing即只改变线性层参数的fine-tune）。

  下图表示不同预训练轮数对fine-tune效果的影响，结果表明，随着预训练的epoch数量的增加，通过两种fine-tune方式得到的模型精度都在不断上升，并没有出现一个饱和的现象。

Comparisons with Previous Results

  下表就是和之前的一些方法的对比，可以看出MAE的性能是最佳的。（训练设置：预训练数据是ImageNet-1K训练集（除了BeiT中的tokenizer使用了250M的DALLE数据进行预训练）。所有的自监督方法都是通过端到端的微调进行评估的。ViT型号为B/16、L/16、H/14。所有模型的输入图像大小都是224×224，ViT-H还额外使用了大小为448×448的输入图像。在这里，MAE重建归一化像素，并使用1600个epoch进行预训练。）
  下图展示了ViT模型在ImageNet-1k和JFT300M数据集上进行有监督预训练，MAE在ImageNet-1k数据集上进行自监督预训练，均使用ImageNet-1k数据集进行fine-tune后模型的性能对比。可以看出，当ViT使用了JFT300M数据集（大小是ImageNet-1k的300倍）进行预训练后，效果稍优于MAE，但两者的训练成本相差巨大，由此说明MAE还是具有一定优势的，同样地，在ViT上如果也使用ImageNet-1k数据集进行预训练，效果就明显差于MAE了。

Partial Fine-tuning

  下图表示了当在ImageNet-1k数据集上进行fine-tune时，调整的Transformer block的个数和模型微调后在该任务上的精度的关系。调整0个block时即表示linear probing，调整24个block时即表示end-to-end fine-tuning。可以看出，当使用linear probing时，MAE的效果是要差于MoCo V3的（MoCo是一种使用对比学习进行自监督预训练的方法，详情可以阅读论文：https://arxiv.org/pdf/1911.05722.pdf），但是当调整了1个或多个block时，MAE的性能就优于MoCo，并且不断上升后到达一个平台。这说明了底层的block学习到的是低级特征，和特定任务关联性不大，但高层的block与特定任务息息相关，因此对这些高层block的微调对模型在特定任务上的性能影响是很大的（可以看到前期曲线上升幅度很大，就是微调了高层block的结果）。

Transfer Learning Experiments

  作者最后进行了模型在迁移学习任务上的性能对比的实验，由于这部分的意义只是为了展示结果，和模型原理关系不大，这里就不一一展开了，把实验结果图放在下面，大家看看就好。

Conclusion

后言

  如果觉得博主写得还不错，请点赞加收藏呀！也欢迎在评论区里提问交流！