【论文阅读】Masked Autoencoders Are Scalable Vision Learners（MAE）

1.Abstract

文章链接：Masked Autoencoders Are Scalable Vision Learners

• 证明了 MAE 是计算机视觉上可扩展的自监督学习
• 我们对输入图像的随机块进行 mask，重建缺失的像素点，这基于两种核心设计
  • 提出了非对称的 encoder-decoder 结构
    • encoder只针对可见的 patches 进行操作
    • 一个轻量级的 decoder，从潜在表示 和 mask tokens 中重建原始图像
  • 对 input image 进行大比例的 mask（eg.75%），产生了不容易的并且有意义的自监督任务
• 上面两种设计是的我们能够高效得训练大模型：我们将训练加速了三倍多，并且提高了准确度。
• 我们的可扩展方法允许学习出泛化性很好的大容量模型：一个 vanilla ViT-Huge model 达到了 87.8%，这是在只使用 ImageNet-1K 数据集上的最好结果
• 在下游任务上迁移的效果超过了有监督的预训练，并且显示了可扩展能力

2.Introduction

现在随着硬件的升级，很多模型都需要大量数据

对数据的大量需求在 NLP 领域通过自监督预训练成功解决了。这些解决方案，是基于 GPT 中的自回归 language model 和 BERT 中的 masked autoencoding。这些概念很简单：移除一部分数据，并学习去预测这些被移除的内容。这些方法现在能够训练包含超过1000亿个参数的可推广的NLP模型。

mask autoencoder的想法（一种更通用的降噪 autoencoder），在计算机视觉领域也是很适用的。

在 BERT 成功之后，尽管人们对这一想法产生极大的兴趣，但是视觉上的 autoencodeing 方法仍然落后与 NLP。那么，视觉和语言上的 masked autoencoding 区别在哪呢？

• 视觉领域，卷积网络是上一个十年的主流。卷积通常在规则的网格上运行，将mask token 或 positional embedding 等“指标”集成到卷积网络中并不简单。然而，这个架构上的差距已经通过 Vision transformer (ViT) 的引入得到了解决，它不再是一个障碍。

• 语言和视觉的信息密度是不同的。语言是人类产生的信号，具有高度语义和信息密度。当训练一个模型来预测每句只有少数缺失单词的句子时，这个任务似乎能诱导复杂的语言理解。相反，图像是具有大量空间冗余的自然信号-- 一个缺失的 patch 可以从邻近的 patch 中覆盖，这些补丁对部件、对象和场景的理解很少。为了克服这种缺点，并鼓励学习有用的特征，我们一个简单有效的策略：随机 make 很大比例的 patch。这个策略大幅度减少了冗余，并且创造了一项具有挑战性的自监督任务，需要超越低级图像统计的整体理解

• 自动编码器的 decoder ，将潜在的表示映射回输入，在重建文本和图像之间扮演着不同的角色。在视觉上，decoder重建像素，因此，它的输出比一般的识别任务具有更低的语义层次。这与语言相反，在语言中，解码器会预测包含丰富语义信息的缺失单词。虽然在 BERT 中，decoder是微不足道的（一个 MLP），我们发现对于图像来说，decoder 的设计在决定（学习的潜在表征的）语义层次上，发挥关键作用。

基于以上分析，对于视觉 representation 的学习，我们提出了一种简单，高效，可扩展形式的 masked autoencoder（MAE）。

我们的 MAE 随机遮住输入图像的一些块，并且在像素空间上重建这些损失的块。这里包含一个非对称的encoder-decoder设计。我们的 encoder 值处理 patchs 的可见部分，而 decoder 是轻量级的，并且从隐含的 representation 和 mask token 中重建输入图像。将 mask token 移动到我们这个小的 decoder 中，大大减少了计算量（因为处理所有 patches 的 decoder 比较小，比较大的 encoder 只处理一小部分 patches）

在这种设计下，高比例的 mask（75%），可以实现双赢：

• 优化了准确度
• 允许 encoder 只处理一小部分 patches

这可以将整个训练前的时间减少3倍或更多，同时也可以减少内存消耗，使我们能够轻松地将MAE扩展到大型模型

我们的 MAE 学习了泛化性很好的大容量模型。使用 MAE pre-training，我们可以在ImageNet-1K上训练数据需求量很大的模型（比如 ViT-Large）。我们在 ImageNet-1K 上 fine-tune ViT-Huge模型，达到了 87.8%的准确度。这超过了使用 ImageNet-1K 的任何方法。

