论文阅读|MAE

研究人员开发了一个非对称编码器 - 解码器架构，其中一个编码器只对可见的 patch 子集进行操作 (即没有被 mask 掉的 token)，另一个简单解码器可以从潜在表征和被 masked 掉的 token 重建原始图像。Decoder 的架构可以是十分轻量化的模型，且具体的架构对模型性能影响很大。研究人员进一步发现，Mask 掉大部分输入图像 (例如 75%) 会产生重要且有意义的自监督任务。结合这两种设计，我们就能高效地训练大型模型：提升训练速度至 3 倍或更多，并提高准确性。

我们的可扩展方法允许学习泛化良好的大容量模型：例如，在仅使用 ImageNet-1K 数据的方法中，vanilla ViT-Huge 模型实现了最佳精度 (87.8%)。下游任务中的迁移性能优于有监督的预训练，并显示出有希望的扩展行为。

Introduction

深度学习见证了能力和容量不断增长的体系结构的爆炸性增长[28、24、47]。在硬件快速发展的帮助下，今天的模型可以很容易地容纳100万张图像[13]，并开始需要数亿张--通常是公开无法访问的--标签的图像[16]。

通过自我监督的预训练，自然语言处理 (NLP) 已成功解决了这种对数据的需求。基于 GPT [40, 41, 4] 中的自回归语言建模和 BERT [14] 中的掩码自动编码的解决方案在概念上很简单：它们删除一部分数据并学习预测删除的内容。这些方法现在可以训练包含超过一千亿个参数的可泛化 NLP 模型 [4]。

MAE 方法严格来讲属于一种去噪自编码器 (Denoising Auto-Encoders (DAE))，去噪自动编码器是一类自动编码器，它破坏输入信号，并学会重构原始的、未被破坏的信号。MAE 的 Encoder 和 Decoder 结构不同，是非对称式的。Encoder 将输入编码为 latent representation，而 Decoder 将从 latent representation 重建原始信号。

Masked autoencoders的想法，一种更通用的去噪自动编码器 [48] 的形式，在计算机视觉中也很自然且适用。事实上，与视觉密切相关的研究 [49, 39] 早于 BERT。然而，尽管随着 BERT 的成功对这一想法产生了浓厚的兴趣，但视觉自动编码方法的进展仍落后于 NLP。我们问：是什么让masked autoencoding在视觉和NLP之间有所不同？我们试图从以下几个方面来回答这个问题：

CV 和 NLP 主流架构不同：直到 ViT (2020.12) 出现之前，CV 的主流架构一直是以卷积网络为主，NLP 的主流架构一直是以 Transformer 为主。卷积核作用在一个个的 grid 上面，直观来讲没法产生像 Transformer 一样的 token 的概念，也就是说如果我们只使用卷积网络，那么 image token 概念的建立就不那么直观。所以，像 Transformer 那样在 token 的基础上进行自监督学习就不太适用，这是第一个难点。然而，随着Vision Transformers(VIT)[16]的引入，这一架构差距已经得到解决，不应再构成障碍。
语言和图片 (视频) 的信息密度不同：语言是人类造就的信号，它 highly semantic ， information-dense。而图片 (视频) 是自然产生的信号，它 heavy spatial redundancy。即挡住图片的一部分 patches，可以很容易地通过看它周围的 patches 而想象出它的样子来。所以，语言和图像，一个信息密度高，一个信息密度低，这是第二个难点。解决的办法是什么呢？作者提出了一个简单的策略：即挡住图片的 patches 的比例高一些。比如之前你挡住一张图片的 30% 的 patches，能够轻松通过周围的 patches 预测出来；那现在如果挡住图片的 90% 的 patches，还能够轻松通过周围的 patches 预测出来吗？
AutoEncoder 里面的 Decoder 部分 (就是将映射得到的中间特征重建为 input 的模块) 在 CV 和 NLP 中充当的角色不同：在 CV 领域，Decoder 的作用是重建 image pixels，所以 Decoder 的输出语义级别很低。在 NLP 领域，Decoder 的作用是重建 sentence words ，所以 Decoder 的输出语义级别很丰富。

在此分析的推动下，我们提出了一种简单、有效且可扩展的Masked AutoEncoder (MAE) 形式，用于视觉表示学习。我们的 MAE 从输入图像中屏蔽随机patch，并在像素空间中重建缺失的补丁。它具有不对称的编码器解码器设计。我们的编码器只对可见的patch子集（没有掩码标记）进行操作，我们的解码器是轻量级的，可以通过潜在表示和masked tokens中重构输入（图 1）。在我们的非对称编码器-解码器中将masked tokens转移到小型解码器会导致计算量大大减少。在这种设计下，非常高的掩蔽率（例如，75%）可以实现双赢：它优化了准确性，同时允许编码器只处理一小部分（例如，25%）的patches。这可以将整体预训练时间减少 3 倍或更多，同时减少内存消耗，使我们能够轻松地将 MAE 扩展到大型模型。

