《MAE: Masked Autoencoders Are Scalable Vision Learners》论文精读笔记

前言

参考视频:MAE 论文逐段精读【论文精读】_哔哩哔哩_bilibili

论文链接:arxiv.org/pdf/2111.06377

Transformer：基于注意力机制的编码器和解码器，在机器翻译任务上相比于 RNN 架构表现更佳。
BERT：使用 Transformer 的编码器，通过自监督的完型填空训练机制，无需标注数据，扩展了 Transformer 的应用。
ViT：将 Transformer 应用于计算机视觉（CV），通过将图片切割成小块进行处理，证明了在大规模数据集上，Transformer 相对于 CNN 在精度上可能更高。
MAE: 可以视作 BERT 在计算机视觉领域的一个版本，基于 ViT，但将++训练拓展到无标注数据，采用完型填空方法来理解图像。++

论文标题: Masked Autoencoders Are Scalable Vision Learners

Scalable（可拓展的） ：指该学习器能够++处理大规模数据++。常见的论文中，若算法高效则会用"efficient"（高效）这个词，若算法处理大规模数据则会用"scalable"。
Vision Learner（视觉学习器） ：相对广泛的定义，没有特指某种分类器，表示这是一个++通用的模型架构++。
Masked（带掩码）：概念来源于BERT，表示通过挖去部分输入（掩码）以预测被掩盖的内容，类似于完型填空的方法。
AUTO : 之前讨论的Transformer和BERT没有使用"auto"前缀，而"auto "在此处指的是"自 "的含义，++强调输入和标号来自同一来源。++在机器学习中，自回归模型等属于自模型，表明在训练过程中，输入（x）和输出（y）之间存在直接联系。

摘要

这篇论文表明，带掩码的自编码器（MAE）是可扩展的自监督学习器，适用于计算机视觉。

我们的MAE方法很简单：随机遮住输入图像的一些块，并重构缺失的像素。

该方法基于两个核心设计:

首先，我们开发了一个非对称的编码器-解码器架构，编码器仅对可见的部分进行操作（不使用掩码标记），而轻量的解码器则从潜在表示和掩码标记中重构原始图像。

其次，我们发现遮住高比例的输入图像（例如75%）能产生一个非平凡且有意义的自监督任务。

这两个设计结合使我们能够有效且高效地训练大模型：我们加速了训练（提高了3倍以上）并提高了准确性。我们可扩展的方法允许学习高容量模型，具有良好的泛化能力：例如，一个普通的ViT-Huge模型在仅使用ImageNet-1K数据的情况下实现了最佳准确率（87.8%）。在下游任务中的迁移表现超越了监督预训练，并显示出良好的扩展性。

编码器只对可见的图像块进行编码，遮住的块不参与编码，遮住大量块后，模型的训练量减少至原来的1/4，因此训练速度加快了3倍或更多。

在最简单的ViT-Huge模型上，只使用ImageNet-1K（100万张图片）数据，准确率达到87.8%。该结果挑战了之前认为需要大量标注数据才能获得良好效果的观点，证明了自监督学习的潜力。

实验效果

下图是测试结果:

第一列是图片80%被遮掩的效果;

第二列表示MAE预测的图像效果;

第三列是imageNet的原始图像。

结论

简单且能够良好扩展的算法是深度学习的核心。在自然语言处理（NLP）领域，简单的自监督学习方法（例如 [47, 14, 48, 4]）使得模型能够从指数级扩展中获益。然而，在计算机视觉中，尽管自监督学习取得了一定进展，实际的预训练范式仍然以监督学习为主（例如 [33, 51, 25, 16]）。在本研究中，我们观察到在ImageNet和迁移学习中，自编码器（作为一种简单的自监督方法，类似于NLP中的技术）提供了可扩展的好处。计算机视觉中的自监督学习现在可能正走上一条与NLP相似的发展轨迹。

