大火的何铠明：MAE——用于计算机视觉的可扩展自监督学习神器

作者：王浩 毕业于北京航空航天大学，人工智能领域优质创作者，CSDN博客认证专家
首发：公众号【3D视觉开发者社区】

导语：

近期，何铠明的新作可谓是火出了圈，毕竟何佬出品必是精品，由何佬提出的的ResNet、Faster RCNN等模型一直被大家学习和研究。如今，何铠明又带来一种用于计算机视觉的可扩展自监督学习器, 称之为”掩码自编码器 (MAE) “，本文将对该视觉学习器原理及实现方法进行解读。

Masked Autoencoders Are Scalable Vision Learners

论文地址：https://arxiv.org/pdf/2111.06377.pdf

摘要

何铠明在CVPR2021提出的一种用于计算机视觉的可扩展自监督学习器—Masked Autoencoders，目前大有获得“CVPR2022最佳论文”的趋势。这篇论文最神奇的地方在于：整篇论文没有公式！颇有“大道至简”的风范。

MAE 方法很简单：屏蔽输入图像的随机块并重建丢失的像素。它主要基于两个核心设计：

1、开发了一个非对称编码器-解码器架构，其中一个编码器只对可见的补丁子集进行操作（没有掩码标记），以及一个轻量级的解码器，可以从潜在表示和掩码标记重建原始图像。

2、发现屏蔽大部分输入图像（例如75%）会产生重要且有意义的自我监督任务。

结合这两种设计能够高效地训练大型模型：

• 加速训练（3倍或更多）并提高模型的精度。

• 可扩展方法允许学习泛化良好的高容量模型：比如，在仅使用ImageNet-1K数据的方法中，vanilla ViT-Huge 模型实现了最佳准确率(87.8%)。下游任务中的传输性能优于有监督的预训练，显示出MAE的可扩展能力。

简介

深度学习见证了能力和容量不断增长的架构的爆炸式增长。在硬件快速增长的帮助下，今天的模型很容易过拟合一百万张图像，并开始需要数亿张——通常是公开无法访问的——标记图像。

这种对数据的需求已经通过自监督预训练在自然语言处理(NLP)中成功解决。基于 GPT中的自回归语言建模和BERT中的掩码自动编码的解决方案在概念上很简单：它们移除了一部分数据并学习预测移除的内容。这些方法现在可以训练包含超过一千亿个参数的可泛化NLP模型。

掩码自动编码器的思想，一种更通用的去噪自动编码器的形式，很自然，也适用于计算机视觉。事实上，在BERT之前，与视觉密切相关的研究。然而，尽管随着BERT 的成功，人们对这个想法产生了浓厚的兴趣，但视觉中自动编码方法的进展落后于 NLP。是什么使视觉和语言之间的掩码自动编码不同？从以下几个角度来回答这个问题：

(i) 直到最近，架构还是不同的。在视觉方面，卷积网络在过去十年中占主导地位。卷积通常在规则网格上运行，将诸如掩码标记或位置嵌入之类的“指标”集成到卷积网络中并非易事。然而，随着 Vision Transformers (ViT) 的引入，这种架构差距已经得到解决，不应再成为障碍。

(ii) 语言和视觉的信息密度不同。语言是人类生成的具有高度语义和信息密集度的信号。当训练一个模型来预测每个句子的几个遗漏单词时，这个任务似乎会导致复杂的语言理解。相反，图像是具有大量空间冗余的自然信号——例如，可以从相邻的补丁中恢复丢失的补丁，而对零件、对象和场景的高级理解很少。为了克服这种差异并鼓励学习有用的特征，展示了一个简单的策略在计算机视觉中效果很好：屏蔽了很大一部分的随机补丁。这种策略在很大程度上减少了冗余并创建了一项具有挑战性的自我监督任务，需要超越低级图像统计的整体理解。要对重建任务有一个定性的认识，请参见Figure 2-4。

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

在此分析的推动下，提出了一种简单、有效且可扩展的掩码自编码器 (MAE) 形式，用于视觉表示学习。MAE 从输入图像中屏蔽了随机补丁，并在像素空间中重建丢失的补丁。它具有非对称编码器解码器设计。编码器只对可见的补丁子集（没有掩码标记）进行操作，解码器是轻量级的，可以从潜在表示和掩码标记中重建输入（Figure 1）。在非对称编码器-解码器中将掩码标记转移到小型解码器会导致计算量大幅减少。

