type
status
password
date
slug
summary
category
URL
tags
icon
Abstract
掩码自编码器( MAE, Masked autoencoders)是视觉领域中的BERT模型。MAE具体操作为随机屏蔽输入image中的patchs(屏蔽高达75%)，再重建丢失的像素。在下游任务的迁移性优于有监督的预训练。它主要基于以下两个核心设计：
- MAE采用不对称的编码器-解码器架构，编码器参数量大于解码器参数量；
- 编码器只作用在可见的patch中，对于屏蔽的patch，编码器不会对它进行编码，这样能够节省一定的计算资源
- 解码器是一个比较轻量的解码器，拿到编码器的输出之后，重构被屏蔽的块
- 遮住了大量的块（比如说把75%的块全部遮住），迫使模型去学习一些更好的特征，从而得到一个较好的自监督训练效果。如果只是遮住几块的话，只需要进行插值就可以出来了，模型可能学不到特别的东西。编码器只需要编码1/4大小的图片，可以降低计算量，训练加速3倍或者以上。
一、Introduction
带掩码的自编码器在CV领域仍然落后与NLP领域的原因：
- 架构不同。在CV领域，CNN在过去十年中一直占据主导地位，并且在CNN中添加掩码令牌和位置编码都有困难，后边的ViT已经解决了这个问题。
- 信息密度不同。语言的信息密度非常高，但图片的信息密度较低。在语言中，训练一个模型去预测一个句子中缺失的几个词就会诱发对整个句子的复杂性的理解；而在视觉中，利用周围的patches去重建被屏蔽的patch只产生很少的理解（只需要进行插值就可以重建）。因此，采用高比例随机地屏蔽patches，减少冗余信息，提高对图像的整体理解。
- 解码器不同。在视觉领域，解码器重建像素，因此其输出的语义层级低于常见的识别任务。在语言领域，BERT 中的解码器预测包含丰富语义信息的缺失单词，并且非常简单（如仅含 MLP）。我们发现，对于图像而言，解码器设计对所学习的潜在表示的语义层级起着关键作用。
MAE中具有不对称的编码器-解码器架构，编码器随机地屏蔽掉输入图像的patches，解码器根据潜在表示和掩码令牌恢复缺失的图像。在不对称的编码器-解码器架构中，将掩码令牌加入到轻量级解码器中大大减少了计算量。同时遮掩比（75%）既提高了精度，又使得编码器只需要处理一小部分patches，减少了训练时间和内存损耗。
二．Related Work
Masked Language modeling（掩码语言模型）和自回归模型，例如GPT的自回归语言模型和BERT的掩码语言建模在NLP领域都取得了很大的成功。这些方法都是保留输入序列的一部分，去预测缺失的部分。这些方法具有很强的扩展性且已经证明可以广泛应用到各种下游任务中去。
Auto encoders(自编码器)是学习表示中一种经典的方法。其中编码器将输入映射到潜在表示，解码器恢复输入信息。去噪自编码器是一种破坏输入信号，学习恢复输出信号的自编码器。MAE就是一种去噪自编码器。
掩码图像编码方法通过对被掩码破坏的图像进行学习来获取表示。掩码作为去噪自动编码器中的一种噪声类型。Context Encoder使用卷积网络对大面积缺失区域进行修复。iGPT对像素序列进行操作并预测未知像素。ViT 论文研究了用于自监督学习的掩码图像块预测任务。最近，BEiT提出预测离散token。
计算机视觉领域的自监督学习方法引起了广泛关注，这些方法通常侧重于不同的预训练前置任务。
三、Approach
我们的掩码自动编码器（MAE）与所有自动编码器一样，包含一个将观测信号映射到潜在表示的编码器，以及一个从潜在表示重建原始信号的解码器。与经典自动编码器不同的是，我们采用非对称设计：编码器仅仅处理未被屏蔽掉的信号，大大减少了计算复杂度，而轻量级解码器根据潜在表示和掩码标记重建原始信号。下图展示了这一思路。
- 将图像划分成 patches：(B,C,H,W)->(B,N,PxPxC)；
- 对各个 patch 进行 embedding(实质是通过全连接层)，生成 tokens，并加入位置信息(position embeddings)：(B,N,PxPxC)->(B,N,dim)；
- 根据预设的掩码比例(paper 中提倡的是 75%)，使用服从均匀分布的随机采样策略采样一部分 tokens 送给 Encoder(喂给Encoder的图像块乱序，但是位置信息有序)，另一部分“扔掉”(mask 掉)；先加入了位置信息再进行mask。
- 将 Encoder 编码后的 tokens 与 加入位置信息后的 masked tokens 按照原先在 patch 形态时对应的次序拼在一起，然后喂给 Decoder 玩(如果 Encoder 编码后的 token 的维度与 Decoder 要求的输入维度不一致，则需要先经过 linear projection 将维度映射到符合 Decoder 的要求)；
- Decoder 解码后取出 masked tokens 对应的部分送入到全连接层，对 masked patches 的像素值进行预测，最后将预测结果与 masked patches 进行比较，计算 MSE loss
3.1、掩码策略
学习ViT中的操作，将图片打散成若干不重叠的patches。再按照随机分布来对patches进行采样，进而屏蔽掉其他的patches。采用较高的遮掩比例既大大减少的数据冗余，减少计算量，又避免了通过邻近patches简单恢复缺失patch的问题。均匀分布避免了潜在的中心偏差（即图像中心附近掩码块更多的情况）。
3.2、MAE Encoder
编码器是一个仅对未被屏蔽的输入信号进行操作的ViT。和标准的ViT一样，对输入的图片进行线性投影，后加上位置编码，最后通过一系列的Transformer块对数据操作。这里的编码器仅需要处理未被屏蔽掉的信号，大大减少计算量，这就可以再较少的内存和计算量上训练更大的编码器。
3.3、MAE Decoder
MAE 解码器的输入是完整的图像块（包括编码后的可见图像块以及掩码标记），如图 1 所示。每个掩码标记是一个共享的可学习向量，表示需要预测的缺失图像块。我们向所有标记（屏蔽掉和未屏蔽掉的）添加位置嵌入；因为，如果没有位置嵌入，掩码标记将无法获取其在图像中的位置信息。
MAE 解码器仅在预训练阶段用于执行图像重建任务（下游任务中仅使用编码器生成图像表示）。因此，解码器架构可以独立于编码器进行灵活设计。我们尝试了非常小的解码器，其宽度和深度均小于编码器。通过这种非对称设计，完整的标记集合仅由轻量级解码器处理，这显著减少了预训练时间。
当 Decoder 更深/宽时，它本身会拥有更强的重建能力，这样就使得在预训练时 Encoder 能更专注于提取抽象语义层级的特征。
3.4、Reconstruction target
我们的 MAE 通过预测每个掩码图像块的像素值来重建输入图像。解码器输出的每个元素是一个表示图像块像素值的向量。解码器的最后一层是一个线性投影层，其输出通道数等于图像块的像素值数量。解码器的输出被重新塑形以形成重建图像。我们的损失函数在像素空间中计算重建图像与原始图像之间的均方误差（MSE）。与 BERT 类似，我们仅对掩码图像块计算损失。
我们还研究了一种变体，其重建目标是每个掩码图像块的归一化像素值。具体来说，我们计算一个图像块中所有像素的均值和标准差，并使用它们对该图像块进行归一化。在实验中，使用归一化像素作为重建目标提高了表示质量。
3.5、Simple implementation
MAE的预训练非常简单，不需要任何稀疏操作。首先，为我们通过线性投影（添加位置嵌入）为每个输入图像块生成一个标记。然后打乱标记顺序，按照遮掩比例屏蔽掉一部分的Patches。随后，使用编码器对未屏蔽的patches进行编码。编码完成后，往已经编码的patches中添加掩码标记，此时不打乱顺序并且使其和目标标记对齐。
四、ImageNet Experiments
先在ImageNet-1K上做自监督预训练，再在同样的数据集上做有标签的监督训练。有两种方法：（1）基于end to end 的微调，允许改模型中所有可学习参数（2）linear probing：只允许改最后一层的线性输出层
4.1、主要特性
我们使用表 1 中的默认设置对 MAE 进行消融实验（见标题），观察到几个有趣的特性。
4.1.1、掩码比例：
图 5 展示了掩码比例的影响，最佳比例高得惊人。75% 的比例对线性探测和微调都很有效。这一现象与BERT形成对比，后者的典型掩码比例为 15%。我们的掩码比例也远高于计算机视觉领域的相关工作（20% 至 50%）。
对于linear probing，准确率随着掩码比例的增加稳步上升，直到达到最佳点，准确率差距可达约 20%（54.6% vs. 73.5%）；对于微调，结果对比例的敏感性较低，掩码比例40-80%都能取得较好效果。图 5 中的所有微调结果均优于从头开始训练的结果（82.5%）。
4.1.2、解码器设计
我们的 MAE 解码器可灵活设计，如表 1a 和 1b 所示。 MAE 的 默认解码器是轻量级的，它有 8 个宽度为 512 维的Transformer块组成，每个标记的计算量仅为 ViT-L（24 块，1024 维）的 9%。因此，尽管解码器处理所有标记，其计算量仍占整体的很小一部分。
- 解码器深度：表 1a 改变解码器深度（Transformer 块数量）。足够深的解码器对linear probing很重要，这可以解释为像素重建任务与识别任务之间的差距：自动编码器的最后几层更专注于重建，而与识别任务的相关性较低。合理深度的解码器可以使潜在表示保持在更抽象的层级，这可使linear probing的准确率提升高达 8%（表 1a，“lin”）。然而，在使用微调时，编码器的最后几层可以通过调整适应识别任务，此时解码器深度对提升微调效果的影响较小（表 1a，“ft”）。
- 解码器宽度：表 1b 研究解码器宽度（通道数）。我们默认使用 512 维，在微调与linear probing中均表现良好。更窄的解码器在微调时也能正常工作。
4.1.3、掩码标记
在编码器中要不要加入被盖住的那些块。结果显示，不加入掩码的patches，精度反而更高一些，而且计算量更少。
如果编码器使用掩码标记，性能会下降：linear probing准确率降低 14%。在这种情况下，预训练与部署之间存在差距——编码器在预训练中输入包含大量掩码标记，而在未损坏的图像中不存在这些标记，这一差距可能导致部署时准确率下降。通过从编码器中移除掩码标记，我们限制编码器始终处理真实图像块，从而提高准确率；此外还可以大幅减少训练计算量。在表 1c 中，我们将整体训练 FLOPs 减少了 3.3 倍，这使我们的实现中的实际训练速度提升了 2.8 倍。
4.1.4、重建目标
使用归一化像素可提高准确率，这种每块归一化增强了局部对比度。表明，高频分量在我们的方法中是有用的。
4.1.5、数据增强
使用仅裁剪的增强（固定大小或随机大小，均包含随机水平翻转）即可表现良好。这一特性与对比学习及相关方法截然不同，后者严重依赖数据增强。
使用仅裁剪的增强（固定大小或随机大小，均包含随机水平翻转）即可表现良好。添加颜色抖动会降低结果，因此我们在其他实验中不使用它。
令人惊讶的是，即使不使用数据增强（仅中心裁剪，无翻转），MAE 也能表现良好。这一特性与对比学习及相关方法截然不同，后者严重依赖数据增强。有研究表明，仅使用裁剪增强会使 BYOL 和 SimCLR 的准确率分别降低 13% 和 28%。
4.1.6、掩码采样策略
我们在表 1f 中比较了不同的掩码采样策略，如图 6 所示。
- 块级掩码策略（[2] 提出）倾向于移除大的块（图 6 中间）。我们的 MAE 在掩码比例为 50% 时使用块级掩码表现合理，但在比例为 75% 时性能下降。该任务比随机采样更难，训练损失更高，重建结果也更模糊。
- 网格采样策略定期保留每四个块中的一个（图 6 右侧），这是一个更简单的任务，训练损失更低，重建更清晰，但表示质量较低。
- 简单的随机采样对我们的 MAE 效果最佳，它允许更高的掩码比例，在实现更大速度提升的同时保持良好的准确率。
4.1.7、训练时长
到目前为止，我们的消融实验基于 800 epoch 的预训练。图 7 展示了训练时长的影响，准确率随着训练时间的延长稳步提高。事实上，即使在 1600 epoch 时，我们也未观察到 linear probing 准确率饱和。这种行为与对比学习方法（如 MoCo v3）不同，后者的 ViT-L 在 300 epoch 时就已饱和。请注意，MAE 编码器每个 epoch 仅处理 25% 的图像块，而对比学习中编码器每个 epoch 处理 200%（双裁剪）甚至更多（多裁剪）的图像块。
4.2、与先前结果的比较
4.2.1、与自监督方法的比较
在表 3 中，我们比较了自监督 ViT 模型的微调结果。对于 ViT-B，所有方法表现接近；对于 ViT-L，方法之间的差距更大，这表明更大的模型面临的挑战是减少过拟合。
我们的 MAE 可以轻松扩展，并且随着模型增大表现稳步提升。使用 ViT-H（224 尺寸），我们获得了 86.9% 的准确率；通过 448 尺寸的微调，仅使用 IN1K 数据就达到了 87.8% 的准确率。在所有仅使用 IN1K 数据的方法中，之前的最佳准确率是 87.1%（512 尺寸），其基于先进的网络架构。在竞争激烈的 IN1K 基准（无外部数据）中，我们以显著优势超越了现有技术。我们的结果基于原始 ViT，预计先进的网络架构会表现得更好。
与 BEiT 相比，我们的 MAE 更准确，同时更简单、更快。
4.3、部分微调
表 1 显示，linear probing和微调结果基本不相关。根据图 9 显示，仅微调一个 Transformer 块就能将准确率从 73.5% 显著提升至 81.0%。此外，如果我们仅微调“半个”最后一个块（即其 MLP 子块），可以达到 79.1%，远高于inear probing。该变体本质上是微调一个 MLP 头。微调几个块（如 4 或 6 个）可以达到接近全量微调的准确率。
在图 9 中，我们还与对比方法 MoCo v3（提供了 ViT-L 结果）进行了比较。MoCo v3 具有更高的inear probing准确率，但其所有部分微调结果均低于 MAE。微调 4 个块时，差距为 2.6%。尽管 MAE 的表示线性可分性较差，但它们是更强的非线性特征，在微调非线性头时表现良好。
这些观察表明，线性可分性并非评估表示质量的唯一指标。已有研究表明，inear probing与迁移学习性能（如目标检测）的相关性并不强。据我们所知，线性评估在 NLP 预训练基准中并不常用。
五、迁移学习实验
我们使用表3中的预训练模型，对下游任务中的迁移学习进行评估。
5.1、目标检测与分割
我们在COCO数据集上对Mask R-CNN进行端到端微调。ViT主干网络适配于FPN（详见A.3）。我们对表4中的所有条目均采用此方法，报告目标检测的边界框AP（AP box）和实例分割的掩码AP（AP mask）。
与监督预训练相比，我们的MAE在所有配置下表现更优（表4）。使用较小的ViT-B时，MAE比监督预训练高出2.4个点（50.3 vs. 47.9，AP box）；更显著的是，使用较大的ViT-L时，MAE预训练比监督预训练高出4.0个点（53.3 vs. 49.3）。
基于像素的MAE优于或等同于基于标记的BEiT，同时MAE更简单、更快。MAE和BEiT均优于MoCo v3，而MoCo v3与监督预训练表现相当。
5.2、语义分割
我们使用UperNet在ADE20K数据集上进行实验（详见A.4）。表5显示，我们的预训练结果显著优于监督预训练，例如ViT-L提升了3.7个点。我们的像素级MAE也优于标记级BEiT，这些观察结果与COCO实验一致。
5.3、分类任务
表6研究了在iNaturalists[[56]]和Places[[71]]任务上的迁移学习（详见A.5）。在iNat数据集上，我们的方法展现出强大的规模扩展能力：模型越大，准确率提升越显著，结果大幅超越先前最佳水平。在Places数据集上，MAE优于先前基于数十亿图像预训练的最佳结果[[19,40]]。
6、讨论与结论
简单且可扩展的算法是深度学习的核心。在NLP领域，简单的自监督学习方法（如）通过指数级扩展模型规模实现了显著突破。在计算机视觉领域，尽管自监督学习取得了一定进展，但实际应用中的预训练范式仍以监督学习为主（如[[33,51,25,16]]）。本研究中，我们在ImageNet和迁移学习任务中观察到，自动编码器——这种与NLP技术类似的简单自监督方法——能够带来可扩展的性能提升。视觉领域的自监督学习可能正沿着与NLP相似的轨迹发展。