我们也评估迁移学习的对象检测，实例分割，和语义分割。在这些任务中，我们的前训练比监督前训练获得了更好的结果，更重要的是，我们通过扩大模型观察到显著的收获。

这些观察结果与在NLP的自我监督前训练中所看到的结果一致。我们希望它们能使我们的领域探索类似的发展轨迹

3.Approach

我们的 MAE 是简单的 autoencoder 方法，通过一部分观察来重建原始 signal。像所有 autoencoder 一样，我们的方法有一个将观察到的 signal 映射成潜在表达的 encoder，和一个通过潜在表达，重建原始 signal 的 decoder。和传统 autoencoder 不同的是，我们用了非对称的设计，使encoder 只对观察到的部分 signal 进行操作，一个轻量级的 decoder ，通过潜在的表达和 mask tokens 重建完整 signal。

3.1.Masking

像 ViT 那样，我们将图像划分成了规则的，没有重叠的 patches。然后我们对一个子集的 patches 进行采样，并 mask 剩下的部分（删除）。我们的采样规则很简单：无替代得随机采样一些 patches，遵循均匀分布。我们称其为 random sampling

高 mask 比率的随机采样（删除 patches 的比例），大幅度消除了冗余，因此创造了一个不能通过可见邻居 patches 外推来简单解决的任务。均匀的分布防止了潜在的中心偏差(即在图像中心附近有更多的被屏蔽的补丁)。

最终，高度稀疏的输入为设计一个高效的编码器创造了机会

3.2.MAE encoder

我们的 encoder 是一个 ViT，但是只应用在可见的 patches 上。encoder 先通过 Linear Projection编码图片，然后加上位置编码，随后送入一系列 Transformer blocks 来处理。我们的 encoder 只使用了一小部分，masked patched 被丢弃，不使用 mask tokens。这允许我们只使用一小部分运算和内存，来训练一个很大的 encoder。

3.3.MAR decoder

decoder 的输入是整个 tokens 的集合，包含

• 编码好的可见的 patches
• mask tokens

每一个 mask token 都是共享的，学习的向量，表示这里存在一个有待预测的缺失 patch。我们将位置 embedding 添加到这个完整集合中的所有 tokens 中 ;如果不这样做， mask tokens 就没有他们在图像中位置的信息。

解码器还有另一系列Transformer块

decoder 只用在预训练过程中使用。因此，decoder 结构可以以独立于 encoder 设计的方式灵活设计。我们用很小的 decoder进行实验，比 encoder 更小，更窄。例如，我们默认的 decoder 对每个 token 的计算量小于 encoder 的 10%。

在这种非对称设计中，完整的 token 集合只由轻量级 decoder 处理（encoder 只处理可见的 tokens），这大大减少了预训练时间。

3.4.重建目标

我们的 MAE 通过预测为 masked patch 预测每个像素值来重建输入，每个元素在 decoder 中的输出都是表达一个 patch 中像素值的向量。decoder 的最后一层是一个 linear projection，其输出的通道数等于一个 patch 中的像素值的个数。decoder 的输出被 reshape 来组成重建的图。

我们的 loss function 在像素空间上计算了重建图像和原图像的 MSE（平方差）。我们值计算了 masked patches 的 loss，就像 BEiT 那样。

我们还研究了一个变种，他的重建目标是每个 masked patch 的归一化像素值。具体来说，我们计算一个patch中所有像素的均值和标准差，然后用它们来归一化这个patch。在实验中，使用归一化像素作为重建目标提高了表示质量。

3.5.简单的实施

我们的MAE预训练可以高效地实施，而且重要的是，不需要任何专门的稀疏操作。

首先，我们为每个input patch生成一个 token（先使用 Linear Projection，然后加上位置 embedding）。

接下来，我们将随机洗牌 tokens list，并根据屏蔽比率删除 list 的最后一部分。这个过程为encoder产生了一个tokens 的小的集合，和无替换地采样 patches 是一样的。

在 encoding 之后，我们向 encoded patch list 中添加一个 mask tokens list，然后反洗牌整个list（和前面的洗牌操作是相反的）以将所有的tokens和tragets对齐

decoder 在整个 list 上应用（带有添加的位置 embedding）。