我们的 MAE 学习了泛化能力非常高的模型。通过 MAE 预训练，我们可以在 ImageNet-1K 上训练数据密集型模型，如 ViT-Large/-Huge [16]，并提高泛化性能。使用 vanilla ViT-Huge 模型，我们在 ImageNet-1K 上微调时达到 87.8% 的准确率。这优于之前仅使用 ImageNet-1K 数据的所有结果。我们还评估了对象检测、实例分割和语义分割的迁移学习。在这些任务中，我们的预训练比有监督的预训练取得了更好的结果，更重要的是，我们通过扩大模型观察到了显着的收益。这些观察结果与 NLP [14、40、41、4] 中自我监督预训练中的观察结果一致。

Method

所提MAE是一种非常简单的自编码器方案：基于给定部分观测信息对原始信号进行重建 。类似于其他自编码器，所提MAE包含一个将观测信号映射为隐式表达的编码器，一个用于将隐式表达重建为原始信号的解码器。与经典自编码器不同之处在于：我们采用了非对称设计，这使得编码器仅依赖于部分观测信息(无需掩码token信息)，而轻量解码器则接与所得隐式表达与掩码token进行原始信号重建(可参见下图)。

*图1：MAE结构。在预训练期间，遮蔽图像patch的大的随机子集(例如，75%)。编码器被应用于可见patch的小子集。在编码器之后引入masked token，全套编码的patch和masked token由一个小解码器处理，该解码器以像素重建原始图像。在预训练之后，丢弃解码器，并将编码器应用于未被破坏的图像以产生用于识别任务的表示。

Masking

参考ViT，我们将输入图像拆分为非重叠块，然后采样一部分块并移除其余块(即Mask)。我们的采样策略非常简单：服从均匀分布的无重复随机采样 。我们将该采样策略称之为“随机采样”。具有高掩码比例的随机采样可以极大程度消除冗余，进而构建一个不会轻易的被近邻块推理解决的任务 (可参考下面图2-4)。而均匀分布则避免了潜在的中心偏置问题。最后，高度稀疏的输入为设计高效编码器创造了机会，接下来介绍。

MAE encoder

MAE中的编码器是一种ViT，但仅作用于可见的未被Mask的patches。类似于标准ViT，该编码器通过线性投影于位置嵌入对块进行编码，然后通过一系列Transformer block进行处理。然而，由于该编解码仅在较小子集块(比如25%)进行处理，且未用到掩码Token信息。这就使得我们可以用很少的计算量和内存来训练一个非常大的编码器 。

MAE decoder

MAE 解码器的输入是由 (i) 编码器的输出和 (ii) 掩码token组成的完整令牌集。参见图 1。每个掩码标记是一个共享的可学习向量，表示存在要预测的缺失patch。我们将位置嵌入添加到这个完整集合中的所有标记；如果没有这个，掩码token将没有任何关于它们在图像中的位置的信息。解码器有另一个系列的 Transformer block。

MAE 解码器仅在预训练期间用于执行图像重建任务（仅编码器用于生成用于识别的图像表示）。因此，解码器架构可以以独立于编码器设计的方式灵活设计。我们尝试了非常小的解码器，比编码器更窄更浅。例如，我们的默认解码器与编码器相比，每个token的计算量 <10%。采用这种非对称设计，全套token仅由轻量级解码器处理，大大减少了预训练时间。

Reconstruction target

该MAE通过预测每个masked token的像素值进行原始信息重建。解码器输出中的每个元素都是代表patch的像素值向量。解码器的最后一层为线性投影，其输出通道数等于每个patch的像素数量。解码器的输出将通过reshape构建重建图像。损失函数计算像素空间中重建图像和原始图像之间的均方误差 (MSE)，注：类似于BERT仅在masked patches计算损失。

我们同时还研究了一个变种：其重建目标为每个掩码块的归一化像素值 。具体来说，我们计算每个patch的均值与标准差并用于对该块进行归一化，最后采用归一化的像素作为重建目标提升表达能力。

Simple implementation

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

首先，我们为每个输入patch生成一个token（通过添加位置嵌入的线性投影）。

