本文证明了掩蔽自动编码器(MAE)是一种可扩展的计算机视觉自监督学习器。
我们的MAE方法很简单:我们掩盖输入图像的随机补丁,并重建缺失的像素。
它基于两个核心设计:
将这两种设计结合起来,使我们能够高效地训练大型模型:我们加快了训练(3倍或更多),并提高了准确性。
我们的可扩展的方法允许学习高容量的模型,这些模型具有很好的通用性:例如,在使用仅ImageNet-1K数据的方法中,一个普通的ViT-Huge模型获得了最好的准确性(87.8%)。下游任务的迁移性能优于监督前训练,表现出良好的扩展行为。
深度学习见证了能力和容量不断增长的架构的爆炸式发展[28,24,47]。得益于硬件的快速发展,今天的模型可以很容易地超过一百万张图片[13],并且开始需要上亿张标签为[16]的图片(通常是公开不可访问的)。
这种对数据的需求已经通过自我监督的预训练成功地在自然语言处理(NLP)中得到了解决。该解决方案基于GPT中的自回归语言建模[40,41,4]和BERT[14]中的掩码自编码,在概念上很简单:它们删除部分数据,并学会预测删除的内容。这些方法现在可以训练包含超过1000亿个参数[4]的可推广的NLP模型。
掩蔽自动编码器的思想是:一种更通用的去噪自动编码器[48]的形式,在计算机视觉中也很自然和适用。事实上,与视觉密切相关的研究[49,39]早于BERT。然而,尽管BERT的成功引起了人们对这一思想的极大兴趣,但视觉自编码方法的进展滞后于NLP。我们问:是什么让masked的自动编码区别于视觉和语言?我们试图从以下几个方面来回答这个问题:
(i)直到最近,结构都是不同的。在视觉领域,卷积网络[29]在过去十年中占据主导地位。卷积通常在规则的网格上运行,将掩码令牌[14]或位置嵌入[47]等“指标”集成到卷积网络中并不简单。然而,这个架构上的差距已经通过Vision transformer (ViT)[16]的引入得到了解决,应该不再是一个障碍了。
(ii) 语言和视觉的信息密度是不同的。语言是人类产生的高度语义和信息密集的信号。当训练一个模型每句只预测几个遗漏的单词时,这个任务似乎诱发了复杂的语言理解。相反,图像是具有大量空间冗余的自然信号。丢失的补丁可以从邻近的补丁中恢复,而这些补丁对部件、对象和场景的理解很少。为了克服这一差异,并鼓励学习有用的特征,我们展示了一种简单的策略在计算机视觉中很好地工作:mask很大一部分的随机补丁。这种策略在很大程度上减少了冗余,并创建了一个挑战性的自我监督任务,需要整体理解超越低级图像统计。要对我们的重建任务有一个定性的认识,请参见图2 - 4。
(iii) 自动编码器的解码器将潜在表示映射回输入,在文本和图像的重构之间起着不同的作用。在视觉中,解码器重构像素,因此其输出的语义级别低于一般的识别任务。这与语言相反,在语言中,解码器会预测包含丰富语义信息的缺失单词。虽然在BERT中解码器可以是微不足道的(一个MLP)[14],但我们发现,对于图像,解码器的设计在确定学习的潜在表示的语义级别中起着关键作用。
在此分析的驱动下,我们提出了一个简单、有效和可扩展的形式的掩蔽自动编码器(MAE)的视觉表示学习。
我们的MAE学习非常高容量的模型,可以很好地泛化。通过MAE预训练,我们可以在ImageNet-1K上训练数据匮乏的模型,如ViT-Large/-Huge [16] ,从而提高泛化性能。
在普通的ViT-Huge模型中,我们在ImageNet-1K上进行微调时达到了87.8%的精度。这比以前所有只使用ImageNet-1K数据的结果都要好。我们也评估迁移学习的对象检测,实例分割,和语义分割。在这些任务中,我们的前训练比它的监督前训练对手取得了更好的结果,更重要的是,我们观察到显著的收益通过扩大模型。这些观察结果与NLP自我监督前训练中的观察结果一致[14,40,41,4],我们希望它们能使我们的研究领域探索类似的轨迹。
蒙面语言建模及其自回归对应项,如BERT[14]和GPT[40,41,4],是非常成功的NLP预训练方法。
自编码是学习表示的一种经典方法。它有一个将输入映射到潜在表示的编码器和一个重建输入的解码器。例如,PCA和k-means是[25]自动编码器。去噪自动编码器(DAE)[48]是一类自动编码器,它破坏输入信号,并学会重建原始的、未损坏的信号。一系列的方法可以被认为是在不同的破坏下的一个广义DAE,例如,掩蔽像素[49,39,6]或去除颜色通道[59]。我们的MAE是去噪自编码的一种形式,但与传统的DAE有许多不同之处。
掩蔽图像编码方法从被掩蔽破坏的图像中学习表示。[49]的开创性工作提出了在DAE中作为一种噪声类型的掩蔽。编码器[39]使用卷积网络绘制大量缺失区域。受NLP成功的激励,最近的相关方法[6,16,2]都是基于《Transformers[47]》。iGPT[6]作用于像素序列,并预测未知像素。ViT论文[16]研究了自我监督学习的掩蔽patch预测。最近,BEiT[2]提出预测离散代币[37,43]。
在计算机视觉领域,自我监督学习方法已经引起了人们的极大兴趣,它们通常聚焦于不同的前文本任务,用于预训练[15,50,35,59,38,17]。
最近,对比学习[3,21]已经流行起来,例如,[51,36,22,7],它模拟两个或多个视图之间的图像相似性和不相似性(或仅仅是相似性[20,8])。对比及相关方法强烈依赖于数据扩充[7,20,8]。自编码追求一个概念上的不同方向,它展示了不同的行为,正如我们将要展示的。
我们的掩蔽自动编码器(MAE)是一个简单的自动编码方法,重建原始信号的部分观察。像所有的自动编码器一样,我们的方法有一个编码器,将观察到的信号映射到一个潜在的表示,和一个解码器,从潜在的表示重建原始信号。
与传统的自动编码器不同,我们采用了一种非对称设计,允许编码器只对部分观测信号(没有掩码标记)进行操作,并采用一个轻量级解码器,从潜在表示和掩码标记重新构造完整的信号。图1演示了接下来介绍的想法。
Masking掩蔽。
MAE解码器。
6. MAE解码器的输入是:由(i)编码的可见补丁和(ii)掩码标记组成的完整标记集。参见图1。每个掩码令牌[14]都是一个共享的、学习过的向量,表示存在一个缺失的补丁来预测。我们将位置嵌入添加到这个完整集合中的所有标记中;如果不这样做,遮罩令牌在图像中就不会有关于它们位置的信息。解码器还有另一系列Transformer块。
7. MAE解码器仅在预训练用于执行图像重建任务(仅编码器用于产生图像表示进行识别)。因此,解码器体系结构可以以独立于编码器设计的方式灵活设计。我们实验用非常小的解码器,比编码器窄和浅层低。例如,我们的默认解码器与编码器相比,每个令牌的计算量小于10%。采用这种非对称设计,整套令牌只由轻量级解码器处理,这大大减少了预训练时间。
重建目标。我们的MAE通过预测每个掩蔽块的像素值来重建输入。解码器输出的每个元素都是表示一个patch的像素值向量。解码器的最后一层是线性投影,其输出通道数等于一个patch中的像素值数。解码器的输出被重塑以形成重构图像。我们的损失函数计算重建和原始图像在像素空间的均方误差(MSE)。我们只计算屏蔽补丁的损失,类似于BERT [14].1
我们还研究了一种以每个掩蔽块的归一化像素值为重建目标的变量。具体来说,我们计算一个patch中所有像素的均值和标准差,并用它们来归一化这个patch。在实验中,我们采用归一化像素作为重建目标,提高了图像的表示质量。
简单的实现。我们的MAE预训练可以有效地执行,而且重要的是,不需要任何专门的稀疏操作。
我们使用表1中的默认设置消除MAE。
观察到几个有趣的性质。
屏蔽率。图5显示了掩码比例的影响。最优比率高得惊人。75%的比例对于线性探测和微调都是有利的。这种行为与BERT[14]的典型掩蔽比为15%形成了对比。我们的掩蔽比也远远高于计算机视觉领域的相关研究[6,16,2](20% - 50%)。
该模型推断出缺失的补片会产生不同的,但似乎可信的输出(图4)。
它解释了物体和场景的格式塔,这些格式塔不能简单地通过扩展线条或纹理来完成。我们假设这种类似推理的行为与使用表征的学习有关。
图5还显示了线性探测和微调结果遵循不同的趋势。对于线性探测,精度随着掩蔽比的增加而稳步增加,直到最佳点:精度差距高达20% (54.6% vs. 73.5%)。对于微调,结果不太敏感的比率,和广泛的掩蔽比率(40-80%)工作良好。图5中的所有微调结果都优于从头开始训练(82.5%)。
(图5。屏蔽率。高掩蔽比(75%)对于微调(上)和线性探测(下)都很好。在本文所有图中,y轴为ImageNet-1K验证精度(%)。)
译码器的设计。我们的MAE解码器可以灵活地设计,如表1a和1b所示。
表1a改变了解码器的深度(变压器块的数量)。一个足够深的解码器对线性探测是重要的。这可以用像素重建任务和识别任务之间的差距来解释:自动编码器的最后几层更专门用于重建,但与识别不太相关。一个合理的深度解码器可以解释重构专门化,使潜在的表示处于更抽象的级别。这种设计可以在线性探测方面提高8%(表1a, ’ lin ')。然而,如果使用微调,编码器的最后一层可以调整以适应识别任务。解码器深度对改善微调的影响较小(表1a,“ft”)。
有趣的是,我们使用单块解码器的MAE在微调时表现强劲(84.8%)。请注意,单个Transformer块是传播从可见令牌到掩码令牌的信息的最小需求。这样的小解码器可以进一步加快训练速度。
在表1b中,我们研究了解码器宽度(信道数)。我们默认使用512-d,它在微调和线性探测下表现良好。较窄的解码器也能很好地进行微调。
总的来说,我们默认的MAE解码器是轻量级的。它有8个块,宽度为512-d(表1中的灰色)。与ViT-L 相比,每个令牌只有9%的FLOPs(24个块,1024-d)。因此,当解码器处理所有令牌时,它仍然只是整个计算的一小部分。
面具令牌Mask token.。我们的MAE的一个重要设计是跳过编码器中的掩码令牌[M],并将其应用到轻量级解码器中。表1c研究了这一设计。
如果编码器使用掩码标记,则性能更差:在线性探测中,其精度下降14%。在这种情况下,在预训练和部署之间有一个间隙:这个编码器在预训练的输入中有很大一部分掩码令牌,这在未损坏的图像中是不存在的。这一差距可能会降低部署的准确性。通过从编码器中移除掩码令牌,我们约束编码器始终看到真实的补丁,从而提高精度。
此外,通过跳过编码器中的掩码标记,大大减少了训练的计算量。在表1c中,我们将总的训练FLOPs减少了3.3×。在我们的实现中,这导致了2.8倍的墙壁时钟加速(见表2)。墙壁时钟加速甚至更大(3.5 - 4.1倍),对于更小的解码器(1块),更大的编码器(ViT-H),或两者都是。请注意,当掩蔽率为75%时,加速可以达到>4×,部分原因是自注意复杂度是二次的。此外,内存大大减少,这可以使训练更大的模型或加快更多的大批量训练。在时间和记忆效率方面,MAE适合于训练非常大的模型。
重建目标。我们在表1d中比较了不同的重建目标。到目前为止,我们的结果是基于没有(每个补丁)标准化的像素。使用具有归一化的像素可以提高精度。这种每片正态化可以局部增强对比度。在另一种变体中,我们在patch空间中执行PCA,并使用最大的PCA系数(这里是96)作为目标。这样做会降低交流准确度。两个实验都表明高频分量在我们的方法中是有用的。
我们还比较了预测tokens的MAE变体,tokens是BEiT[2]中使用的目标。特别对于这个变体,我们使用经过DALLE预处理的dVAE[43]作为标记器,后面是[2]。在这里,MAE解码器使用交叉熵损失预测骰子中的令牌。与非规范化像素相比,这种标记化将微调精度提高了0.4%,但与规范化像素相比没有任何优势。它也降低了线耳探测的准确性。在§5中,我们进一步证明了在迁移学习中,标记化是不必要的。
我们基于像素的MAE比标记化简单得多。dVAE标记器需要多一个训练前阶段,这个阶段可能需要额外的数据(250M图像[43])。dVAE编码器是一个大型卷积网络(ViT-L的40% FLOPs),并增加了不小的开销。使用像素就不会遇到这些问题。
数据增强。表1e研究了数据增强对我们MAE预训练的影响。
我们的MAE可以很好地使用只收割的增强,无论是固定大小还是随机大小(都有随机水平翻转)。添加颜色抖动会降低实验结果,因此我们没有在其他实验中使用它。
令人惊讶的是,即使不使用数据增强(只使用中心裁剪,不翻转),我们的MAE也表现得很好。这一特性与对比学习及相关方法有很大的不同[51,22,7,20],后者严重依赖于数据扩充。[20]观察到,仅使用作物增效法分别降低了BYOL[20]和SimCLR[7]的准确性13%和28%。此外,没有证据表明对比学习可以在没有增强的情况下工作:图像的两种视图是相同的,很容易满足一个简单的解决方案。
在MAE中,数据增强的作用主要是通过随机掩蔽来实现的。每次迭代的masks都是不同的,因此不管数据是否增强,它们都会生成新的训练样本。pretext任务由于掩蔽而变得困难,并且需要较少的增强来正则化训练。
面具抽样策略Mask sampling strategy.。在表1f中,我们比较了不同的掩码采样策略,如图6所示。
[2]中提出的逐块屏蔽策略倾向于删除较大的块(图6中)。我们的分段掩蔽的MAE在50%的比例下工作得相当好,但在75%的比例下降解。这个任务比随机抽样更难,因为观察到更高的训练损失。重建图像也变得模糊了。
我们还研究了网格抽样,它有规律地保持每四个补丁中的一个(图6右边)。这是一项比较容易完成的任务,训练损失也比较小。重建更清晰。但是,表示质量较低。
简单的随机抽样最适合我们的MAE。它允许更高的掩蔽比,这提供了更大的加速效益,同时也享受良好的准确性。
训练计划。到目前为止,我们的消融是基于800epoch前的训练。图7显示了训练计划长度的影响。随着训练时间的延长,准确率稳步提高。事实上,我们甚至在1600epoch也没有观察到线性探测精度的饱和。这种行为不同于对比学习方法,例如MoCo v3[9]在vitl的300个epoch饱和。请注意,MAE编码器每个epoch只看到25%的补丁,而在对比学习中,编码器每个epoch看到200%(两茬)或更多(多茬)的补丁。
与自我监督方法的比较。在表3中,我们比较了自监督ViT模型的微调结果。对于ViT-B,所有方法都很接近。对于ViT-L,方法之间的差距更大,这表明对于更大的模型来说,一个挑战是减少过拟合。
我们的MAE可以很容易地扩大规模,并且比更大的模型显示出稳定的改进。使用ViT-H(224尺寸),我们获得了86.9%的准确率。通过448尺寸的微调,我们实现了87.8%的精度,仅使用IN1K数据。基于先进的网络,在所有仅使用IN1K数据的方法中,先前的最佳精度为87.1%(512大小)[56]。在竞争激烈的IN1K基准测试(没有外部数据)中,我们比最先进的技术提高了不少。我们的结果是基于普通的ViT,我们预计高级网络的性能会更好。
与BEiT[2]相比,我们的MAE更准确、更简单、更快。我们的方法重建像素,与BEiT预测令牌相反:BEiT报告说,当用ViT-B.重建像素时,[2]退化了1.8%我们不需要dVAE预训练。此外,我们的MAE比BEiT快得多(3.5× / epoch),原因如表1c所示。
表3中的MAE模型经过了1600个epoch的预训练以获得更好的精度(图7)。即便如此,如果在相同的硬件中训练,我们的总预训练时间比其他所有方法都要短。例如,使用ViT-L,对于1600个epoch,我们的MAE的训练时间是31小时,而MoCo v3的训练时间是36小时,对于300个epoch[9],使用相同的128个TPU-v3核心。
与有监督的预训练进行比较。在原始的ViT论文[16]中,当在IN1K中训练时,ViT-L 退化。参见图8所示。我们改进的监督配方对于从零开始的训练效果更好(图8,“我们的impl.”;见A.2),但精度是饱和的。
我们的MAE预训练,仅使用IN1K,可以更好地泛化:从零训练的增益更大的容量模型。这与[16]中JFT-300M监督前训练的趋势相似。这个比较表明,我们的MAE可以帮助放大模型尺寸。
表1显示,linear probing线性探测和微调结果在很大程度上是不相关的。在过去的几年中,linear probing线性探测已经成为一种流行的协议;然而,它错过了追求强大但非线性特征的机会——这确实是深度学习的优势。作为中间立场,我们研究了一个局部微调协议:微调最后几层,冻结其他几层。这一方案也在早期的著作中使用,如[54,59,35]。
图9显示了结果。值得注意的是,只需微调一个Transformer块,就可以将精度从73.5%显著提高到81.0%。此外,如果我们只微调最后一个块的“一半”(即它的MLP子块),我们可以得到79.1%,这比线性探测好得多。这种变体本质上是对MLP头的微调。微调几个块(如4块或6块)可以达到不错的精度,与冻结的主干相比,这仍然是一个小的微调头。
在图9中,我们还与MoCo v3[9]进行了比较,这是一种与ViT-L结果可用的对比方法。它比我们的MAE具有更高的线性探测精度。但是,它的部分微调结果都比我们的差。当调整4个块时,差距是2.6%。这些结果表明,MAE表征具有较弱的线性可分性,但具有较强的非线性特征,在非线性磁头调谐时表现良好。
这些观察结果表明,线性可分性并不是评价表征质量的唯一指标。还观察到(如[8]),线性探测与迁移学习性能(如对象检测)没有很好的相关性。据我们所知,线性评估在NLP中并不经常用于培训前的基准评估。
我们评估了迁移学习在COCO[32]上的对象检测和分割,在ADE20K[60]上的语义分割。我们使用表3中的预训练模型。
目标检测与分割Object detection and segmentation.。我们微调Mask R-CNN[23]端到端的COCO。ViT骨干适合与FPN[31]一起使用(参见附录A.3)。我们将此对象检测系统应用于表4中的所有条目。我们box AP用于对象检测,掩码AP用于实例分割。
与监督前训练相比,我们的MAE在所有配置下都表现得更好(表4)。在ViT-B较小的情况下,我们的MAE比监督前训练高2.4分(50.3比47.9,APbox)。更重要的是,在ViT-L较大的情况下,我们的MAE训练前比监督训练前表现好4.0分(53.3比49.3)。
基于像素的MAE优于或与基于令牌的BEiT相当,而MAE更简单和更快。MAE和BEiT都优于MoCo v3, MoCo v3与监督前训练相当。
Semantic segmentation.语义分割。我们在ADE20K上的实验使用了跟随[2]代码的UperNet[52]。详情见A.4。表5显示,我们的MAE显著提高了vitl的迁移结果,比监督训练前对照者(53.6 vs 49.9)好3.7个点。基于像素的MAE优于基于令牌的BEiT。这些观测结果与COCO的观测结果一致。
Pixels vs. tokens.像素和令牌。表6给出了作为MAE重建目标的像素与令牌的全面比较。虽然使用dVAE令牌比使用未归一化像素更好,但在统计上与在我们研究的所有任务和模型中只使用归一化像素相似。这再次表明,标记化对于我们的MAE不是必要的。
可扩展性好的简单算法是深度学习的核心。在NLP中,简单的自我监督学习方法(例如,[40,14,41,4])可以从指数缩放模型中获益。在计算机视觉中,尽管在自我监督学习方面取得了进展,但实际的训练前范式仍主要受到监督(如[28,44,24,16])。在本研究中,我们在ImageNet和迁移学习中观察到自动编码器(一种简单的自我监督方法,类似于nlp中的技术)提供了可扩展的好处。视觉中的自我监督学习现在可能正走上与NLP相似的轨迹。
另一方面,我们注意到,图像和语言是不同性质的信号,必须仔细处理这种差异。图像只是光的记录,没有语义分解成文字的视觉模拟。我们不是试图删除对象,而是重新移动那些很可能不会形成语义段的随机补丁。同样,我们的MAE重建像素,像素不是语义实体。然而,我们观察到(如图4),我们的MAE推断出复杂的、整体的重建,表明它已经学会了许多视觉概念,即语义。我们假设这种行为是通过MAE内部丰富的隐藏表示发生的。我们希望这一观点将激励未来的工作。
更广泛的影响。提出的方法基于训练数据集的学习统计数据预测内容,因此将反映这些数据中的偏见,包括具有负面社会影响的偏见。模型可能生成不存在的内容。这些问题值得进一步的研究和考虑,当建立在这项工作的基础上生成图像。