【MAE】一种高效的自监督训练器

MAE

Paper : Masked Autoencoders Are Scalable Vision Learners

Author : Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár , Ross Girshick

Affiliation : Facebook AI Research (FAIR)

Publication : CVPR-2022

1 前言

MAE—随机mask掉一部分patches并重建这部分丢失的像素，可看作是一个可伸缩的(scalable)自监督学习器，能极大加速模型的训练速度并提升准确率。 下游迁移性能优于有监督预训练，并有良好的的scaling（可伸缩性）。

1.1 动机

深度学习的发展需要越来越多的数据和越来越大的模型，而这种对数据的需求已通过自监督预训练在自然语言处理（NLP）中成功得到解决。这些解决方案基于 GPT中的自回归语言建模和 BERT 中的屏蔽自动编码，（概念上很简单：它们删除一部分数据并学习预测删除的内容。）

MAE的想法是一种更通用的去噪自动编码器，并且在视觉上相关的研究是在BERT之前，但视觉自动编码方法的进展却落后于 NLP。

于是何凯明大佬说：

视觉和语言之间的屏蔽自动编码有何不同？

(i) 直到最近，架构还是不同的。在视觉领域，卷积网络 在过去十年中占据主导地位。卷积通常在规则网格上运行，将掩码标记或位置嵌入等“指标”集成到卷积网络中并不简单。然而，这一架构差距已通过 Vision Transformers (ViT)的引入得到解决，并且不应再成为障碍。

(ii) 语言和视觉之间的信息密度不同。语言是人类生成的信号，具有高度语义和信息密度。当训练模型仅预测每个句子中的几个缺失单词时，此任务似乎会引发复杂的语言理解。相反，图像是具有大量空间冗余的自然信号，例如，可以从邻近的补丁中恢复丢失的补丁，而几乎没有高层次的对于部分，物体和场景的理解。为了克服这种差异并鼓励学习有用的特征，我们证明了一种在计算机视觉中效果很好的简单策略：屏蔽很大一部分的随机补丁。这种策略很大程度上减少了冗余，并创造了一项具有挑战性的自我监督任务，需要超越低级图像统计的整体理解。

（iii）自动编码器的解码器将潜在表示映射回输入，在重建文本和图像之间发挥着不同的作用。在视觉中，解码器重建像素，因此其输出的语义级别低于常见的识别任务。这与语言相反，在语言中，解码器预测包含丰富语义信息的缺失单词。虽然在 BERT 中，解码器可能很简单（MLP），但我们发现对于图像，解码器设计在确定学习的潜在表示的语义级别方面起着关键作用。

在这种分析的推动下，提出了一种简单、有效且可扩展的掩码自动编码器（MAE）形式，用于视觉表示学习。

2 方法

2.1实现

MAE是一种简单的自动编码方法，它根据原始信号的少部分观测值重建原始信号。**像所有的自动编码器一样，有一个编码器将观察到的信号映射到潜在表示，还有一个解码器从潜在表示重构原始信号。与经典的自动编码器不同，采用了一种非对称设计，允许编码器仅对部分观察信号（无掩码标记）进行操作，并采用一种轻量级解码器，从潜在表示和掩码标记重新构造完整信号。**图1说明了下面介绍的想法。

MASKING

按照ViT，将图像划分为规则的非重叠的patches。然后，对一个子集的patches进行采样，并MASK（即移除）剩余的面片。采样策略：在不替换的情况下，按照均匀分布对随机patches进行采样。称之为“随机抽样”。

具有高掩蔽率的随机采样在很大程度上消除了冗余，因此产生了一个无法通过从可见相邻面片外推来轻松解决的任务。均匀分布可防止潜在的中心偏移（即，图像中心附近被mask更多）。最后，高度稀疏的输入为设计高效编码器创造了机会。

MAE编码器

编码器使用ViT，但仅适用于可见的、没被mask的patches。就像在标准ViT中一样，编码器通过添加positional embeddings的Linear Projection来得到embedded patches，然后通过一系列transformer block处理结果集。然而，编码器只在整个集合的一小部分上运行。去除masked patches；不使用mask tokens。(也就是说直接去除而非填充零值)，这使MAE能够仅用一小部分计算和内存来训练非常大的编码器。全套由轻量级解码器处理，如下所述

MAE解码器

MAE解码器的输入是由：（i）经过encoder编码的可见patches；（ii）masked tokens。每个masked tokens都是一个共享的学习向量，表示要预测的缺失patches。向这个完整集合中的所有tokens添加positional embedding。解码器具有另一系列Transformer模块。

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

重构目标

MAE通过预测每个masked patches的像素值来重构输入值。解码器输出中的每个元素都是表示patches的像素值向量。解码器的最后一层是一个linear projection，其输出通道的数量等于每个patch中包含的像素数量。解码器的输出形成重构图像。损失函数计算像素空间中重建图像和原始图像之间的均方误差（MSE）：只计算被mask的patches的损失，类似于BERT。

作者还研究了一种变体，其重建目标是每个masked patches的归一化像素值。具体来说，计算每个patches中所有像素的平均值和标准偏差，并使用它们来规范化patches。在实验中，使用归一化像素作为重构目标提高了表示质量。

简单的实现

MAE预训练可以被高效实现，并不需要任何专门的稀疏化操作。首先，为每个输入patch生成一个对应token（即添一个positional embedding的Linear projection，正如ViT所做的那样）。接下来，随机shuffle这些token的列表，并根据dropout rate删除列表的最后那部分。此过程为encoder生成一小部分（未被mask）的token，而无需更换采样的patches。编码后，将masked token列表附加到已编码的patch列表中，并unshuffle此完整列表，以将所有token与其目标对齐。解码器应用于这个完整的、添加了poisitional embedding的列表。如前所述，这不需要稀疏操作，仅引入了可忽略不计的开销，因为shuffle和unshuffle的操作速度很快。

实验

与有监督/无监督的预训练方法进行比较

• 比MOCOv3（对比学习）的预训练finetune效果更好

  

• 在目标检测任务中，ViT与FPN搭配更好，与有监督的预训练相比，MAE要更优。

  

总结

MAE是一种很好的自监督训练器，其训练的预训练模型，比有监督的、MOCO自监督的，在下游任务finetune上都取得更好的效果