【源头活水】CV预训练MAE（Masked AutoEncoder）

“问渠那得清如许，为有源头活水来”，通过前沿领域知识的学习，从其他研究领域得到启发，对研究问题的本质有更清晰的认识和理解，是自我提高的不竭源泉。为此，我们特别精选论文阅读笔记，开辟“源头活水”专栏，帮助你广泛而深入的阅读科研文献，敬请关注。

来源：知乎—大师兄

地址：https://zhuanlan.zhihu.com/p/448407149

最近何恺明大神的残差网络突破了10万的的引用量，而他的MAE又引起了业内极大的关注，这里将带来MAE的详细解读。继2017年Transformer[2]被提出之后，预训练成为了NLP领域主流的研究方向，经典算法有GPT，BERT[3]等。GPT是经典的自回归（Auto-Regressive）预训练模型，而BERT是经典的掩码模型，或者说是去噪自编码（Denosing AutoEncoder）语言模型。DeepMind的image GPT（iGPT）将预训练任务迁移到了计算机视觉方向，它的预训练任务借鉴了GPT系列，即通过自回归（Auto-Regressive）的方式根据保留的图像上半部分逐像素的预测图像的下半部分的方式，并且通过微调和线性探测两个方法验证了iGPT在图像分类任务上可以起到明显的效果提升的作用。iGPT表明自回归训练任务是可以迁移计算机视觉方向的，那么自编码语言模型是否也能迁移到计算机视觉上呢？这里要介绍的Masked AutoEncoder（MAE）给出了肯定的答案。MAE的最核心的思想是通过对图片中的patch进行随机掩码，然后通过未被掩码的区域预测被掩码掉的区域，进行使得模型学习图像的语义特征。

01

Motivation

虽然预训练在NLP上正发展的如火如荼，但是在计算机视觉方向却鲜有文章，究其原因，论文中给出了三个重要的点。

1. 模型架构不同：在过去的几十年，计算机视觉被卷积神经网络所垄断着，卷积是一个基于划窗的算法，它和其它嵌入（位置嵌入等）的融合比较困难，直到Transformer的提出才解决了这个问题。

2. 信息密度不同：文本数据是经过人类高度抽象之后的一种信号，它的信息是密集的，所以仅仅预测文本中的几个被掩码掉的单词就能很好的捕捉文本的语义特征。而图像数据是一个信息密度非常小的矩阵，其中包含着大量的冗余信息，而且像素和它周围的像素仅仅在纹理上就有非常大的相似性，恢复被掩码的像素并不需要太多的语义信息。

3. 解码器的作用不同：在BERT的掩码语言模型任务中，预测被掩码掉的单词是需要解码器了解文本的语义信息的。但是在计算机视觉的掩码预测任务中，预测被掩码的像素往往对图像的语义信息依赖的并不严重。

基于这三个动机，作者设计了基于掩码自编码器（Masked AutoEncoder，MAE）的图像预训练任务。MAE的先对图像的Patch进行掩码，然后通过模型还原这些掩码，从事实现模型的预训练。MAE的核心是通过75%的高掩码率来对图像添加噪音，这样图像便很难通过周围的像素来对被掩码的像素进行重建，迫使编码器去学习图像中的语义信息。

02

掩码机制

在ViT中，输入图像被分成若干个互不覆盖的Patch，然后每个Patch使用Transformer独立的计算embedding。在MAE中，作者也是以Patch为单位对图像进行掩码的。MAE的掩码的生成方式是先对patch进行shuffle，然后根据掩码率来对这个shuffle的patch列表进行掩码。对于掩码率，作者在微调和线性探测上实验了若干组值，实验结果表明75%的掩码率在微调和线性探测上都表现的比较突出，如图1。掩码的实现方式如下：

class RandomMaskingGenerator:
    def __init__(self, input_size, mask_ratio):
        if not isinstance(input_size, tuple):
            input_size = (input_size,) * 2
        self.height, self.width = input_size
        self.num_patches = self.height * self.width
        self.num_mask = int(mask_ratio * self.num_patches)

    def __repr__(self):
        repr_str = "Maks: total patches {}, mask patches {}".format(
            self.num_patches, self.num_mask
        )
        return repr_str

    def __call__(self):
        mask = np.hstack([
            np.zeros(self.num_patches - self.num_mask),
            np.ones(self.num_mask),
        ])
        np.random.shuffle(mask)
        return mask # [196]

对于掩码方式，作者实验了随机掩码，掩码掉一整块（BEit[6]的策略）和Grid掩码三种方式，从图2的实验结果来看，随机掩码的效果最稳定。

图2：MAE在三种掩码策略上的模型效果，左侧是不同的掩码策略，右侧是还原的效果