另一方面，我们注意到图像和语言是不同性质的信号，这一差异必须谨慎对待。图像仅仅是记录的光，没有像词汇那样的语义分解。我们不是试图去移除对象，而是随机去除那些最有可能不形成语义片段的补丁。同样，我们的MAE重构的是像素，而不是语义实体。尽管如此，我们观察到（例如，图4）我们的MAE能够推断出复杂的整体重构，表明它已学习到许多视觉概念，即语义。我们假设这种行为是通过MAE内部丰富的隐藏表示实现的。我们希望这一观点能够激励未来的研究。

引言

深度学习见证了不断增长的能力和容量的架构的爆炸性发展 [33, 25, 57]。借助硬件的快速进步，当前模型可以轻松过拟合一百万张图像 [13]，并开始需求数亿张（通常是公共不可访问的）标注图像 [16]。这种对数据的渴求在自然语言处理（NLP）中通过自监督预训练得到了成功的解决。基于GPT [47, 48, 4] 的自回归语言建模和BERT [14] 的掩码自编码技术的解决方案在概念上很简单：它们去除部分数据并学习预测被移除的内容。这些方法现已能够训练出包含超过一千亿参数的可推广NLP模型 [4]。

掩码自编码器的思想（作为更一般性去噪自编码器的一种形式 [58]）在计算机视觉中同样自然且适用。事实上，与BERT紧密相关的视觉研究 [59, 46] 早于BERT提出。然而，尽管在BERT成功后这一理念引起了极大关注，视觉中的自编码方法的进展仍滞后于NLP。

我们提出：掩码自编码在视觉和语言之间有什么不同？我们试图从以下几个角度来回答这个问题：

1. 架构的差异：直到最近，视觉架构与语言架构不同。在视觉领域，卷积网络 [34] 在过去十年中占主导地位 [33]。卷积通常在规则网格上操作，将"指示符"（如掩码符号 [14] 或位置嵌入 [57]）集成到卷积网络中并不简单。然而，这一架构差距已通过视觉Transformer（ViT） [16] 的引入得到解决，不再是障碍。

在Transformer中，mask是一个特定的词，会一直保留下来，跟别的能够做区分。

但在卷积网络中，把一个像素换成一个特定的值，在卷积窗口扫描的时候，没办法把这块替换的像素单独提取出来。导致掩码信息在后面比较难还原。

现在vit可以解决这个问题。

2. 信息密度的差异：语言是人类生成的信号，具有高度的语义性和信息密集性。当训练模型预测句子中少数缺失单词时，这一任务似乎会引发复杂的语言理解。相反，图像是具有严重空间冗余的自然信号------例如，缺失的补丁可以从相邻的补丁中恢复，几乎不需要对部分、物体和场景有高层次的理解。为克服这一差异以促使学习有用特征，我们展示了一种简单的策略在计算机视觉中表现良好：对大量随机补丁进行掩码。这一策略大幅降低冗余，并创造了一个需要超越低层次图像统计的整体理解的具有挑战性的自监督任务。有关我们重构任务的定性感受，请参见图2-4。

在自然语言里面一个词就是一个语义的实体，包含很多的信息，所以完形填空是一个比较难的任务。

在图片中，每个像素是冗余的，如果简单去除某一个像素，可以通过邻域像素进行插值还原出来。所以作者想把非常高比率的随机块都给去除掉，大大降低了图像的冗余性。使得模型去看一个全局信息，而不是关注学习局部。

如图2，3，4，在 ImageNet 和 COCO 数据集上,MAE 插值出来图片的效果确实可以通过一些局部的一些很稀疏的块得到一个全局的图片的重构。

3. 自编码器的解码器角色不同：自编码器的解码器将在隐含表示映射回输入的过程中，在文本和图像重构中扮演不同角色。在视觉中，解码器重构的是像素，因此其输出的语义水平低于常见的识别任务。这与语言形成对比，在语言中，解码器预测的缺失单词包含丰富的语义信息。虽然在BERT中，解码器可以是简单的多层感知器（MLP） [14]，但我们发现对于图像，解码器的设计在决定学习到的隐含表示的语义水平方面起着关键作用。

