Masked Autoencoders Are Scalable Vision Learners

摘要

本文提出了一种用于计算机视觉的可扩展自监督学习方法：MAE（Masked Autoencoders，掩膜自编码器）：掩蔽输入图像的随机区块并重建丢失的像素。两个核心理念为：
非对称的编码-解码架构：
编码器的输入为没有被mask的子块；解码器为轻量级（解码器仅在图像重建的预训练中起作用，因此解码器设计可以独立于编码器，且灵活和轻量级），输入为编码器的输入和被mask部分的位置信息，输出为待重建的丢失像素的值。
高比例mask情况下的自监督：
高比例指75%；作者认为这对计算机而言不容易，但有意义。
通过这两个设计，可以学习泛化良好的高容量模型。
论文地址：Masked Autoencoders Are Scalable Vision Learners

问题的提出

目前的模型已经能够消化百万张图像，开始想以亿为计的图像出发。在NLP领域，这种需求已经通过自监督预训练来解决。掩蔽自编码器是一种通用的去噪自编码器，但在计算机视觉方面的发展弱于NLP。因此作者尝试从架构差异、信息密度差异和自编码器的解码器的作用几个方面来研究这种差异形成的原因。
对于图像来说，图像是自然信号，拥有大量的空间冗余。例如，一个缺失的 patch 可以根据相邻的 patch 恢复，而不需要对其他部分、对象和场景有很多的高级理解。为了利用这一点，论文提出了：一个简单的策略在计算机视觉中也能非常有效——掩蔽很大一部分随机 patch。且对于图像，解码器的设计对于学到的潜在表示的语义水平起着关键作用。
基于以上分析，研究者提出了一种简单、有效且可扩展的掩蔽自编码器（MAE）用于视觉表征学习。该 MAE 从输入图像中掩蔽了随机 patch 并重建像素空间中缺失的 patch。它具有非对称的编码器 - 解码器设计。其中，编码器仅对 patch 的可见子集（没有掩码 token）进行操作，解码器则是轻量级的，可以从潜在表征和掩码 token 中重建输入。

结构

网络结构如图1。

图1

基本流程：

• 对图片切分 patch， 随机挑选少部分（比如文中25%）作为网络输入
• 输入通过 encoder 得到对应编码后的 encoded patches
• 将 encoded patches 还原到对应的原始位置，并在缺失的部分补上 masked patches 送入decoder， 每个 decoder 预测对应 patch 的图像像素点
• 计算预测的像素和原始图片的像素之间 MSE 作为 loss
• 取训练完的模型的 encoder 部分作为下游任务的 basemodel 并在下游任务下 finetune

Masking

将输入图像拆分为非重叠块，使用服从均匀分布的无重复随机采样并移除其余块(即Mask)。如图2。

图2

MAE Encoder

MAE中的编码器是一种ViT，但仅作用于可见的未被Mask的块。类似于标准ViT，该编码器通过线性投影于位置嵌入对块进行编码，然后通过一系列Transformer模块进行处理。

MAE Decoder

MAE解码器的输入包含：(1) 编码器的输出；(2) 掩码token。每个掩码Token共享的可学习向量用于指示待预测遗失块。同时，对所有token添加位置嵌入信息。解码器同样包含一系列Transformer模块。

Reconstruction target

MAE通过预测每个掩码块的像素值进行原始信息重建 。解码器的最后一层为线性投影，其输出通道数等于每个块的像素数量。编码器的输出将通过reshape构建重建图像。损失函数则采用了MSE。

实验

论文在ImageNet-1K数据集上进行自监督预训练，然后再通过监督训练评估预训练模型的表达能力。下表是baseline MAE方法的实验结果，可以看到经过MAE预训练后finetune的效果要超过直接从头训练（84.9 vs 82.5）：

可视化效果：

Masking ratio：

Mask比例在75%时，linear probing和fine-tuning下效果才能最好。信息密度越大，我们需要mask的比例越少；信息密度越稀疏，我们需要mask的比例越多。
不同Decoder深度(transformer blocks数量）和宽度（channels数量）的实验结果：