【论文精读】MAE

摘要

       将掩码重建任务从nlp引入到cv，提出非对称掩码自编码器。

框架

概述

       如上图，本文提出掩码自编码器，即将给定原始信号的部分观测值的情况下重建原始信号，编码器将观察到的部分信号(没有掩码标记)映射到潜在表示，采用轻量级的解码器从潜在表示重建原始信号，模型采用了非对称设计。

掩码

       类似ViT，将图像划分为规则的非重叠patch，按照均匀分布对无替换的随机patch进行采样，即随机抽样。随后屏蔽（删除patch）其余未被采样的patch。


       采用高屏蔽率（删除patch的比率）的随机抽样可以在很大程度消除了冗余，从而创建了一个不能通过可见的邻近patch外推而轻松解决的任务（如上几幅图）。而均匀分布可以防止中心归纳偏好（图像中心附近有更多的掩码块）。

MAE编码器

       编码器为原始ViT，且只应用未屏蔽的patch，并采用线性投影计算这些patch的patch embedding，并添加position embedding，然后通过一系列Transformer块处理结果集。

MAE解码器

       如图1，解码器的输入是完整的patch集，包括编码器输出的未屏蔽patch的特征token和mask tokens。其中，每个mask tokens都是一个共享的、可学习的向量，表示要预测的缺失patch。随后向这个完整patch集合添加position embedding。

       解码器仅在预训练任务中重建图像，而其余的下游任务形式多种多样，实际应用时不用解码器。所以解码器的设计和编码器是解耦的，解码器可以设计得简单、轻量（比编码器更窄、更浅。窄：对应通道数；浅：对应深度）。

目标重建

       MAE通过预测每个掩码patch的像素值来重建输入。解码器的最后一层为线性投影，其输出通道的数量等于patch中的像素值的数量，其输出代表输入patch的像素值向量，并将输出重塑以形成重建的图像。损失函数为计算像素空间中重建图像和原始图像之间的均方误差(MSE)。且仅在掩码patch上计算损失。

       本文还研究了一种变体，其重建目标是每个掩码 patch的归一化像素值。即计算一个掩码 patch中所有像素的均值和标准差，然后用来归一化这对应的掩码patch作为目标。实验中，使用归一化像素作为重建目标可以提高表示质量。

实现

• 将输入图像划分为patches$(B,C,H,W)\to (B,N,P\times P\times C)$
• 对所有patch计算patch embedding，生成tokens，并加入position embeddings，维度变换为$(B,N,P\times P\times C)\to (B,N,dim)$
• 根据预设的掩码比例(75%)，使用服从均匀分布的随机采样策略采样一部分前一步骤得到的tokens送给编码器，另一部分被mask掉
• 将编码器编码后的tokens与masked tokens按照原先在patch形态时对应的次序拼在一起，输入给解码器
• 解码器解码后取出masked tokens对应的部分送入到全连接层，对masked patches的像素值进行预测，最后将预测结果与masked patches进行比较，计算MSE loss

实验

ImageNet实验

       实验在ImageNet-1K（IN1K）训练集上采用ViT进行自监督预训练，然后进行监督训练，评估端到端微调或线性检测的表现。

实验配置

       ViT架构实验遵循标准的ViT架构。其有一个Transformer块的堆栈，其中每个块由一个多头自注意力块和一个MLP块组成，两者都具有LayerNorm（LN）。编码器以LN结尾。由于MAE编码器和解码器宽度不同，在编码器之后采用线性投影层对其进行调整匹配。MAE向编码器和解码器输入添加了位置嵌入，没有使用相对位置或层缩放。

       从编码器输出中提取特征以进行微调和线性探测。由于ViT有一个类标记，为了适应这种设计在MAE预训练中向编码器输入添加了一个辅助token。此token被视为在线性探测和微调中训练分类器的类token，实验证明MAE在没有这个token的情况下(使用平均池化)也同样表现得很好。

       上图为预训练配置。相对于ViT官方源码，没有使用color jittering、drop path、gradient clip。

       上图为微调及部分微调配置。微调使用layer-wise learning rate decay，部分微调只微调编码器的最后几个层。

       上图为线性检测配置。线性检测和微调有不同，regularization对线性检测可能会损失模型性能，因此舍弃了一些regularization strategies，例如mixup、cutmix、drop path、color jittering。并对输入执行normalization,同时在预训练特征层和线性分类层之间加上一个额外的BN，可以实现特征的标准化。