在NLP中，还原的是词，是高阶的语义信息，所以BERT使用了最简单的全连接层就能还原。

自编码器的解码器，在计算机视觉中，因为要还原的输入是原始的像素格，是一个低层次的表示。所以这里需要更大的一个解码器。

基于这项分析，我们提出了一种简单、有效且可扩展的掩码自编码器（MAE）用于视觉表征学习。

我们的MAE++随机掩码输入图像的补丁++，并在像素空间中重构缺失的补丁。其具有不对称的编码器-解码器设计。

我们的++编码器仅在可见的补丁子集上操作（不包含掩码符号）++ ，而解码器则轻量级，并根++据隐含表示及掩码符号重构输入（见图1）。将掩码符号转移到小型解码器中++，使得计算量大幅减少。

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

我们的MAE学习了非常高容量的模型，具有良好的泛化能力。通过MAE预训练，我们可以在ImageNet-1K上训练数据需求模型，如ViT-Large/-Huge [16]，并获得改进的泛化性能。使用普通的ViT-Huge模型，我们在ImageNet-1K的微调中达到了87.8%的准确率。这优于所有仅使用ImageNet-1K数据的先前结果。我们还评估了在目标检测、实例分割和语义分割任务上的迁移学习。在这些任务中，经过预训练的MAE比其监督预训练的对应模型取得了更好的结果，更重要的是，我们观察到通过扩大模型规模获得了显著的提升。这些观察与NLP中自监督预训练的结果一致 [14, 47, 48, 4]，我们希望这将使我们的领域探索类似的发展轨迹。

方法

我们的掩码自编码器（MAE）是一种简单的自编码方法，能够根据部分观测重构原始信号。与所有自编码器一样，我们的方法包含一个++将观测信号映射到隐含表示++ 的编码器 和一个++从隐含表示重构原始信号++ 的解码器。

与经典自编码器不同，我们采用了一种不对称的设计 ，使++编码器仅在部分观测信号（不包含掩码标记）上操作，而轻量级解码器则从隐含表示和掩码标记中重构完整信号。++图1展示了接下来要介绍的这一思想。

掩码

我们按照ViT [16] 的方法，将图像划分为规则的非重叠补丁。然后我们++抽样一部分补丁，并对剩余的补丁进行掩码（即移除）++。

我们的抽样策略相当简单：我们随机抽样补丁，且不放回，遵循均匀分布。我们称之为"随机抽样"。高掩码率的随机抽样 （即移除补丁的比例）++大大消除了冗余，从而创造了一个无法通过从可见邻近补丁进行外推轻松解决的任务++ （参见图2-4）。均匀分布 可以++防止潜在的中心偏倚（即更多被掩码的补丁集中在图像中心）++。最后，高度稀疏的输入为设计高效的编码器提供了机会，接下来将介绍。

MAE编码器

我们的编码器是一个ViT [16]，但仅应用于可见的、未掩码的补丁。与标准ViT一样，我们的编码器通过线性投影和添加位置嵌入来嵌入补丁，然后通过一系列Transformer块处理得到的集合。然而，我们的编码器仅在一个小的子集（例如，25%）上操作。被掩码的补丁被移除；不使用掩码标记。这使我们能够用更少的计算和内存训练非常大的编码器。完整的补丁集由一个轻量级的解码器处理，具体描述如下。

MAE解码器

MAE解码器的输入是由（i）编码的可见补丁和（ii）掩码标记组成的完整标记集（见图1）。

每个掩码标记 [14] 是一个共享的、学习得到的向量，表示要预测的缺失补丁的存在。我们为这个完整的标记集中的所有标记添加位置嵌入；没有位置嵌入，掩码标记将无法知道它们在图像中的位置。解码器有另一系列的Transformer块。

++MAE解码器仅在预训练期间用于执行图像重构任务（仅使用编码器生成图像表示用于识别）。++

因此，解码器的架构可以灵活设计，独立于编码器的设计。我们实验了非常小的解码器，其规模比编码器窄且浅。例如，我们的默认解码器的计算量小于编码器的10%。通过这种不对称的设计，只有轻量级的解码器处理完整的标记集，这显著减少了预训练时间。