如前所述，不需要进行稀疏操作。这个简单的实现引入的开销可以忽略不计，因为洗牌和反洗牌操作非常快。

4.ImageNet Experiments

我们在 ImageNet-1K训练集上做了自监督预训练。然后我们通过 1.端到端 fine-tuning 2.linear probing（冻结 encoder，只加一个线性分类器来训练） 来评价 representation 效果。我们使用的是对于单个 224 x 224 的裁剪图片的 top-1 validation accuracy 。

使用了 ViT-Large 作为 backbone，并提供了一些实验结果

可以看到，经过 MAE 预训练后 finutune 的效果要超过直接从头训练。这里 fine-tuning使用 50 epochs（从头训练使用 200 epochs），表明 fine-tuning 的准确度主要依赖于预训练

5.Main Properties

5.1.Masking ratio

75% ration 在 linear probing 和 fine-tuning上都很好。这与 BERT 是相反的，它的典型的 masking ratio 是 15%

上图也说明了，linear probing 和 fine-tuning 的趋势是不同的。

• 对于 linear probing 来说，准确度随着 masking ratio稳定上升，直到到到达最佳点：准确率差了20%。

• 对于 fine-tuning来说，结果对于 ration 不是那么敏感，在一个比较大范围的 masking ratio 内效果都不错。Fine-tuning 的所有结果都比从头开始训练要好。

• 上图表示，使用 75% masking ratio进行预训练，但是应用在更高的 masking ratio输入上

该模型推断出缺失的 patch 会产生不同的、但貌似合理的输出。它搞清楚对象和场景的形态意义，这不能简单地通过延伸线条或纹理来完成。我们假设这种类似推理的行为与学习有用的表征有关。其实就是说 MAE 学到的 representation 是有用的表征。

5.2.Decoder design

（a） 代表 decoder 的深度（Transformer blocks个数）对 fine-tuning 和 linear probing 的影响

一个足够深的 decoder 对于 linear probing来说是很重要的。

• 这可以用像素重建任务和识别任务之间的差距来解释: autoencoder 的最后几层更专门用于重建，与识别不太相关。如果decoder比较深，那么decoder就有足够的容量学习到重建能力，这样encoder可以更专注于提取特征。
• 在 linear probing 中，可以提升 8% 的准确度。

但是，如果使用的是 fine-tuning，encoder 的最后几层会被调节得适应识别任务。decoder 的深度对于提升 fine-tuning影响不那么大。

有趣的是，使用一个单层的 decoder，之后 fine-tuning 就能达到 84.8%的准确度，这么小的 decoder能够进一步加速训练。

（b）代表 decoder 的宽度（通道数）

默认情况下，我们使用 521 维，这在两种方法下都能表现得很好。更窄的 decoder 在 fine-tuning 下也能工作得很好。

综上，我们默认的 MAE decoder 是轻量的，有 8 blocks 和 512-d 宽度。它只有 9% FLOPS(每秒浮点运算次数)，相较于 ViT-L（24 blocks，1024-d）。因此，尽管 decoder 处理了所有的 tokens，他在总的计算量中，仍然是很小的一部分

5.3.Mask token

MAE 一个重要的设计就是，在 encoder 中跳过了 mask token，在后面的轻量级 decoder 中又加入了。

注：w/o 是 without 的意思，w/ 是 with 的意思

如果encoder也处理 mask tokens，在 linear probing 上效果会降低 14%，如上表。

这种情况下，在预训练和应用时会有差别。预训练时，输入会含有大比例的 mask tokens，而这些在未损坏的图像中不会出现。通过移除make tokens，我们限制了 encoder 只能看到真的 patches，以此提升准确度。

移除 mask tokens 还能训练时的运算，另外内存使用也大大减少，这可以支持训练更大的模型，或者通过更大的 batch 来训练，以加速。

5.4.Reconstruction target

我们比较了不同的重建目标。目前为止，我们的结果是基于没有（每个 patch）做标准化的像素。使用标准化的像素能够提高准确度。这种 pre-patch 的标准化增强了局部的对比度。

在另一个变种中，我们用 PCA 处理了 patch，并用最大的 PCA 系数作为目标。这样做降低准确度。

两个实验都说明了，在我们的方法中，高频成分（pixel）是有用的