在这种设计下，非常高的掩蔽率（例如 75%）可以实现双赢：它优化了准确性，同时允许编码器仅处理一小部分（例如 25%）的补丁。这可以将整体预训练时间减少 3 倍或更多，同样减少内存消耗，能够轻松地将 MAE 扩展到大型模型。

MAE 学习了非常高容量的模型，可以很好地泛化。通过 MAE 预训练，可以在 ImageNet-1K 上训练 ViT-Large/Huge等数据饥渴模型，并提高泛化性能。使用普通的 ViT-Huge 模型，在 ImageNet-1K 上微调时达到了 87.8% 的准确率。这优于之前仅使用 ImageNet-1K 数据的所有结果。

相关工作

掩码语言建模及其自回归对应物，例如 BERT 和 GPT，是非常成功的 NLP 预训练方法。这些方法保留一部分输入序列并训练模型来预测缺失的内容。这些方法已被证明具有出色的扩展性，并且大量证据表明，这些预训练的表示可以很好地推广到各种下游任务。

自编码是学习表征的经典方法。它有一个将输入映射到潜在表示的编码器和一个重构输入的解码器。例如，PCA 和 k-means 是自动编码器。去噪自动编码器 (DAE) 是一类自动编码器，它破坏输入信号并学习重建原始的、未破坏的信号。一系列方法可以被认为是不同损坏下的广义 DAE，例如，屏蔽像素 或去除颜色通道。MAE 是一种去噪自动编码形式，但在许多方面与经典 DAE 不同。

自监督学习方法对计算机视觉产生了浓厚的兴趣，通常侧重于不同的预训练任务。最近，对比学习很流行，它对两个或多个视图之间的图像相似性和不相似性进行建模。对比和相关方法强烈依赖于数据增强。自动编码追求概念上不同的方向，它表现出将要展示的不同行为。

方法

掩码自编码器 (MAE) 是一种简单的自编码方法，可以在给定部分观察的情况下重建原始信号。与所有自动编码器一样，方法有一个编码器，将观察到的信号映射到潜在表示，以及一个解码器，从潜在表示重建原始信号。与经典的自动编码器不同，采用非对称设计，允许编码器仅对部分观察到的信号进行操作（没有掩码标记）和一个轻量级解码器，从潜在表示和掩码标记重建完整信号。Figure 1 说明了这个想法，接下来将逐个介绍。

掩码

遵循 ViT，将图像划分为规则的非重叠补丁。然后对补丁的子集进行采样并屏蔽（即删除）剩余的补丁。采样策略很简单：按照均匀分布对随机补丁进行采样，无需替换。简单地将其称为“随机抽样”。

具有高掩蔽率（即移除补丁的比率）的随机采样在很大程度上消除了冗余，从而创建了一项无法通过从可见相邻补丁外推来轻松解决的任务（见Figure 2-4）。均匀分布可防止潜在的中心偏差（即在图像中心附近有更多被屏蔽的补丁）。最后，高度稀疏的输入为设计高效编码器创造了机会，接下来介绍。

MAE 编码器

编码器是 ViT ，但仅应用于可见的、未屏蔽的补丁。就像在标准 ViT 中一样，编码器通过添加位置嵌入的线性投影嵌入补丁，然后通过一系列 Transformer 块处理结果集。然而，编码器只对整个集合的一小部分（例如 25%）进行操作。被屏蔽的补丁被移除；不使用掩码令牌。能够仅使用一小部分计算和内存来训练非常大的编码器。全套由轻量级解码器处理，如下所述。

MAE 解码器

MAE 解码器的输入是完整的令牌集，包括编码的可见补丁和掩码令牌。参见Figure 1。每个掩码标记是一个共享的、学习的向量，表示存在要预测的缺失补丁。为这个完整集合中的所有标记添加了位置嵌入；如果没有这个，掩码令牌将没有关于它们在图像中的位置的信息。解码器有另一系列的 Transformer 模块。

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

重建目标

MAE 通过预测每个掩码补丁的像素值来重建输入。解码器输出中的每个元素都是一个表示补丁的像素值向量。解码器的最后一层是线性投影，其输出通道的数量等于补丁中像素值的数量。解码器的输出被重新整形以形成重建的图像。损失函数计算像素空间中重建图像和原始图像之间的均方误差(MSE)。

