【ICCV 2022】（MAE）Masked Autoencoders Are Scalable Vision Learners

何凯明一作文章：https://arxiv.org/abs/2111.06377

本文的出发点：是BERT的掩码自编码机制：移除一部分数据并对移除的内容进行学习。mask自编码源于CV但盛于NLP，恺明对此提出了疑问：是什么导致了掩码自编码在视觉与语言之间的差异？尝试从不同角度进行解释并由此引申出了本文的MAE。

恺明提出一种用于计算机视觉的可扩展自监督学习方案Masked AutoEncoders(MAE)。所提MAE极为简单：对输入图像的随机块进行mask并对遗失像素进行重建。它基于以下两个核心设计：

我们设计了一种非对称编解码架构，其中解码器仅作用于可见块(无需mask信息)，而解码器则通过隐表达与mask信息进行原始图像重建；
我们发现对输入图像进行高比例mask(比如75%被mask掉)可以产生一项重要且有意义的自监督任务。

上述两种设计促使我们可以更高效的训练大模型：我们加速训练达3x甚至更多，同时提升模型精度。所提方案使得所得高精度模型具有很好的泛化性能：仅需ImageNet-1K，ViT-Huge取得了87.8%的top1精度 。下游任务的迁移取得了优于监督训练的性能，证实了所提方案的可扩展能力。

整体架构：

图 1：我们的 MAE 架构。

1、在预训练期间，图像patches 的大量随机子集（例如 75%）被mask。

2、Encoder应用于可见patches的小子集。

3、将encoder后的所有patches 和masked tokens由小型decoder处理，该解码器以像素为单位重建原始图像。

4、预训练后，decoder被丢弃，encoder应用于未损坏的图像（完整的patches）以进行识别任务。

方法：

我们的MAE是一种简单的自动编码方法，可以在给定部分观察的情况下重建原始信号。与所有AE一样，我们的方法具有将观察到的信号映射到 latent representation 的encoder，以及从 latent representation 重建原始信号的decoder。与经典AE不同，我们采用非对称设计，允许编码器仅对部分观察到的信号（没有掩码标记）进行操作，并采用轻量级解码器，从 latent representation 和 mask tokens 中重建完整信号。图 1 说明了接下来介绍的想法。

Masking

我们的采样策略非常简单：服从均匀分布的无重复随机采样 。我们将该采样策略称之为“随机采样”。具有高掩码比例的随机采样可以极大程度消除冗余，进而构建一个不会轻易的被近邻块推理解决的任务。而均匀分布则避免了潜在的中心偏置问题。【提升自监督学习任务的难度】

MAE Encoder

编码器用于生成用于：encoder image + 识别的完整图像的representations

MAE中的编码器是一种ViT，但仅作用于未被Mask的块。类似于标准ViT，该编码器通过线性投影于位置嵌入对块进行编码，然后通过一系列Transformer模块进行处理。然而，由于该编解码仅在未mask的patches(比如25%)进行处理。这就使得我们可以用很小的计算资源训练一个非常大的编码器 。

MAE Decoder：

MAE解码器的输入包含：(1) encoded visible patches；(2) mask token。如Figure1所示，每个mask token共享的可学习向量，它用于指示待预测mask块。我们还会对所有token添加位置嵌入信息。解码器同样包含一系列Transformer模块。

MAE decoder 仅在预训练期间用于执行图像重建任务。因此，可以以独立于编码器设计的方式灵活地设计解码器架构。我们用非常小的解码器进行实验，比encoder更窄、更浅。例如，我们的默认decoder 与encoder 相比，每个token 的计算量 < 10%。通过这种非对称设计，全套 token 仅由轻量级解码器处理，这显着减少了预训练时间。

Reconstruction target

基本：我们的 MAE 通过预测每个mask patch的像素值来重建输入。解码器输出中的每个元素是vector of pixel values representing a patch。解码器的最后一层是线性投影，其输出通道的数量等于patch 中像素值的数量。decoder的输出从a reconstructed image被reshaped。我们的损失函数计算像素空间中重建图像和原始图像之间的均方误差（MSE）。我们仅在mask patches上计算损失，类似于 BERT [14]（计算所有像素的损失会导致精度略有下降）。

变体：我们还研究了一种变体，其重建目标是每个masked patch的归一化像素值。具体来说，我们计算补丁中所有像素的平均值和标准差，并使用它们来标准化该 patch。使用归一化像素作为重建目标可以提高我们实验中的表示质量。

Simple implementation

Our MAE pre-training can be implemented efficiently, and importantly, does not require any specialized sparse operations. First we generate a token for every input patch (by linear projection with an added positional embedding).

Next we randomly shuffle the list of tokens and remove the last portion of the list, based on the masking ratio.

This process produces a small subset of tokens for the encoder and is equivalent to sampling patches without replacement.

After encoding, we append a list of mask tokens to the list of encoded patches, and unshuffle this full list (inverting the random shuffle operation) to align all tokens with their targets. The decoder is applied to this full list (with positional embeddings added). As noted, no sparse operations are needed. This simple implementation introduces negligible overhead as the shuffling and unshuffling operations are fast.
我们的 MAE 预训练可以高效地实施，而且重要的是，不需要任何专门的稀疏操作（sparse operations）。

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

接下来，我们根据 masking ratio 随机打乱 token 列表并删除列表的最后部分。此过程为编码器生成一小部分tokens ，相当于对不带替换的patches进行采样。

编码后，我们将掩码令牌列表附加到编码补丁列表中，并取消对这个完整列表的洗牌（反转随机洗牌操作）以使所有令牌与其目标对齐。解码器应用于此完整列表（添加了位置嵌入）。如前所述，不需要稀疏操作。由于洗牌和解洗操作速度很快，这种简单的实现引入的开销可以忽略不计。