Masked Autoencoders Are Scalable Vision Learners
郑重声明:原文参见标题,如有侵权,请联系作者,将会撤销发布!
CVPR 2022
Abstract
本文表明,掩码自编码器(masked autoencoders, MAE)是用于计算机视觉的可扩展自监督学习器。我们的MAE方法很简单:我们屏蔽输入图像的随机补丁并重建丢失的像素。它基于两个核心设计。首先,我们开发了一种非对称的编码器-解码器架构,其中一个编码器仅对可见的补丁(patch)子集进行操作(没有掩码标记(token)),另一个轻量级解码器从潜在表示(latent representation)和掩码标记重建原始图像。其次,我们发现掩码高比例的输入图像,例如75%,产生了一个非平凡且有意义的自监督任务。将这两种设计结合起来,使我们能够高效地训练大型模型:我们加快训练速度(3倍或更多)并提高精度。我们的可扩展方法允许学习具有良好泛化能力的高容量模型:例如,在仅使用ImageNet-1K数据的方法中,原始ViT-Huge模型的准确率最高(87.8%)。下游任务中的转移性能优于监督预训练,并显示出良好的缩放行为。
1. Introduction
深度学习见证了能力和容量不断增长的架构的爆炸[33, 25, 57]。在硬件快速增长的帮助下,如今的模型可以很容易地过度填充一百万张图像[13],并开始需要数亿张经常无法公开获取的标记图像[16]。
这种对数据的需求已经在自然语言处理(NLP)中通过自监督的预训练成功解决。基于GPT中的自回归语言建模[47, 48, 4]和BERT中的掩码自编码[14]的解决方案在概念上很简单:它们删除了部分数据,并学习预测删除的内容。这些方法现在能够训练包含超过1000亿个参数的可推广NLP模型[4]。
掩码自编码器的概念是一种更通用的去噪自编码器[58],它是自然的,也适用于计算机视觉。事实上,在BERT之前,视觉方面的密切相关研究[59, 46]。然而,尽管随着BERT的成功,人们对这一想法产生了极大的兴趣,但视觉自编码方法的进展却落后于NLP。我们问:是什么使视觉和语言之间的掩码自编码不同?我们试图从以下角度回答这个问题:
(i)直到最近,架构都有所不同。在视觉上,卷积网络[34]在过去十年中占据主导地位[33]。卷积通常在规则网格上运行,将掩码标记[14]或位置嵌入[57]等“指示符”集成到卷积网络中并不简单。然而,通过引入视觉Transformer(ViT)[16],解决了这种架构上的差距,不应再成为障碍。
(ii)语言和视觉的信息密度不同。语言是人类产生的高度语义和信息密集的信号。当训练一个模型来预测每一句话只会漏掉几个单词时,这项任务似乎能诱导人们对语言的理解。相反,图像是具有严重空间冗余的自然信号,例如,可以从相邻的补丁中恢复丢失的补丁,而对部分、目标和场景的高级理解很少。为了克服这一差异并鼓励学习有用的特征,我们证明了一种简单的策略在计算机视觉中很有效:掩码很高比例的随机补丁。该策略在很大程度上减少了冗余,并创建了一个具有挑战性的自监督任务,需要对低级别图像统计之外的整体理解。要对重建任务有一个定性的了解,请参见图2-4。
(iii)自编码器的解码器将潜在表示映射回输入,在重建文本和图像之间扮演不同的角色。在视觉中,解码器重建像素,因此其输出的语义级别低于普通识别任务。这与语言不同,在语言中,解码器预测包含丰富语义信息的缺失单词。尽管在BERT中,解码器可能是微不足道的(MLP)[14],但我们发现,对于图像,解码器设计在确定学习的潜在表示的语义水平方面起着关键作用。
在这种分析的驱动下,我们提出了一种简单有效且可扩展的掩码自编码器(MAE)形式,用于视觉表征学习。我们的MAE从输入图像中掩码随机补丁,并在像素空间中重建缺失的补丁。它具有非对称编码器设计。我们的编码器仅对可见的补丁子集(没有掩码标记)进行操作,而我们的解码器是轻量级的,并从潜在表示和掩码标记中重建输入(图1)。在我们的非对称编码器-解码器中,将掩码标记转移到小解码器会导致计算的大幅减少。在这种设计下,非常高的掩码率(例如,75%)可以实现双赢:它优化了精度,同时允许编码器仅处理一小部分(例如,25%)的补丁。这可以将整个预训练时间减少3倍或更多,并同样减少内存消耗,使我们能够轻松地将MAE扩展到大型模型。
我们的MAE学习的高容量模型具有很好的通用性。通过MAE预训练,我们可以在ImageNet-1K上训练像ViT Large/-Huge[16]这样的数据饥饿模型,并提高泛化性能。使用朴素ViT-Huge模型,在ImageNet-1K上进行微调时,我们获得了87.8%的准确率。这优于仅使用ImageNet-1K数据的所有先前结果。我们还评估了目标检测、实例分割和语义分割方面的迁移学习。在这些任务中,我们的预训练取得了比受监督的预训练对应工作更好的结果,更重要的是,我们观察到通过扩大模型而获得的显著收益。这些观察结果与NLP自监督预训练中的观察结果一致[14, 47, 48, 4],我们希望这些观察结果将使我们的领域能够探索类似的轨迹。
2. Related Work
掩码语言建模(Masked language modeling)和其自回归对应物,例如BERT[14]和GPT[47, 48, 4],是NLP预训练的非常成功的方法。这些方法保留了输入序列的一部分,并训练模型来预测缺失的内容。这些方法已被证明具有很好的规模[4],大量证据表明,这些预先训练的表示可以很好地推广到各种下游任务。
自编码(Autoencoding)是学习表示的经典方法。它有一个将输入映射到潜在表示的编码器和一个重构输入的解码器。例如,PCA和k-means是自编码器[29]。去噪自编码器(denoising autoencoder, DAE)[58]是一类损坏输入信号并学习重建原始未损坏信号的自编码器。一系列方法可以被认为是不同破坏下的通用DAE,例如,掩码像素[59, 46, 6]或移除颜色通道[70]。我们的MAE是一种去噪自编码形式,但在许多方面不同于经典的DAE。
掩码图像编码(Masked image encoding)方法从被掩码破坏的图像中学习表示。[59]的开创性工作将掩码作为DAE中的一种噪声类型。上下文编码器(Context Encoder)[46]使用卷积网络修复大的缺失区域。受NLP成功的激励,最近的相关方法[6, 16, 2]基于Transformer[57]。iGPT[6]对像素序列进行操作,并预测未知像素。ViT论文[16]研究了用于自监督学习的掩码补丁预测。最近,BEiT[2]提出预测离散标记[44, 50]。
自监督学习(Self-supervised learning)方法对计算机视觉产生了极大的兴趣,通常集中在不同的预训练伪装(pretext)任务上[15, 61, 42, 70, 45, 17]。最近,对比学习[3, 22]已经很流行,例如[62, 43, 23, 7],它对两个或多个视图之间的图像相似性和不相似性(或仅相似性[21, 8])进行建模。对比和相关方法强烈依赖于数据扩充[7, 21, 8]。自编码追求一个概念上不同的方向,它表现出我们将呈现的不同行为。
3. Approach
我们的掩码自编码器(MAE)是一种简单的自编码方法,它在给定部分观测的情况下重建原始信号。像所有的自编码器一样,我们的方法有一个将观察到的信号映射到潜在表示的编码器,以及一个从潜在表示重构原始信号的解码器。与经典的自编码器不同,我们采用了一种非对称设计,允许编码器仅对部分观察到的信号进行操作(无掩码标记),并采用了一个轻量级解码器,从潜在表示和掩码标记重建完整信号。图1说明了接下来介绍的想法。
Masking. 根据ViT[16],我们将图像划分为规则的非重叠补丁。然后我们对一个子集的补丁进行采样,并掩码(即移除)其余的补丁。我们的采样策略很简单:我们按照均匀分布,在没有替换的情况下对随机补丁进行采样。我们简单地称之为“随机采样”(random sampling)。
具有高掩码率(即,去除补丁的比率)的随机采样在很大程度上消除了冗余,从而创建了一个无法通过从可见相邻补丁外推(extrapolation)来轻松解决的任务(见图2-4)。均匀分布防止了潜在的中心偏移(即,图像中心附近的更多掩码补丁)。最后,高度稀疏的输入为设计一个高效的编码器创造了机会,下面将介绍。
MAE encoder. 我们的编码器是一个ViT[16],但仅适用于可见且无掩码的补丁。就像在标准ViT中一样,我们的编码器通过添加位置嵌入的线性投影来嵌入补丁,然后通过一系列Transformer块来处理结果集。然而,我们的编码器仅对整个集合的一小部分(例如,25%)进行操作。移除掩码补丁;不使用掩码标记。这使我们能够训练非常大的编码器,只需很少的计算和内存。整个集合由一个轻量级解码器处理,如下所述。
MAE decoder. MAE解码器的输入是由(i)编码的可见补丁和(ii)掩码标记组成的全组标记。参见图1。每个掩码标记[14]都是一个学到的共享向量,指示要预测的缺失补丁的存在。我们将位置嵌入添加到此完整集合中的所有标记;没有这一点,掩码标记将没有关于它们在图像中的位置的信息。解码器具有另一系列的Transformer块。
MAE解码器仅在预训练期间用于执行图像重建任务(仅编码器用于产生用作识别的图像表示)。因此,能够以独立于编码器设计的方式灵活地设计解码器架构。我们实验了非常小的解码器,比编码器更窄、更浅。例如,我们的默认解码器与编码器相比,每个标记的计算量小于10%。采用这种不对称设计,全套标记仅由轻量级解码器处理,这大大减少了预训练时间。
Reconstruction target. 我们的MAE通过预测每个掩码补丁的像素值来重建输入。解码器输出中的每个元素都是表示补丁的像素值向量。解码器的最后一层是线性投影,其输出通道的数量等于单个补丁中像素值的数量。解码器的输出被调整成重构图像。我们的损失函数计算像素空间中重建图像和原始图像之间的均方误差(MSE)。我们只计算掩码补丁上的损失,类似于BERT[14]。1
我们还研究了一个变体,其重建目标是每个掩码补丁的归一化像素值。具体来说,我们计算一个补丁中所有像素的平均值和标准差,并使用它们来归一化这个补丁。在我们的实验中,使用归一化像素作为重建目标提高了表示质量。
Simple implementation. 我们的MAE预训练可以有效实现,而且重要的是,不需要任何专门的稀疏操作。首先,我们为每个输入补丁生成一个标记(通过添加位置嵌入的线性投影)。接下来,我们随机混洗(randomly shuffle)标记列表,并根据掩码率移除列表的最后一部分。此过程为编码器生成一小部分标记,相当于在不替换的情况下对补丁进行采样。编码后,我们将掩码标记列表附加到已编码补丁列表中,并取消混洗(unshuffle)该完整列表(反转随机洗牌操作),以将所有标记与其目标对齐。解码器应用于此完整列表(添加了位置嵌入)。如前所述,不需要稀疏操作。这种简单的实现引入了可忽略的开销,因为混洗和取消混洗操作很快。
1 仅在掩码补丁上计算损失与计算所有像素损失的传统去噪自编码器[58]不同。这种选择纯粹是结果驱动的:计算所有像素的损失会导致精度略有下降(例如,~0.5%)。
4. ImageNet Experiments
我们在ImageNet-1K(IN1K)[13]训练集上进行自监督预训练。然后我们进行监督训练,以(i)端到端微调(end-to-end fine-tuning)或(ii)线性探测(linear probing)来评估表示。我们报告了单个224×224裁剪的最高验证精度。详情见附录A.1。
Baseline: ViT-Large. 我们使用ViT Large(ViT-L/16)[16]作为消融研究的骨干。ViT-L非常大(比ResNet-50[25]大一个数量级),并且倾向于过拟合。以下是从零开始训练的ViT-L与从我们的基准MAE微调的ViT-L之间的比较:
我们注意到,从头开始训练受监督的ViT-L是非常重要的,需要一个具有强大正则化的良好配置(82.5%,见附录A.2)。即使如此,我们的MAE预训练也有很大的改进。这里微调仅适用于50个回合(而不是从零开始的200个回合),这意味着微调精度严重依赖于预训练。
4.1. Main Properties
我们使用表1中的默认设置消融MAE(见标题)。观察到了几个有趣的特性。
Masking ratio. 图5显示了掩码率的影响。最佳比率高得惊人。75%的比率对于线性探测和微调都很好。这种行为与BERT[14]形成对比,BERT的典型掩码率为15%。我们的掩码率也远高于计算机视觉相关研究[6, 16, 2]中的掩码率(20%至50%)。
该模型推断出缺失的补丁,以产生不同但似乎合理的输出(图4)。它能理解目标和场景的格式塔(gestalt),而这不能简单地通过延伸线条或纹理来完成。我们假设这种类似推理的行为与学习有用的表示有关。
图5还显示了线性探测和微调结果遵循不同的趋势。对于线性探测,准确度随着掩码率的增加而稳步增加,直到达到最佳点:准确度差距高达20%(54.6% vs. 73.5%)。对于微调,结果对比率不太敏感,并且范围广泛的掩码率(40-80%)工作良好。图5中的所有微调结果都优于从头开始的训练(82.5%)。
Decoder design. 我们的MAE解码器可以灵活设计,如表1a和1b所示。
表1a改变了解码器深度(Transformer块的数量)。足够深的解码器对于线性探测很重要。这可以通过像素重建任务和识别任务之间的差距来解释:自编码器中的最后几层更专门用于重建,但与识别的相关性较小。一个合理深度的解码器可以解释重建的专门化,将潜在的表示留在更抽象的层次。这种设计可使线性探测提高8%(表1a,“lin”)。然而,如果使用微调,则可以调整编码器的最后一层以适应识别任务。解码器深度对改善微调影响较小(表1a,“ft”)。
有趣的是,我们的单块解码器MAE可以在微调的情况下表现出色(84.8%)。注意,单个Transformer块是将信息从可见标记传播到掩码标记的最低要求。这样的小解码器可以进一步加快训练速度。
在表1b中,我们研究了解码器宽度(信道数)。我们默认使用512维,它在微调和线性探测下表现良好。更窄的解码器也可以很好地进行微调。
总体而言,我们的默认MAE解码器是轻量级的。它有8个块,宽度为512维(表1中的灰色)。与ViT-L相比,每个标记只有9%的FLOP(24块,1024维)。因此,当解码器处理所有标记时,它仍然是整个计算的一小部分。
Mask token. 我们的MAE的一个重要设计是跳过编码器中的掩码标记[M],然后将其应用于轻量级解码器。表1c研究了该设计。
如果编码器使用掩码标记,它的性能会更差:在线性探测中,它的精度会下降14%。在这种情况下,预训练和部署之间存在差距:该编码器在预训练中的输入中有很大一部分掩码标记,这在未损坏的图像中不存在。此间隙可能会降低部署的准确性。通过从编码器中移除掩码标记,我们约束编码器始终看到真实的补丁,从而提高准确性。
此外,通过在编码器中跳过掩码标记,我们大大减少了训练计算。在表1c中,我们将总体训练FLOP减少了3.3倍。这导致在我们的实现中,墙时钟速度提高了2.8倍(见表2)。对于较小的解码器(1块)、较大的编码器(ViT-H)或两者,墙时钟的加速比甚至更大(3.5–4.1倍)。注意,对于75%的掩码率,加速比可以大于4倍,部分原因是自注意力复杂度是二次的。此外,内存大大减少,可以训练更大的模型或通过大批量训练加快速度。时间和记忆效率使我们的MAE有利于训练非常大的模型。
Reconstruction target. 我们比较了表1d中的不同重建目标。到目前为止,我们的结果是基于没有(逐补丁)归一化的像素。使用归一化的像素可以提高精度。这种逐补丁归一化增强了局部对比度。在另一个变体中,我们在补丁空间中执行PCA,并使用最大的PCA系数(此处为96)作为目标。这样做会降低准确性。两个实验都表明高频分量在我们的方法中是有用的。
我们还比较了预测标记的MAE变体,即BEiT[2]中使用的目标。特别是对于该变体,我们使用DALLE预训练dVAE[50]作为标记器,如[2]。这里,MAE解码器使用交叉熵损失来预测标记索引。与未归一化的像素相比,这种标记化将微调精度提高了0.4%,但与归一化像素相比没有优势。它还降低了线性探测精度。在§5中,我们进一步表明,标记化在迁移学习中是不必要的。
我们基于像素的MAE比标记化简单得多。dVAE标记器需要一个额外的预训练阶段,这可能取决于额外的数据(250M图像[50])。dVAE编码器是一个大型卷积网络(ViT-L的40%FLOP),并增加了非平凡的开销。使用像素不会出现这些问题。
Data augmentation. 表1e研究了数据增强对MAE预训练的影响。
我们的MAE使用仅裁剪的增强效果很好,无论是固定大小还是随机大小(都具有随机水平翻转)。添加颜色抖动会降低结果,因此我们不会在其他实验中使用它。
令人惊讶的是,即使没有使用数据扩充(只有中心裁剪,没有翻转),我们的MAE也表现得很好。这一特性与对比学习和相关方法[62, 23, 7, 21]显著不同,后者严重依赖数据增强。据观察[21],对于BYOL[21]和SimCLR[7],仅使用裁剪增强将准确性分别降低13%和28%。此外,没有证据表明对比学习可以在没有增强的情况下工作:图像的两个视图是相同的,可以很容易地满足一个简单的解决方案。
在MAE中,数据增强的作用主要通过随机掩码(接下来将被消除)来实现。每个迭代的掩码都是不同的,因此它们生成新的训练样本,而不考虑数据扩充。伪装任务很难完成,需要较少的强化来规范训练。
Mask sampling strategy. 在表1f中,我们比较了不同的掩码采样策略,如图6所示。
[2]中提出的逐块掩码策略倾向于去除大块(图6中间)。我们的分块掩码MAE在50%的比率下工作得相当好,但在75%的比率下退化。这项任务比随机采样更难,因为观察到了更高的训练损失。重建也更加模糊。
我们还研究了网格采样,它定期保持每四个补丁中的一个(图6右侧)。这是一项更容易的任务,训练损失也更低。重建更为清晰。然而,表示质量较低。
简单的随机采样对我们的MAE最有效。它允许更高的掩码率,这提供了更大的加速优势,同时也享受到良好的精度。
Training schedule. 到目前为止,我们的消融是基于800小时的预训练。图7显示了训练规划长度的影响。随着训练时间的延长,准确度稳步提高。事实上,即使在1600个回合,我们也没有观察到线性探测精度的饱和。这种行为与对比学习方法不同,例如,MoCo v3[9]在ViT-L的300个回合饱和。注意,MAE编码器每个回合仅看到25%的补丁,而在对比学习中,编码器在每个回合看到200%(双裁剪)或甚至更多(多裁剪)的补丁。
4.2. Comparisons with Previous Results
Comparisons with self-supervised methods. 在表3中,我们比较了自监督ViT模型的微调结果。对于ViT-B,所有方法都表现得很好。对于ViT-L,方法之间的差距更大,这意味着更大的模型面临的挑战是减少过拟合。
我们的MAE可以很容易地扩大规模,并从更大的模型中显示出稳定的改进。我们使用ViT-H(224尺寸)获得了86.9%的准确度。通过对448尺寸的微调,我们仅使用IN1K数据就实现了87.8%的准确率。基于先进网络,在所有仅使用IN1K数据的方法中,先前的最优精度为87.1%(512尺寸)[67]。我们在IN1K(无外部数据)这一极具竞争力的基准中,以非凡的优势超越了最新技术。我们的结果基于普通ViT,我们希望高级网络的性能会更好。
与BEiT[2]相比,我们的MAE更准确,同时更简单、更快。我们的方法重建像素,与预测标记的BEiT相反:当使用ViT-B重建像素时,BEiT报告了1.8%的退化[2]。2 我们不需要dVAE预训练。此外,由于表1c中所研究的原因,我们的MAE比BEiT快得多(每个回合3.5倍)。
表3中的MAE模型经过1600个回合的预训练,以获得更好的精度(图7)。即使如此,当在相同的硬件上进行训练时,我们的预训练总时间比其他方法短。例如,在128个TPU-v3核上训练ViT-L,我们的MAE训练时间为31小时(1600个回合),MoCo v3训练时间为36小时(300个回合)[9]。
2 我们也观察到了BEiT和ViT-L的退化:它产生85.2%(标记)和83.5%(像素),从官方代码中复制。
Comparisons with supervised pre-training. 在最初的ViT论文[16]中,当在IN1K中训练时,ViT-L会退化。我们实现的监督训练(见A.2)效果更好,但准确性饱和。见图8。
我们的MAE预训练(仅使用IN1K)可以更好地泛化:对于高容量模型,从头开始训练的收益更大。它遵循类似于[16]中JFT-300M监督预训练的趋势。这一比较表明,我们的MAE可以帮助扩大模型尺寸。
4.3. Partial Finetuning
表1显示了线性探测和微调结果在很大程度上是不相关的。在过去几年中,线性探测一直是一种流行的规约(protocol);然而,它错过了追求强大但非线性特征的机会,这确实是深度学习的优势。作为中间地带,我们研究了一种部分微调规约:微调最后几层,同时冻结其他层。该方案也用于早期工作,例如[65, 70, 42]。
图9显示了结果。值得注意的是,仅对一个Transformer块进行微调,将准确率从73.5%显著提高到81.0%。此外,如果我们只微调最后一个块的“一半”(即其MLP子块),我们可以获得79.1%,比线性探测要好得多。该变体本质上是对MLP头进行微调。微调几个块(例如,4或6)可以实现接近完全微调的精度。
在图9中,我们还与MoCo v3[9]进行了比较,这是一种与ViT-L结果的对比方法。MoCo v3具有更高的线性探测精度;然而,它的所有部分微调结果都比MAE差。调整4个块时,差距为2.6%。虽然MAE表示的线性可分离性较低,但它们是更强的非线性特征,在调整非线性头时表现良好。
这些观察结果表明,线性可分性不是评估表示质量的唯一指标。还观察到(例如,[8]),线性探测与迁移学习性能(例如,对象检测)没有很好的关联。据我们所知,NLP中不经常使用线性评估作为基准预训练。
5. Transfer Learning Experiments
我们使用表3中的预训练模型评估下游任务中的迁移学习。
Object detection and segmentation. 我们在COCO[37]上端对端微调Mask R-CNN[24]。ViT主干适用于FPN[36](见A.3)。我们将此方法应用于表4中的所有条目。我们报告用于目标检测的框AP和用于实例分割的掩码AP。
与监督预训练相比,我们的MAE在所有配置下都表现得更好(表4)。使用较小的ViT-B,我们的MAE比监督预训练高2.4分(50.3 vs. 47.9,APbox)。更重要的是,使用更大的ViT-L,我们的MAE预训练优于监督预训练4.0分(53.3 vs. 49.3)。
基于像素的MAE优于或等同于基于标记的BEiT,而MAE更简单、更快。MAE和BEiT都优于MoCo v3,MoCo v3与监督预训练不相上下。
Semantic segmentation. 我们使用UperNet[63]对ADE20K[72]进行了实验(见A.4)。表5显示,我们的预训练比监督预训练显著提高了结果,例如ViT-L提高了3.7分。我们基于像素的MAE也优于基于标记的BEiT。这些观察结果与COCO中的一致。
Classification tasks. 表6研究了iNaturalists[56]和Places[71]任务的迁移学习(见A.5)。在iNat上,我们的方法显示出强大的缩放行为:使用更大的模型,精度会显著提高。我们的结果大大超过了以前的最优结果。在Places上,我们的MAE优于之前的最佳结果[19, 40],这些结果是通过对数十亿张图像进行预训练获得的。
Pixels vs. tokens. 表7比较了像素与标记作为MAE重建目标。虽然使用dVAE标记比使用非标准化像素更好,但在统计上,它与在我们测试的所有情况下使用归一化像素相似。这再次表明,标记化对于我们的MAE来说是不必要的。
6. Discussion and Conclusion
扩展性好的简单算法是深度学习的核心。在NLP中,简单的自监督学习方法(例如[47, 14, 48, 4])能够从指数缩放模型中获益。在计算机视觉中,尽管在自监督学习方面取得了进展,但实际的预训练范式主要受到监督(例如[33, 51, 25, 16])。在这项研究中,我们在ImageNet和迁移学习中观察到,自编码器——一种类似于NLP技术的简单自监督方法——提供了可扩展的好处。视觉中的自监督学习现在可能正走上与NLP类似的轨道。
另一方面,我们注意到图像和语言是不同性质的信号,必须仔细处理这种差异。图像仅仅是记录下来的光,没有语义分解成文字的视觉模拟。我们不是试图移除对象,而是移除很可能不形成语义片段的随机补丁。同样,我们的MAE重构像素,这些像素不是语义实体。然而,我们观察到(例如,图4),我们的MAE推断出复杂的整体重建,这表明它已经学习了许多视觉概念,即语义。我们假设这种行为是通过MAE内部丰富的隐藏表示发生的。我们希望这一观点将激励未来的工作。
Broader impacts. 所提出的方法基于训练数据集的学习统计数据预测内容,因此将反映这些数据中的偏差,包括具有负面社会影响的偏差。模型可能生成不存在的内容。这些问题值得进一步研究和考虑,以在这项工作的基础上生成图像。
A. Implementation Details
A.1. ImageNet Experiments
ViT architecture. 我们遵循标准的ViT架构[16]。它有一堆Transformer块[57],每个块由一个多头自注意力块和一个MLP块组成,两者都具有LayerNorm(LN)[1]。编码器以LN结束。由于MAE编码器和解码器具有不同的宽度,我们在编码器之后采用线性投影层来匹配它。我们的MAE将位置嵌入[57](正弦-余弦版本)添加到编码器和解码器输入。我们的MAE不使用相对位置或层缩放(在[2]的代码中使用)。
我们从编码器输出中提取特征用于微调和线性探测。由于ViT有一个类标记[16],为了适应这种设计,在我们的MAE预训练中,我们向编码器输入添加了一个辅助虚拟标记。该标记将被视为用于在线性探测和微调中训练分类器的类标记。我们的MAE在没有此标记的情况下也能很好地工作(具有平均池)。
Pre-training. 默认设置如表8所示。我们不使用颜色抖动、下降路径(drop path)或梯度裁剪。我们使用xavier_uniform[18]来初始化所有Transformer块,遵循ViT的官方代码[16]。我们使用线性lr缩放规则[20]:lr=base_lr×batchsize/256。
End-to-end fine-tuning. 我们的微调遵循监督ViT训练的常见做法。默认设置如表9所示。我们在[2]之后使用逐层lr衰减[10]。
Linear probing. 我们的线性分类器训练参考[9]。见表10。我们观察到,线性探测需要与端到端微调非常不同的配置。特别是,正则化通常对线性探测有害。参考[9],我们禁用了许多常见的正则化策略:我们不使用mixup[69]、cutmix[68]、下降路径[30]或颜色抖动,我们将权重衰减设置为零。
在训练经典线性分类器(例如,SVM[11])时,对分类器输入进行归一化是一种常见的做法。类似地,在训练线性探测分类器时,对预训练的特征进行归一化是有益的。参考[15],我们采用了没有仿射变换的额外BatchNorm层[31](affine=False)。该层应用于编码器产生的预训练特征,并且在线性分类器之前。我们注意到,该层不会破坏线性特性,并且可以在训练后被吸收到线性分类器中:它本质上是一个重参数化的线性分类器。3 引入该层有助于校准消融中不同变体的特征量值,以便它们可以使用相同的设置而无需进一步的lr搜索。
3 或者,我们可以预先计算特征的均值和标准差,并使用归一化特征来训练线性分类器。
Partial fine-tuning. 我们的MAE部分微调(§4.3)遵循表9中的设置,除了我们调整微调回合的数量。我们观察到,调整更少的块需要更长的时间表。我们将微调回合的数量设置为{50, 100, 200},并为每个调整的块数使用最佳回合。
A.2. Supervised Training ViT-L/H from Scratch
我们发现,在ImageNet-1K上从头开始训练受监督的ViT-L/H是不容易的。训练不稳定。虽然对于较小的模型已经有了公开可用的实现的强大基准[53],但对于较大的ViT-L/H的配置尚未探索。直接将之前的配置应用于这些较大的模型是行不通的。训练期间经常观察到NaN损失。
我们在表11中提供了我们的配置。我们使用了0.3的wd、4096的大批量和长时间的预热(warmup),遵循了最初的ViT[16]。我们参考[6]使用β2=0.95。我们使用表11中列出的归一化,并禁用其他归一化,如[64]所示。所有这些选择都是为了提高训练稳定性。我们的配置可以在无NaN损失的情况下完成训练。ViT-L的准确率为82.6%(81.5%w/o EMA),ViT-H的准确度为83.1%(80.9%w/o EMA)。如果不使用EMA,ViT-L和ViT-H都显示过拟合趋势。
作为副产物,我们的ViT-B配置具有82.3%的准确度(82.1%w/o EMA),而[53]中为81.8%。
A.3. Object Detection and Segmentation in COCO
我们在Mask R-CNN[24]中调整了原始ViT以使用FPN主干[36]。ViT有一堆Transformer块,它们都以单个比例(例如步幅16)生成特征图。我们将堆栈平均分成4个子集,并应用卷积对中间特征图进行上采样或下采样,以产生不同的比例(步长4、8、16或32,与标准ResNet[25]相同)。FPN建立在这些多尺度图上。
为了在不同方法之间进行公平比较,我们为表4中的每个条目(包括所有竞争对手)搜索超参数。我们搜索的超参数是学习率、权重衰减、下降路径速率和微调回合。我们将随特定配置一起发布代码。有关完整的模型和训练细节,以及其他实验,请参见[35]。
A.4. Semantic Segmentation in ADE20K
我们在[2]的语义分割代码之后使用UperNet[63]。我们对100个回合进行端到端微调,批大小为16。我们为表5中的每个条目(包括所有竞争对手)搜索最佳lr。
[2]的语义分割代码使用相对位置偏差[49]。我们的MAE预训练不使用它。为了公平比较,我们仅在迁移学习期间启用相对位置偏差,初始化为零。我们注意到,我们的BEiT复制在预训练和微调中都使用了相对位置偏差,遵循它们的代码。
A.5. Additional Classification Tasks
我们遵循表9中的iNaturalist和Places微调设置(表6)。我们调整每个数据集的lr和微调回合。
B. Comparison on Linear Probing Results
在§4.3中,我们已经表明,线性探测精度和微调精度在很大程度上是不相关的,它们对线性可分性有不同的关注。我们注意到,现有的掩码图像编码方法在线性探测(例如,对比学习)中的竞争力通常较低。为了完整性,在表12中,我们比较了线性探测精度与基于掩码的方法。
我们使用ViT-L的MAE具有75.8%的线性探测精度。这比以前的基于掩码的方法要好得多。另一方面,它仍然落后于该规约下的对比方法:例如,MoCo v3[9]对ViT-L的线性探测精度为77.6%(图9)。
C. Robustness Evaluation on ImageNet
在表13中,我们评估了我们的模型对ImageNet验证集的不同变体的鲁棒性。我们使用在原始ImageNet上微调的相同模型(表3),只在不同的验证集上运行推断,没有任何专门的微调。表13显示了我们的方法具有很强的缩放行为:增加模型大小有显著的好处。增加图像大小有助于除IN-C以外的所有集合。我们的结果在很大程度上超过了以前(专业系统)的最佳结果。
相比之下,监督训练的表现要差得多(表13底部;A.2中描述的模型)。例如,使用ViT-H,我们的MAE预训练在IN-A方面比受监督的对应工作好35%(68.2% vs. 33.1%)。