还研究了一个变体，其重建目标是每个蒙版块的归一化像素值。具体来说，计算一个补丁中所有像素的均值和标准差，并使用它们来规范化这个补丁。在实验中，使用归一化像素作为重建目标可以提高表示质量。

简单的实现

MAE预训练可以高效实施，重要的是，不需要任何专门的稀疏操作。首先，为每个输入补丁生成一个标记（通过线性投影和添加的位置嵌入）。接下来，根据掩蔽率随机打乱令牌列表并删除列表的最后一部分。此过程为编码器生成一小部分标记，相当于采样补丁而无需替换。编码后，将掩码标记列表附加到编码补丁列表中，并解散这个完整列表（反转随机洗牌操作）以将所有标记与其目标对齐。解码器应用于此完整列表（添加了位置嵌入）。如前所述，不需要稀疏操作。这个简单的实现引入的开销可以忽略不计，因为shuffle和unshuffling操作很快。

ImageNet实验

在ImageNet-1K (IN1K) 训练集上进行自我监督的预训练，然后进行有监督的训练，以通过端到端微调或线性探测来评估表示。报告了单个224×224裁剪的top-1 验证精度。

基线：ViT-大，在消融研究中使用ViT-Large (ViT-L/16) 作为主干，ViT-L 非常大（比 ResNet-50 大一个数量级）并且容易过拟合。以下是从头开始训练的ViT-L 与从基线 MAE 微调的比较：

scratch, original	scratch, our impl.	baseline MAE
76.5	82.5	84.9

注意到从头开始训练有监督的 ViT-L 是很重要的，并且需要一个具有强正则化的良好配方。尽管如此， MAE预训练还是有很大的改进，此处微调仅针对 50 个epoch（而从头开始 200 个），这意味着微调精度在很大程度上取决于预训练。

4.1 主要属性

掩蔽比
Figure 5 显示了掩蔽比的影响，最佳比率高得惊人。75% 的比率对线性探测和微调都有利。这种行为与 BERT 形成对比，后者的典型掩蔽率为 15%。掩蔽率也远高于计算机视觉相关工作中的掩蔽率（20% 到 50%）。该模型推断缺失的补丁以产生不同但合理的输出（Figure 4），它使物体和场景的格式塔变得有意义，这不能简单地通过延伸线条或纹理来完成。

假设这种类似推理的行为与学习有用的表征有关。Figure 5 还显示线性探测和微调结果遵循不同的趋势，对于线性探测，准确率随着掩蔽率稳步增加，直到最佳点：准确率差距高达 20%（54.6% 对 73.5%）。对于微调，结果对比率不那么敏感，并且很宽范围的掩蔽比率 (40-80%) 效果很好。Figure 5 中的所有微调结果都优于从头开始训练 (82.5%)。

解码器设计
MAE 解码器可以灵活设计，如Table 1a和 1b所示。Table 1a 改变了解码器深度（Transformer 块的数量）。足够深的解码器对于线性探测很重要。这可以通过像素重建任务和识别任务之间的差距来解释：自动编码器中的最后几层更专门用于重建，但与识别的相关性较低。一个合理深度的解码器可以解释重建专业化，将潜在表示留在更抽象的层次上。这种设计可使线性探测提高 8%（Table 1a，“lin”）。

但是，如果使用微调，则可以调整编码器的最后几层以适应识别任务。解码器深度对改进微调的影响较小（Table 1a，‘ft’）。有趣的是，带有单块解码器的 MAE 可以通过微调 (84.8%) 表现得很好。请注意，单个 Transformer 块是将信息从可见标记传播到掩码标记的最低要求。这么小的解码器可以进一步加快训练速度。在Table 1b中，研究了解码器宽度（通道数）。默认使用 512-d，它在微调和线性探测下表现良好。更窄的解码器也适用于微调。总的来说，默认的 MAE 解码器是轻量级的。它有 8 个块，宽度为 512-d（Table 1 中的灰色）。与 ViT-L（24 个区块，1024-d）相比，每个代币只有 9% 的 FLOP。因此，虽然解码器处理所有令牌，但它仍然是整体计算的一小部分。

掩码令牌
MAE 的一个重要设计是跳过编码器中的掩码标记 [M]，稍后将其应用到轻量级解码器中。Table 1c 研究了这种设计。