重构目标

我们的MAE++通过预测每个被掩码补丁的像素值来重构输入++ 。解码器输出的每个元素 是一个代表补丁的像素值向量。

解码器的最后一层是一个线性投影 ，其++输出通道数等于补丁中像素值的数量++。

解码器的输出被重塑形成重构的图像。

我们的损失函数计算重构图像和原始图像在像素空间中的均方误差（MSE） 。我们++仅对被掩码的补丁计算损失++，类似于BERT [14]。

我们还研究了一种变体，其重构目标是每个被掩码补丁的归一化像素值 。具体来说，我们计算补丁中所有像素的均值和标准差，并用它们来归一化该补丁。在我们的实验中，使用归一化像素作为重构目标提高了表示质量。

简单实现

我们的MAE预训练可以高效实现，重要的是，它不需要任何专门的稀疏操作。

首先，我们为每个输入补丁生成一个标记（通过线性投影和添加位置嵌入）。
接下来，我们随机打乱标记列表，并根据掩码比例移除列表的最后部分。这个过程生成一个小的标记子集供编码器使用，相当于无放回地抽样补丁。
编码后，我们将掩码标记的列表附加到编码补丁列表中，并将其解除打乱（反转随机打乱操作），以将所有标记与其目标对齐。
解码器应用于这个完整列表（添加位置嵌入）。

如前所述，无需稀疏操作。这种简单的实现引入的开销微不足道，因为打乱和解除打乱操作非常快速。

ImageNet实验

我们在 ImageNet-1K (IN1K) [13] 训练集上进行自监督预训练。然后，我们进行监督训练，以评估表示，采用 (i) 端到端微调或 (ii) 线性探测。我们报告单个 224×224 裁剪的 top-1 验证准确率。详细信息见附录 A.1。

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

训练方式	准确率 (%)
从头训练	76.5
原始 [16] 从头训练	82.5
基线 MAE	84.9

我们注意到，从头训练监督 ViT-L 并非易事，需要良好的配方和强的正则化（82.5%，见附录 A.2）。即便如此，我们的 MAE 预训练仍然带来显著的提升。这里的微调仅进行了 50 个周期（而从头训练为 200 个周期），这意味着微调的准确率在很大程度上依赖于预训练。

主要特性

我们使用表 1 中的默认设置对 MAE 进行消融研究（见说明）。观察到几个有趣的特性。

掩码比例。图 5 显示了掩码比例的影响。最佳比例出乎意料地高。75%的掩码比例在线性探测和微调中表现良好。与 BERT [14] 的典型掩码比例 15% 相比，这一行为截然不同。我们的掩码比例也远高于计算机视觉相关工作中的比例 [6, 16, 2]（20% 至 50%）。该模型推断缺失的块以生成不同但合理的输出（图 4）。它能够理解对象和场景的整体特征，这不能仅通过延伸线条或纹理来完成。我们假设这种类推行为与学习有用的表示相关。

图 5 还显示，线性探测和微调结果的趋势不同。对于线性探测，准确率随掩码比例稳步增加，直到甜点期：准确率差距高达约 20%（54.6% 对比 73.5%）。对于++微调，结果对掩码比例的敏感性较低++ ，40% 至 80% 的广泛掩码比例表现良好。图 5 中的++所有微调结果均优于从头训练++（82.5%）。

微调：这里的微调是end-to-end，整个模型的参数都是可学习的。

linear probing:改最后一层的线性输出层。

解码器设计。我们的 MAE 解码器可以灵活设计，如表 1a 和 1b 所示。

表 1a 变化了解码器的深度（Transformer 块的数量）。对于++线性探测来说，足够深的解码器是重要的++ 。这可以通过像素重构任务与识别任务之间的差距来解释：自编码器的最后几层更专注于重构，但与识别的相关性较低。一个合理深度的解码器可以考虑重构专门化，使潜在表示处于更抽象的层次。这种设计在线性探测中可以获得高达 8% 的提升（表 1a，'lin'）。然而，如果使用微调，编码器的最后几层可以调整以适应识别任务。++解码器深度对微调的提升影响较小++（表 1a，'ft'）。