03

模型介绍

MAE的网络结构如图3所示，它是一个非对称的Encoder-Decoder架构的模型，Encoder架构是采用了ViT[5]提出的以Transformer为基础的骨干网络，它的基于Patch的输入正好可以拿来作为掩码的基本单元。MAE的Decoder是一个轻量级的结构，它在深度和宽度上都比Encoder小很多。MAE的另一个非对称的表现在Encoder仅将未被掩码的部分作为输入，而Decoder将整个图像的Patch（掩码标志和Encoder编码后的未被掩码patch的图像特征）作为输入，下面我们详细介绍这两个模块的细节。

图3：MAE的网络结构

3.1 Encoder

MAE的编码器结构借鉴了ViT，如图4所示。它首先将图片转换成一系列不重合的Patch，假设Patch的大小为    ，那么我们可以使用    个大小为    ，步长为    的卷积核来将图像转换成长度为    的特征向量，这便实现了图像Patch的线性映射（Linear Projection）。除了图像编码外，ViT中还引入了1-D位置编码来反应每个Patch在图像中的实际位置，因为2-D的位置编码并没有带来明显的效果提升。有了输入之后，ViT的骨干网络便是一系列Transformer堆叠而成的结构，最后接一个softmax激活函数作为输出层。

图4：ViT的算法流程

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

MAE的Encoder的输入不是整张图，而是只输入了未被掩码的部分，此外MAE的Encoder会为每一个未被掩码Patch计算一个特征向量。因为我们将输入减少了75%，因此训练速度也提高了75%，或者我们也可以在保持相同训练速度的同时引入更多的参数来获得更好的表现效果。作者也尝试了再Encoder中加入掩码图，实验结果表明这个方式不仅更慢了，而且在线性探测上准确率下降了14%。

最终，Encoder使用了ViT-L和ViT-H两个不同尺度的Encoder，它的参数细节如表1。

表1：ViT的参数

3.2 Decoder

MAE的输入包含两部分：一个是图像嵌入，它由Encoder编码之后的特征和被掩码的标志特征组合而成；另一个是整个图像的1-D位置编码。Decoder的被掩码的标志是一个共享的且可以学习的模块。

在MAE中，Decoder是独立于Encoder的模块，它只在图像重建的时候使用。MAE的Decoder是一个轻量级的模块，作者在它的实验中尝试了不同的解码器的宽度和深度，最终选择了8层，节点数为512的Transformer。实验结果表明这是一组效果最好的参数，无论增加还是减少参数量，模型的效果都会变差，如图5。

图5：MAE的Decoder的参数量的对照实验

从图5中我们还可以看出Decoder的参数量对模型的最终效果影响并不大，即使很小的参数量也和最后的模型表现差距不大。这表明了MAE的Encoder已经提取到了足够Decoder还原图像所需要的语义特征。

3.3 其它优化

重构目标：作者尝试了重构原始图像和重构归一化之后的图像，实验结果表明重构归一化的图像可以带来0.5%的准确率的提升。

数据增强：MAE只使用随机裁剪和水平翻转的数据扩充策略。实验结果表明颜色变化（Color Jitter）相关的增强会带来准确率的下降。因为MAE的随机Mask已经引入了很大的随机性，过度的数据增强可能会带来模型的训练难度的增加。

04

总结

这是一篇经典的何恺明风格的文章，简单且有效。它告诉了我们通过图像重建进行图像预训练的可行性。MAE能够简单且有效的一个重要原因是恺明大神对视觉和文本的深刻洞察能力，他在论文中指出了文本和图像在信息密度上的不同以及BERT的MLM用在图像上的问题，进而提出了解决这些问题的方案。MAE这篇文章也算是开了使用重构图像进行图像预训练的新坑，相信不久就会有更多的基于MAE的文章发表出来，不知道有朝一日MAE的引用量也会像残差网络那样突破10万。

Reference

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." arXiv preprint arXiv:2111.06377 (2021).

[2] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017.

[3] Devlin, Jacob, et al. "Bert: Pre-training of deep bidirectional transformers for language understanding." arXiv preprint arXiv:1810.04805 (2018).

[4] Chen, Mark, et al. "Generative pretraining from pixels." International Conference on Machine Learning. PMLR, 2020.

[5] Dosovitskiy, Alexey, et al. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

[6] Bao, Hangbo, Li Dong, and Furu Wei. "BEiT: BERT Pre-Training of Image Transformers." arXiv preprint arXiv:2106.08254 (2021).

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“源头活水”历史文章

更多源头活水专栏文章，

请点击文章底部“阅读原文”查看

分享、在看，给个三连击呗！