如果编码器使用掩码标记，它的性能会更差：在线性探测中其精度下降 14%。在这种情况下，预训练和部署之间存在差距：该编码器在预训练的输入中具有很大一部分掩码标记，这在未损坏的图像中不存在。这种差距可能会降低部署的准确性。通过从编码器中移除掩码标记，限制编码器始终看到真实的补丁，从而提高准确性。

此外，通过跳过编码器中的掩码标记，大大减少了训练计算。在Table 1c 中，将整体训练 FLOP 减少了 3.3 倍。这导致实现中 2.8 倍挂钟加速（见Table 2）。对于较小的解码器（1 块）、较大的编码器 (ViT-H) 或两者，挂钟加速甚至更大（3.5-4.1 倍）。请注意，对于 75% 的掩蔽率，加速可以大于 4 倍，部分原因是自注意力复杂度是二次的。此外，内存大大减少，可以训练更大的模型或通过大批量训练加快速度。时间和内存效率使 MAE 有利于训练非常大的模型。

重建目标
在Table 1d 中比较了不同的重建目标。到目前为止，结果基于没有（每个补丁）归一化的像素。使用归一化像素可提高准确性。这种逐块归一化在局部增强了对比度。在另一个变体中，在补丁空间中执行 PCA 并使用最大的 PCA 系数（此处为 96）作为目标。这样做会降低准确性。两个实验都表明高频分量在方法中很有用。

还比较了预测令牌的 MAE 变体，即 BEiT 中使用的目标。特别是对于这个变体，使用 DALLE 预训练的 dVAE 作为标记器，遵循。这里 MAE 解码器使用交叉熵损失预测令牌索引。这种标记化与非标准化像素相比将微调精度提高了 0.4%，但与标准化像素相比没有优势。它还降低了线性探测精度。在第 5 节中，进一步表明在转移学习中不需要标记化。

基于像素的 MAE 比标记化简单得多。dVAE 分词器需要一个额外的预训练阶段，这可能取决于额外的数据。dVAE 编码器是一个大型卷积网络（ViT-L 的 40% FLOPs）并增加了非平凡的开销。使用像素不会遇到这些问题。

数据增强
Table 1e 研究了数据增强对 MAE 预训练的影响。MAE 使用仅裁剪增强效果很好，无论是固定大小还是随机大小（两者都具有随机水平翻转）。添加颜色抖动会降低结果，因此不会在其他实验中使用它。

令人惊讶的是，即使没有使用数据增强（只有中心裁剪，没有翻转）， MAE 也表现得很好。这一特性与对比学习和相关方法截然不同，后者严重依赖于数据增强。据观察，对于 BYOL 和 SimCLR ，使用仅裁剪增强分别将准确度降低了 13% 和 28%。此外，没有证据表明对比学习可以在没有增强的情况下工作：图像的两个视图是相同的，可以很容易地满足一个简单的解决方案。在 MAE 中，数据增强的作用主要是通过随机掩码来实现的（下一个消融）。每次迭代的掩码都不同，因此无论数据增强如何，它们都会生成新的训练样本。掩蔽使借口任务变得困难，并且需要较少的增强来规范训练。

掩码采样策略
在Table 1f 中，比较了不同的掩码采样策略，如Figure 6 所示。

还研究了网格采样，它定期保留每四个补丁中的一个（Figure 6 右）。这是一项更简单的任务，并且具有更低的训练损失。重建更加清晰。但是，表示质量较低。

简单随机抽样最适合 MAE。它允许更高的掩蔽率，这提供了更大的加速优势，同时还享有良好的准确性。

训练进度
到目前为止，消融是基于 800 epoch 的预训练。Figure 7 显示了训练计划长度的影响。随着训练时间的延长，准确率稳步提高。事实上，即使在 1600 个时期，也没有观察到线性探测精度的饱和。这种行为与对比学习方法不同，例如，MoCo v3 在 ViT-L 的 300 个 epoch 时饱和。请注意，MAE 编码器每个 epoch 只能看到 25% 的补丁，而在对比学习中，编码器每个 epoch 看到 200%（twocrop）甚至更多（multi-crop）补丁。

4.2. 与之前结果的比较

1、与自监督方法的比较。