消融实验

       上图实验mask比例。观察到微调和线性检测时，mask比例逐渐升高，性能更好。 线性检测近似线性增涨，而微调则是mask比例在30%~40%时激增，而后就倾向于饱和。

       线性检测时为mask比例越高，预训练时得到的编码器越强，这部分在下游任务中不再被训练的了，所以其性能就随着mask比例的增加呈线性增涨的趋势。

       微调时，结果对mask比例不太敏感，大范围的屏蔽比率(40- 80％)都可以很好地工作。最优mask比例是75%。如上图可视化结果。

       上图a显示解码器深度对微调和线性探测任务的影响。观察到足够深的解码器对线性探测很重要，因为自编码器的最后几层更专注于重建，与识别不太相关，预训练与线性检测之间存在gap。故合理深度的解码器（8层）可以使潜在表示的语义信息更抽象，可以改进8％的线性探测精度。对于微调，解码器深度对改进微调的影响较小，故只使用单层解码器的MAE也可以在微调中表现良好(84.8％)，这么小的解码器可以进一步加快训练速度。

       图b实验解码器宽度，也得出与a类似的结论，通道数为512时微调和线性探测精度最优。故编码器默认是8个blocks，通道数是512。

       图c实验编码器是否接受mask tokens。实验表明，encoder如果接收mask tokens，性能会降低（线性检测降低14%），因此编码器只接收visible tokens，既能提升性能，又能降低计算量，且加速训练。decoder越小/encoder越大，加速越明显，MAE的时间和存储效率使其适合训练大型模型。详细实验配置、结果如上图。

       图d实验重建目标方式。实验表明使用normalization（per patch）的效果比不使用更好。对patch执行PCA，然后将前96个主成分作为重构目标，效果并不好。使用BEiT的dVAE作为目标，效果也不太好。

       图e实验数据增强方法。实验表明，MAE只使用裁剪增强时就可以表现很好，添加颜色抖动会降低结果。同时发现不使用数据增强，MAE甚至性能更好。对MAE来说，每一个轮次masks的token都不同，随机masking扮演了类似于数据增强的角色。

       图f实验不同的掩码采样策略。观察到block掩码策略倾向于删除较大的块（上图中），基于block掩码的MAE在50％的比例下工作得相当好，但在75％的比例下退化。该任务比random采样更难，具有更高的训练损失，重建结果也比较模糊。 其次grid采样（上图右）是一个更容易的任务，具有更低的训练损失。然而，其重建更加尖锐。表示质量较低。故简单random采样最适合MAE，其允许更高的掩蔽率，这提供了更大的加速效益，同时也具有良好的准确性。

       上图为训练epoch数量对模型精度的影响。观察到，随着训练时间的延长准确率稳步提高。在1600个epoch训练后，线性检测精度依然没有达到饱和。

对比实验

       上图为与先前自监督方法比较了自监督ViT模型的微调结果。对于ViT-B，所有方法的性能都很好，对于ViT-L，方法之间的差距更大。此外，观察到MAE可以轻松扩展，并从更大的模型中显示出稳定的改进。使用ViT-H获得了86.9％ 的准确率，将ViT-H微调为448大小，只使用IN1K数据就实现了87.8％的准确率。

       与BEiT相比，MAE具有更高的准确性、更简单、更快的计算速度。MAE重建像素，与预测token的BEiT形成对比。

       上图为与有监督预训练ViT的比较。

部分微调实验

       微调ViT最后几层，同时冻结其他层进行训练实验。

       上图显示结果。只微调一个Transformer块就可以将精度从73.5％提高到81.0％，如果只微调最后一个块的一半（其MLP子块），可以得到79.1％的精度。微调一部分块(4或6)可以实现接近完全微调的精度。

       与MoCo v3进行比较。MoCo v3具有更高的线性探测精度，但其所有的部分微调结果都比MAE差。当调优4个区块时，差距为2.6％。虽然MAE表示的线性可分性较差，但它们是更强的非线性特征。

迁移学习实验

       上图在COCO上对Mask RCNN端到端进行微调，ViT骨干经过调整用于FPN。这种方法应用于上图中的所有条目，指标为目标检测的box AP和实例分割的mask AP。

       与有监督的预训练相比，MAE在所有配置下的表现都更好。在较小的ViT-B下，MAE比有监督预训练高2.4个点。对于更大的ViT-L，MAE预训练比有监督预训练高出4.0个点。 基于像素的MAE优于或与基于token的BEiT相当，而MAE更简单、更快，MAE和BEiT都优于MoCo v3, MoCo v3与有监督预训练相当。

       上图使用UperNet在ADE20K实例分割任务上进行实验。观察到，MAE预训练比有监督预训练的结果提高了3.7个点。基于像素的MAE也优于基于token的BEiT。

       上图为在iNaturalists和Places实验分类任务。在iNat上，MAE方法下随着ViT模型的扩大，精度有了很大的提高，大大超过了之前的最佳结果。在Places上，MAE的性能超过了之前的最佳结果，之前的这些结果通过对数十亿张图像进行预训练获得。

       上图实验比较像素和dVAE作为MAE重建目标的结果。虽然使用dVAE token比使用非归一化像素要好，但其在统计上与使用归一化像素相似。表明token化对MAE是不必要的。

鲁棒性评估

       上图评估了模型在不同变体的ImageNet验证集上的鲁棒性。使用在原始ImageNet上进行微调的相同模型(表3)，并仅在不同的验证集上运行推理。

       观察到除IN-C外，MAE在增加模型大小后具有显著的收益。增大图像大小在所有集合中都有帮助，也大大超过了以前的最佳结果。 相比之下，有监督训练的表现要差得多。例如使用ViT-H，MAE预训练在IN-A上比有监督的对应模型好35％。

可视化

       上图显示了在imagenet和COCO上的随机样本重建可视化结果三元组，其中左为掩码图像、中为MAE重建、右为真实值。掩码率为75％。

reference

He, K. , Chen, X. , Xie, S. , Li, Y. , Dollár, Piotr, & Girshick, R. . (2021). Masked autoencoders are scalable vision learners.