我们还尝试了预测 token 的变种（像 BEiT 那样），使用 DALLE pre-trained dVAE 当做 tokenizer。这里，MAE decoder 使用交叉熵损失来预测 token indices。这种方法在 fine-tuning 上比非标准化 pixels 提高了 0.4%，但是不如 标准化 pixels 好。而且，降低了 linear probing 的准确度。这样还要预训练一个 dVAE tokenizer。使用 pixels 就不会有这些问题

直接使用 pixels 不仅效果好，而且省事，计算量小。

5.5.Data augmentation

MAE 在只进行裁剪的增强下，无论是固定尺寸，还是随机尺寸（都有随机的水平翻转）抖工作得很好。增加 color jit 会降低效果。令人惊讶的是，我们的MAE表现得很好，即使没有数据增强（只有中心裁剪，没有翻转）。相比之下，对比学习往往很依赖于数据增强。

在 MAE 中，random masking 发挥了数据增强的作用。在每次迭代中，mask 都是不同的，因此不管数据是否增强，都会产生新的训练样本。

5.6.Mask sampling strategy

MAE 在基于 block 的 masking（50% ratio）时，工作效果时合理的，但是在 75% ratio 时就削弱了。这比随机采样更加困难，因为我们观察到了更高的 training loss。并且，重建是更模糊的。

我们也尝试了基于 grid 的采样，每四个 patches 保留一个。这是更简单的任务，training loss更低了，重建更加清晰。但是，representation 的质量更差了。

MAE 上，随机采样效果最好，它允许更大比率的 masking，这在保持准确度的同时，获得了比较好的加速效果。

5.7.Training schedule

至此，我们前面的实验都是基于 800 epoch 的预训练的。上图显示了 epoch 的影响。准确度随着更长的训练，稳步提升。实际上，我们在 1600 epochs 没有观察到 linear probing 准确度的饱和。

这与对比学习方法是很不同的，比如 MoCo v3 在 300 epochs 就饱和了。

注意到，MAE encoder 每个 epoch 只能看到 25% patches，在对比学习中，encoder 每个 patch 可以看到 200%（两次裁剪），或者更多 patches。

6.与以前结果的对比

6.1.与自监督方法对比

对于 ViT-B，所有方法都比较接近。对于 ViT-L，方法之间的差距就比较大了。这表明，对于更大的模型，一个挑战是减少过拟合。

MAE可以轻易的扩展到更大模型并具有稳定的性能提升。比如：ViT-H取得了86.9%的精度，在448尺寸微调后，性能达到了87.8% ，超越了此前VOLO的最佳87.1%(尺寸为512)。注：该结果仅使用了ViT，更优秀的网络表达可能会更好

更 BEiT相比，我们的 MAE 更简单，更快并且更准确。我们的方法是重建 pixels，而BEiT 是预测 tokens：在使用 ViT-B 重建 pixels 时，会有 1.8% 的衰退。而且我们不用预训练 tokenizer。

6.2.与有监督预训练相比

可以看到，MAE 的效果只比 在 JFT300M 上悬链的 ViT 差一些，可能是训练数据量不够的原因。同时可以看出，MAE 的 pre-training，只使用 IN1K，有很好的泛化性：对于大容量的模型，从头开始训练的收益更大。这表明 MAE 有助于增大模型的 size。

7.Partial Fine-tuning

可以看出，linear probing 和 fine-tuning 的结果是不相关的。linear probing近年来很流行，但是他忽略了很有用但是非线性的特征，这些特征是能够增强深度学习的。两种方法折中，我们fine-tune后面的一些层，冻结其他层。

只 fine-tuning 一个 Transformer block 将 acc 从 73.5% 提升到 81.0%。如果我们只 fine-tuning 最后一个 block 的一半（transformer 的 MLP 子块），可以得到 79.1% 的 acc，这比 linear probing好得多。fine-tuning 一部分块能得到接近全 fine-tuning的结果。

我们也比较了 MoCo v3，他有更高的更高的 linear probing acc，但是其所有的 部分fine-tuning的结果都不如 MAE。当fine-tuning 4个 blocks时，差距达到 2.6%。

8.Transfer Learning Experiments

同时，论文也对比了MAE训练的encoder在下游任务（检测和分割）的迁移能力，同等条件下，MAE均能超过有监督训练或者其它无监督训练方法

一些需要进一步理解的点：

• linear probing
• pretext task
• 对比学习