在Table 3 中，比较了自监督 ViT 模型的微调结果。对于 ViT-B，所有方法的表现都很接近。对于 ViT-L，方法之间的差距更大，这表明更大模型的挑战是减少过度拟合。

MAE 可以轻松扩展，并且从更大的模型中已经显示出稳定的改进。使用 ViT-H（224 尺寸）获得了 86.9% 的准确率。通过使用 448 尺寸进行微调，仅使用 IN1K 数据实现了 87.8% 的准确率。在仅使用 IN1K 数据的所有方法中，以前的最佳准确度为 87.1%（512 大小），基于高级网络。在竞争激烈的 IN1K 基准（无外部数据）中，以显着的优势改进了最先进的技术。结果基于 vanilla ViT，预计高级网络的性能会更好。

与 BEiT 相比， MAE 更准确，同时更简单、更快。与预测令牌的 BEiT 相比，方法重建像素：BEiT 报告在使用 ViT-B.2 重建像素时降低了 1.8% 不需要 dVAE 预训练。此外，由于Table 1c 中研究的原因， MAE 比 BEiT 快得多（每个 epoch 3.5 倍）。

Table 3 中的 MAE 模型经过 1600 个 epoch 的预训练以获得更好的准确性（Figure 7）。即便如此，如果它们在相同的硬件中进行训练，总预训练时间比所有其他方法都要少。例如，对于 ViT-L，使用相同的 128 个 TPU-v3 内核， MAE 的训练时间为 31 小时，1600 个时期，而 MoCo v3 的训练时间为 36 小时，300 个时期 。

2、与监督预训练的比较。

在最初的 ViT 论文 中，ViT-L 在 IN1K 中训练时性能下降。参见Figure 8。改进的监督配方更适合从头开始训练，但准确度已饱和。MAE 预训练，仅使用 IN1K，可以更好地泛化：对于更高容量的模型，从头开始训练的收益更大。它遵循类似于 中 JFT-300M 监督预训练的趋势。这个比较表明MAE 可以帮助扩大模型大小。

4.3、局部微调

Table 1 显示线性探测和微调结果在很大程度上不相关。线性探测在过去几年中一直是一种流行的协议；然而，它错过了追求强大但非线性特征的机会——这确实是深度学习的优势。作为中间立场，研究了部分微调协议：微调最后几层，同时冻结其他层。

Figure 9 显示了结果。值得注意的是，仅微调一个 Transformer 块即可将准确度从 73.5% 显着提高到 81.0%。此外，如果只微调最后一个块（即它的 MLP 子块）的“一半”，可以获得 79.1%，比线性探测要好得多。这种变体本质上是微调 MLP 头。微调几个块（例如 4 或 6 个）可以获得不错的精度，与冻结的主干相比，这仍然是一个很小的微调头。

在Figure 9 中，还与 MoCo v3 进行了比较，这是一种具有 ViT-L 结果的对比方法。它比 MAE 具有更高的线性探测精度。然而，它所有的局部微调结果都比差。调整 4 个块时差距为 2.6%。这些结果表明，MAE 表示的线性可分性较差，但它们具有更强的非线性特征，并且在调整非线性头时表现良好。

这些观察结果表明，线性可分性不是评估表征质量的唯一指标。还观察到线性探测与迁移学习性能（例如，用于对象检测）的相关性不是很好。在 NLP 中不经常使用线性评估来对预训练进行基准测试。

总结

作者从图像和语言的信号本质特征出发，考虑通过使用NLP领域中较为成熟的技术提高视觉自监督框架的性能，同时认真分析和处理二者之间的差距。

MAE 重建像素，它们不是语义实体。MAE 推断出复杂的整体重建，表明它已经学习了许多视觉概念，即语义。

通过最近的研究趋势，我们不得不感叹，CV和NLP融合的趋势越来越明显了。Transformer yyds！

版权声明：本文为奥比中光3D视觉开发者社区特约作者授权原创发布，未经授权不得转载，本文仅做学术分享，版权归原作者所有，若涉及侵权内容请联系删文。

3D视觉开发者社区是由奥比中光给所有开发者打造的分享与交流平台，旨在将3D视觉技术开放给开发者。平台为开发者提供3D视觉领域免费课程、奥比中光独家资源与专业技术支持。点击加入3D视觉开发者社区，和开发者们一起讨论分享吧~

或可微信关注官方公众号 3D视觉开发者社区 ，获取更多干货知识哦。