MAE——Masked Autoencoders

Posted on 2023-01-09 Edited on 2022-01-09 In CV

MAE：Masked Autoencoders

code: https://github.com/facebookresearch/mae

1、概述

以一定比例随机mask掉图片中的一些图像块然后重建这些部分像素值。

主要特点有两个：

非对称的编、解码器设计
使用较高(如 75%)的掩码率(mask 比例)

第1点所述的“非对称”主要体现在 输入形式 与 网络结构 上：编码器(Encoder)仅对可见(un-masked)的图像块进行编码，而解码器(Decoder)的输入则是所有的图像块；同时，Decoder 可以是比较轻量的(比如 Encoder 通常是多层堆叠的 Transformer，而 Decoder 仅需较少层甚至1层就 ok)。这也表明 Encoder 与 Decoder 之间是解耦的。

第2点是该工作的一个重要发现：不同于 NLP，在 CV 中可能要配合较高的 mask 比例才能作为“有效”的自监督代理任务。“有效”指的是任务本身足够困难，这样模型才能学到有效的潜在特征表示。

由于 Encoder 仅处理 un-masked 的 patch(占所有输入的少数)，因此，尽管其本身网络结构比较重载，但依然能够高效训练，特别是对于大模型，能够加速3倍以上，同时配合较高的掩码率，还能够涨点。

作者从一个问题出来解释选择这样mask策略以及模型设计的理由，“为什么 masked autoencoding 在CV中应用比较少相较于NLP？” 作者提炼了以下三点：

架构差异：之前CNN的架构不适合使用
信息密度不同: 语言的信息密度比较高，将句子中的少量词语抹去再让模型预测被抹去的这些词是比较困难的任务；但对于图像则相反，它在空间中是冗余的，对于图片中的某个部分，模型很容易由其相邻的图像块推断出来，因此在CV中，mask的比例应该更高，才能使任务本身具有足够的挑战性，从而使模型学到良好的潜在特征表示。
解码的目标不一致：CV 和 NLP 在解码器的设计上应该有不一样的考虑：NLP 解码输出的是对应被 mask 掉的词语，本身包含了丰富的语义信息；而 CV 要重建的是被 mask 掉的图像块(像素值)，是低语义的。因此，NLP 的解码器可以很简单，比如 BERT，严格来说它并没有解码器，最后用 MLP 也可以搞定。因为来自编码器的特征也是高度语义的，与需要解码的目标之间的 gap 较小；而 CV 的解码器设计则需要“谨慎”考虑了，因为它要将来自编码器的高级语义特征解码至低级语义层级。

2、具体方法

MAE的特点主要是：高掩码率的随机 mask 策略、非对称的编、解码器设计以及重建的目标是像素值

2.1 masked 策略

沿袭 ViT 的做法，将图像分成一块块(ViT 中是 16x16 大小)不重叠的 patch，然后使用服从均匀分布(uniform distribution)的采样策略对这些 patches 随机采样一部分，同时 mask 掉余下的另一部分。被 mask 掉的 patches 占所有 patches 的大部分(实验效果发现最好的比例是 75%)，它们不会输入到 Encoder。

作者实验发现：无论是finetune，还是 fine-tune 还是 linear-probe(微调方法，将最后一层替换成线性层，微调时冻结其他层，只训练这个线性层)，75%都是一个比较好的比例。

2.2 Encoder

Encoder只处理 un-masked 的patches。Encoder可以是Vit或者其他backbone。图像划分成patch采用Vit的做法。

先将图像从 (B,C,H,W) reshape 成 (B,N,PxPxC)，其中 N 和 P 分别为 patch 数量和 patch 大小( $H/P \times W/P$)，也就是将3通道的图像转换成 N 个维度大小为 PxPxC 的向量；然后，通过线性映射(linear projection，可以是全连接层)将其嵌入(embed)到指定的维度空间大小，记为 ‘dim’(从 PxPxC project 到 dim)，转换成为 token(B,N,dim)；最后再加上位置嵌入(position embedding)，从而为各个 patch 添加位置信息。位置嵌入是所有图像共享的、可学习的，shape 与每张图的 token 相对应，即：(N,dim)。

由于 un-masked patches 占所有 patches 的少数，计算消耗和空间需求都减少了，因此可以训练很大的 Encoder。

2.3 Decoder

Decoder 不仅需要处理经过 Encoder 编码的 un-masked 的 tokens，还需要处理 masked tokens。但请注意，masked token 并非由之前 mask 掉的 patch 经过 embedding 转换而来，而是可学习的、所有 masked patches 都共享的1个向量，对，仅仅就是1个！

通过 position embedding 来区分各个 masked patch 对应的 token。

2.4 loss

MAE 预训练任务的目标是重建像素值，并且仅仅是 masked patches 的像素值，也就是仅对 mask 掉的部分计算 loss，而 loss 就是很大众的 MSE。为何仅计算 mask 部分的 loss？实验结果发现这样做模型的性能会更好，而如果对所有 patches 都计算 loss 的话会掉点。