有趣的是，我们的 MAE 在微调时使用单块解码器仍能表现良好（84.8%）。注意，单个 Transformer 块是将信息从可见标记传播到掩码标记的最低要求。这样的一个小解码器可以进一步加速训练。在表 1b 中，我们研究了解码器的宽度（通道数量）。我们默认使用 512-d，在线性探测和微调下表现良好。较窄的解码器在微调中也能良好工作。总体而言，我们的默认 MAE 解码器是轻量级的。它有 8 个块和 512-d 的宽度（表 1 中的灰色）。与 ViT-L (24 块，1024-d) 相比，每个标记的 FLOPs 仅为 9%。因此，虽然解码器处理所有标记，但仍然是整体计算的一小部分。

掩码标记 。我们 MAE 的一个重要设计是跳过编码器中的掩码标记 [M]，并在轻量级解码器中应用它。表 1c 研究了这一设计。++如果编码器使用掩码标记，其性能更差++ ：在线性探测中准确率下降了 14%。在这种情况下，预训练与部署存在差距：这个编码器在预训练时输入中有大量掩码标记，而在未损坏的图像中并不存在。这种差距可能会降低部署中的准确性。通过从编码器中移除掩码标记，我们强制编码器始终看到真实的块，从而提高准确性。此外，++通过在编码器中跳过掩码标记，我们大大减少了训练计算。++在表 1c 中，我们将整体训练 FLOPs 降低了 3.3 倍。这在我们的实现中导致了 2.8 倍的壁钟加速（见表 2）。对于较小的解码器（1 块）、较大的编码器（ViT-H）或两者的组合，壁钟加速甚至更大（3.5--4.1 倍）。注意，对于 75% 的掩码比例，速度提升可能超过 4 倍，这部分是因为自注意力复杂度是平方关系。此外，显著减少的内存使得训练更大的模型成为可能，或通过大批量训练进一步加速。时间和内存效率使得我们的 MAE 在训练非常大的模型时具有优势。

重构目标 。我们在表 1d 中比较了不同的重构目标。我们迄今的结果基于未进行（每块）归一化的像素。++使用归一化的像素可以提高准确性++。这种每块归一化增强了局部对比度。在另一个变体中，我们在块空间中执行 PCA，并使用最大 PCA 系数（这里为 96）作为目标。这样做降低了准确性。这两个实验表明，高频成分在我们的方法中是有用的。我们还比较了一个 MAE 变体，该变体预测标记，这是 BEiT [2] 中使用的目标。特别地，对于这个变体，我们使用 DALLE 预训练的 dVAE [50] 作为分词器，遵循 [2]。在这种情况下，MAE 解码器使用交叉熵损失预测标记索引。与未归一化的像素相比，这种标记化提高了微调准确性 0.4%，但与归一化像素没有优势。它还降低了线性探测的准确性。在 §5 中我们进一步表明，标记化在迁移学习中并不是必需的。我们的基于像素的 MAE 比标记化要简单得多。dVAE 分词器需要额外的预训练阶段，这可能依赖于额外的数据（250M 图像 [50]）。dVAE 编码器是一个大型卷积网络（占 ViT-L FLOPs 的 40%），并增加了非平凡的开销。使用像素不会遭遇这些问题。

