MAE 论文笔记

Masked Autoencoders Are Scalable Vision Learners（2021.12.2-v2）

链接：   paper
      code

---

Abstract：Masked autoencoders（MAE，masked自编码器），可扩展自监督学习器。屏蔽输入图像的随机图块(patch)，重建缺失的像素。其基于两个核心设计。1).一种不对称编码器-解码器架构，编码器只在可见的图块（无mask tokens）子集上操作，轻量级解码器在潜在表示和mask tokens上重建原始图像。2).高比例mask输入图像（如75%）。加快了大模型训练速度（3倍或以上），并提高了准确性。可扩展于泛化良好的大体量模型，如ViT-Huge仅使用ImageNet-1K精度达87.8%。在下游任务中的迁移训练性能优于监督预训练，且可缩放(放大效果更明显，同NLP)。
 

一、Introduction

  GPT¹中的基于自回归语言建模和BERT²中的masked自编码，删除了部分数据，并学会预测删除的内容。以此可训练包含1000多亿参数的NLP模型³。

  MAE是一种通用的去噪自编码器⁴，在BERT前就有对其的视觉研究⁵。随BERT的成功，CV上的进展落后于NLP。为什么其在视觉和语言之间有所不同？论文从以下角度展开：

1. 体系结构不同，indicators（“指示符”；如mask tokens⁶或positional embeddings⁷）集成到规则的卷积网络较难。而ViT的出现撼动了CNN的主导地位，架构不再是问题。

2. 信息密度不同。语言是人类产生的信号，具有高度的语义和信息密集性。当训练模型预测有几个缺失单词的句子时，似乎可以诱导复杂的语言理解。相反，图像是具有大量空间冗余的自然信号（如缺失的patch可以相邻的patch作填充），对部分、对象和场景的高级理解很少。故对大部分的随机patch进行mask，大量减少冗余，并创建了一个具挑战性的自监督任务，即需在高于低级图像统计的层面进行全面理解。（重建任务质量参考图1-3）

3. 自编码器的解码器将潜在表现映射回输入。在视觉中，解码器重建像素，故其输出的语义级别低于常见的识别任务。而在语言中，解码器预测包含丰富语义信息的缺失单词。虽然在BERT中，解码器可能不重要(MLP)²，但对于图像，解码器设计在学习潜在表现（representations）的语义层面上很重要。

图1，ImageNet验证集上的示例结果。左为masked image(80%，39/196)（左）、中为MAE重建†，右为原图。更多示例见附录。
注：†，在可见图块上没有计算损失，模型输出质量较差。可以用图块填充来质量(此处未使用，以便展示效果)。

图2，使用在ImageNet上训练的MAE(图1权重)在COCO验证集上的示例结果。右侧两个重建，虽与ground truth不同，但语义上合理。

图3，使用预训练的MAE在ImageNet验证集上进行重建，遮蔽率为75%时效果最佳。

二、 Related Work

  掩码语言建模及其对应自回归，如BERT²和GPT¹。训练预测缺失内容，其均可扩展，预训练适配下游任务。

  自编码含一个将输入映射到潜在表现的编码器和一个重建输入的解码器(如PCA、k-means)⁸。去噪自编码器(DAE)⁴是一类破坏输入信号并学习重建原始未破坏信号的自编码器。一系列方法可以被认为是在不同破坏下的广义DAE，如掩蔽像素⁹或去除颜色通道¹⁰。MAE是其一种形式，但在许多方面与经典DAE不同。

  掩码图像编码方法，即从被掩码的图像中学习表现。将掩码作为DAE中的一种噪声类型¹¹。Context Encoder¹²使用卷积网络修复大范围缺失。受NLP成功启发，最近基于Transformer⁷的相关方法¹³。iGPT¹⁴对像素序列进行运算，并预测未知像素。基于ViT¹⁵研究用于自监督学习的掩码图块预测。最近，BEiT¹⁶提出预测离散tokens¹⁷。

  自监督学习方法在CV领域的工作¹⁸。对比学习¹⁹流行(如²⁰)，其对两个或多个视图之间的图像相似性和不相似性(或仅相似性²¹)进行建模。对比及其相关方法依赖数据增强²²。而自编码追求一个概念上不同的方向。

三、 Approach

  上图为MAE架构。预训练期间，图像图块的较大(如75%)随机子集被屏蔽掉。编码器被应用于可见图块，后引入mask tokens，再由一个小解码器处理，以像素重建原始图像。经过预训练后，去掉解码器，编码器在完整图像上进行识别任务。

  MAE不对称设计，高掩蔽比(图块)随机采样以消除冗余，均匀分布防止了潜在的中心偏移(图块集中于图像中心)，高度稀疏的输入使编码器高效。编码器为 ViT。

  解码器仅在预训练期间用于执行图像重建任务，其架构可以根据编码器设计进行灵活调整。窄且浅的解码器较编码器，每个token的计算量少于10%，大大减少了预训练时间。其输入是由编码的可见图块和mask tokens组成的tokens。每个mask tokens⁶都是一个共享的学习向量，表示要预测的缺失图块。位置embeddings添加于所有tokens上。

  重建目标。MAE通过预测每个遮蔽图块的像素值来重建输入。解码器输出中的每个元素都是表示一个图块的像素值的向量。解码器的最后一层是线性投影，其输出通道的数量等于图块中像素值的数量，其输出被重构为图像。计算重建图像和原始图像的像素空间中的均方误差(MSE)损失，类似BERT²只计算屏蔽图块上的损失。不同于传统的去噪自动编码器⁴计算所有像素上的损失，其选择纯粹由结果驱动，会导致精度略有下降(如～0.5%)。变体，重建目标改为每个掩蔽图块的归一化(图块中所有像素的平均值和标准差)像素值，提高了重建质量。

  简单的实现。先对每个输入图块生成一个token(通过添加位置embedding的线性投影)。后随机打乱token列表，并根据掩码率删除末尾。编码后，将mask tokens的列表append到编码图块的列表中，对其unshuffle(反转随机混洗操作)，以将所有tokens与其目标对齐，再对添加了位置embeddings的列表进行解码。不需要稀疏操作，因shuffle和unshuffle操作很快，其开销可忽略。

 
 
 
ps：下为论文内文献索引，供拓展使用。

---

1. paper[42，43，4] ↩︎ ↩︎

2. paper[14] ↩︎ ↩︎ ↩︎ ↩︎

3. paper[4] ↩︎

4. paper[53] ↩︎ ↩︎ ↩︎

5. paper[54，41] ↩︎

6. paper[14] ↩︎ ↩︎

7. paper[52] ↩︎ ↩︎

8. paper[26] ↩︎

9. paper[54，41，6] ↩︎

10. paper[64] ↩︎

11. paper[54] ↩︎

12. paper[41] ↩︎

13. paper[6，16，2] ↩︎

14. paper[6] ↩︎

15. paper[16] ↩︎

16. paper[2] ↩︎

17. paper[39，45] ↩︎

18. paper[15，55，37，64，40，17] ↩︎

19. paper[3，22] ↩︎

20. paper[56，38，23，7] ↩︎

21. paper[21，8] ↩︎

22. paper[7，21，8] ↩︎