接下来，我们根据掩码率随机打乱(random shuffle）token序列并根据掩码比例移除最后一部分token。此过程为编码器生成一小部分token，相当于在不替换的情况下对patch进行采样。

完成编码后，我们在编码块中插入掩码token并反置换(unshuffle)得到全序列token以便于与target进行对齐；

最后，将解码器作用于上述全序列token（添加了位置嵌入）。

如前所述，不需要稀疏操作。此外，shuffle与unshuffle操作非常快，引入的计算量可以忽略。

ImageNet Experiments

我们在ImageNet-1K上进行自监督的预训练。然后我们采用有监督的训练去评估表征，通过两种方式：端到端的微调，允许改整个模型所有的可学习的参数；线性探测,允许改最后一层的线性输出层。我们汇报了在224×224图像上的Top1的准确率。具体细节如下：

Baseline: ViT-Large

我们使用 ViT-Large (ViT-L/16) [16] 作为我们消融研究的主干。 ViT-L 非常大（比 ResNet-50 [24] 大一个数量级）并且容易过拟合。以下是从头开始训练的 ViT-L 与从我们的baseline MAE 微调的对比：

scratch, original: 76.5, ViT 所有的内容在 ImageNet-1K上训练,结果不稳定，200epoches

scratch, our impl: 82.5 加入 strong regularization A.2 ViT 文章说 ViT需要很大的数据才行小一点的数据 + 合适的正则化 ✔
baseline MAE: 84.9 先使用 MAE 做预训练，再在 ImageNet 上微调 50 epoches 数据集没变化，预训练和微调都是 ImageNet MAE 纯从图片上学到不错的信息

Main Properties

Decoder design.

(a) 解码器的深度，即多少个 Transformer 块
ft（fine-tune所有可学习的权重都一起调）：效果和深度关系不大 84.x，效果更好（但耗费更贵）
lin（linear-probe只调最后一个线性层）：深度深一点好

(b) 解码器的宽度，每一个 token 表示成一个多长的向量
512 最佳

总体而言，我们默认的MAE解码器是轻量级的。它有8个块，宽度为512-d(表1中为灰色)。与VIT-L相比，它的每个token只有9%的FLOPS(24个数据块，1024-d)。因此，尽管解码器处理所有token，但它仍然只占整个计算的一小部分。

Mask token.

Reconstruction target.

（d）构的目标：minimize MSE(pixel)
w/o + normalization(均值为0 方差为 1)：效果和w/差不多，但计算量更小，更简单
w/ + normalization(均值为0 方差为 1) ：引入块归一化可以进一步提升模型精度
PCA （做一次降维）
dVAE（BEiT 的做法，通过 ViT 把每一个块映射到一个离散的 token，像 BERT 一样的去做预测）

Data augmentation.

(e）固定大小的裁剪（crop）
随机大小的裁剪（crop）:效果最好
裁剪 + 颜色变化
结论：MAE 对数据增强不敏感

Mask sampling strategy.

(f)随机采样（采样率75 %）：最简单最好
块采样（采样率50 %）
块采样（采样率75 %）
网格采样（采样率75 %）

Masking ratio

掩码比例为75%有益于两种监督训练方式(端到端微调与linear probing)。这与BERT的行为截然相反，其掩码比例为15%。

与此同时，从上图可以看到：端到端微调与linear probing两种方式存在不同的趋势：

对于linear probing而言，模型性能随掩码比例非常稳定的提升直到达到最高点，精度差约为20%；
对于微调来说，模型性能再很大范围的掩码比例内均极度不敏感，而且所有微调结果均优于linear probing方式。

Training schedule.

下图给出了不同训练机制的性能对比(此时采用了800epoch预训练)，可以看到：更长的训练可以带来更定的精度提升 。作者还提到：哪怕1600epoch训练也并未发现linear probing方式的性能饱和。这与MoCoV3中的300epoch训练饱和截然相反 ：在每个epoch，MAE仅能看到25%的图像块；而MoCoV3则可以看到200%，甚至更多的图像块。

Comparisons with Previous Results

上表给出了所提MAE与其他自监督方案的性能对比，从中可以看到：

对于ViT-B来说，不同方案的性能非常接近；对于ViT-L来说，不同方案的性能差异则变大。这意味着：对更大模型降低过拟合更具挑战性 。
MAE可以轻易的扩展到更大模型并具有稳定的性能提升。比如：ViT-H取得了86.9%的精度，在448尺寸微调后，性能达到了87.8% ，超越了此前VOLO的最佳87.1%(尺寸为512)。注：该结果仅使用了ViT，更优秀的网络表达可能会更好。

和BEiT对比，我们MAE更加准确，同时更简单更快。我们方法是重构像素，对比BEiT是重构令牌tokens：如果BEiT使用ViT-B重构像素，将会有1.8%性能下降。我们不需要dVAE预训练，同时比BEiT快3.5倍之多。
MAE模型预训练了1600个epoch达到了更好的效果。即使如此，我们整体的预训练时间也比其他模型要少，在相同的硬件条件下。例如，当训练ViT-L时，我们采用128个TPU-v3核处理，1600个epoch需要31个小时；而MoCo v3进行300个epoch就需要36h。

MAE预训练vs.有监督预训练：JFT效果最好，MAE次之。（可能因为JFT数据集包括的类数远远大于 ImageNet）

Partial Fine-tuning

最后，作者提出了一种 Partial Fine-tuning 的新套路，它不同于之前人们常用的 Linear Probing (只训练最后一层线性分类器的参数) 和 Fine-tuning (训练所有层的参数)。

Partial Fine-tuning 是指 只训练最后模型的若干层的参数。如下图12所示，值得注意的是，Fine-tuning 1个 Transformer Block 可以把 Accuracy 从73.0%提升至81.0%。此外，如果我们只Fine-tuning 1个 Transformer Block 的 MLP block，我们可以得到79.1% ，比 Linear Proing 要好得多。

fine-tune层数少，快，精度差
fine-tune层数多，慢，精度好

4-5层较好