数据增强。表 1e 研究了数据增强对我们的 MAE 预训练的影响。我们的 MAE 在使用仅裁剪增强时表现良好，无论是固定大小还是随机大小（两者都有随机水平翻转）。添加颜色抖动会降低结果，因此我们在其他实验中不使用。令人惊讶的是，即使不使用数据增强（仅中心裁剪，无翻转），我们的 MAE 仍然表现良好。这一特性与对比学习和相关方法 [62, 23, 7, 21] 有显著不同，后者严重依赖于数据增强。观察到 [21]，使用仅裁剪增强分别降低了 BYOL [21] 和 SimCLR [7] 的准确性 13% 和 28%。此外，没有证据表明对比学习在没有增强的情况下能正常工作：同一图像的两个视图是相同的，容易满足一个平凡的解决方案。在 MAE 中，数据增强的角色主要由随机掩码执行（下一个部分会消融）。掩码在每次迭代中都是不同的，因此无论数据增强如何，都生成新的训练样本。由于掩码的存在，预训练任务变得更为困难，因此需要更少的增强来正则化训练。

掩码采样策略 。在表 1f 中，我们比较了不同的掩码采样策略，如图 6 所示。++块状掩码策略++ ，最早在 [2] 中提出，倾向于移除大块（图 6 中间）。我们的 MAE 在 50% 的比例下使用块状掩码表现合理，但在 75% 的比例下则有所下降。这一任务比随机采样更困难，因为观察到更高的训练损失。重构质量也更模糊。我们还研究了++网格状采样++ ，该策略定期保留每四个块中的一个（图 6 右）。这是一项更简单的任务，具有较低的训练损失，++重构更清晰。然而，表示质量较低++ 。++简单的随机采样对我们的 MAE 效果最好++ 。它允许更高的掩码比例，提供更大的加速收益，同时也享有良好的准确性。

训练计划。我们迄今的消融基于 800 轮的预训练。图 7 显示了训练计划长度的影响。准确性随着训练时间的延长而稳步提高。实际上，我们甚至在 1600 轮时都没有观察到线性探测准确性的饱和。这种行为与对比学习方法不同，例如 MoCo v3 [9] 在 ViT-L 上的饱和点为 300 轮。请注意，MAE 编码器每轮仅看到 25% 的块，而在对比学习中，编码器每轮看到 200%（双裁剪）甚至更多（多裁剪）的块。

与之前结果的比较

与自监督方法的比较。在表 3 中，我们比较了自监督 ViT 模型的微调结果。对于 ViT-B，所有方法表现相近。对于 ViT-L，各方法之间的差距更大，表明对于更大模型的挑战在于减少过拟合。我们的 MAE 可以轻松扩展，并且从更大的模型中显示出稳定的改进。我们在使用 ViT-H（224 大小）时获得了 86.9% 的准确率。通过使用 448 大小进行微调，我们达到了 87.8% 的准确率，仅使用 IN1K 数据。在所有仅使用 IN1K 数据的方法中，之前的最佳准确率为 87.1%（512 大小） [67]，基于先进网络。我们在 IN1K 这一高度竞争基准中非平凡地提高了状态的表现（无外部数据）。我们的结果基于原始 ViT，我们预计先进网络的性能会更好。与 BEiT [2] 的比较中，我们的 MAE 更准确，同时更简单且更快。我们的方法重构像素，而 BEiT 则预测标记：BEiT 报告在用 ViT-B 重构像素时准确性下降 1.8% [2]。我们不需要 dVAE 预训练。此外，我们的 MAE 在每轮训练中相较于 BEiT 速度更快（3.5 倍），如表 1c 所述。表 3 中的 MAE 模型预训练了 1600 轮以获得更好的准确性（图 7）。即便如此，在相同硬件上训练时，我们的总预训练时间仍低于其他方法。例如，在 128 TPU-v3 核心上训练 ViT-L，我们的 MAE 训练时间为 31 小时（1600 轮），而 MoCo v3 的训练时间为 36 小时（300 轮） [9]。

与监督预训练的比较。在原始 ViT 论文 [16] 中，ViT-L 在 IN1K 上训练时效果下降。我们对监督训练的实现（见 A.2）效果更好，但准确率饱和。见图 8。我们的 MAE 预训练，使用仅 IN1K，能够更好地进行泛化：与从头训练相比，高容量模型的提升更大。它遵循类似于 [16] 中 JFT-300M 监督预训练的趋势。这个比较表明，我们的 MAE 能够帮助扩大模型规模。