TPAMI 2024 | 最强的视觉学习者!对比掩码自编码器
Contrastive Masked Autoencoders are Stronger Vision Learners
题目:对比掩码自编码器是更强的视觉学习者
作者:Zhicheng Huang , Xiaojie Jin , Chengze Lu , Qibin Hou , Ming-Ming Cheng , Dongmei Fu , Xiaohui Shen , and Jiashi Feng
源码:https://github.com/ZhichengHuang/CMAE
摘要
掩码图像建模(MIM)在各种视觉任务上取得了有希望的结果。然而,学习到的表示的有限可区分性表明,要成为一个更强的视觉学习者还有很长的路要走。为了实现这一目标,我们提出了对比掩码自编码器(CMAE),这是一种新的自监督预训练方法,用于学习更全面和有能力的视觉表示。通过精心设计,CMAE 统一了对比学习(CL)和掩码图像模型(MIM),利用它们各自的优势,学习具有强大实例可区分性和局部感知性的表示。具体来说,CMAE 由两个分支组成,其中在线分支是不对称的编码器-解码器,而动量分支是动量更新的编码器。在训练期间,在线编码器从掩码图像的潜在表示重建原始图像以学习整体特征。动量编码器接收完整图像,通过与其在线对应物的对比学习增强特征可区分性。为了使 CL 与 MIM 兼容,CMAE 引入了两个新组件,即像素移动用于生成合理的正视图和特征解码器用于补充对比对的特征。得益于这些新颖设计,CMAE 有效地提高了表示质量和转移性能,超过了其 MIM 对应物。CMAE 在图像分类、语义分割和目标检测的极具竞争力的基准测试中实现了最先进的性能。值得注意的是,CMAE-Base 在 ImageNet 上达到了 85.3% 的 top-1 准确率,在 ADE20k 上达到了 52.5% 的 mIoU,分别超过了之前最佳结果的 0.7% 和 1.8%。
关键词
- 掩码图像建模(MIM)
- 对比学习
- 自监督学习
I. 引言
掩码图像建模(MIM)[23],[30],[56] 最近在自监督学习领域引起了越来越多的关注,这得益于其方法的简单性和学习丰富、全面表示的能力。按照自然语言处理中掩码语言建模的思想[19],它们随机掩码训练图像的大部分补丁,并使用自编码器[31]重建掩码补丁的原始信号(例如,原始像素,离线提取的特征)。[23],[30],[56] 已经表明,这样一个简单的框架在 ImageNet 分类[18]和一些下游任务(如目标检测和语义分割)中优于以前的自监督学习方法。
当我们反思 MIM 的成功时,不可避免地要将其与另一种经过验证并广泛使用的自监督学习方法,即对比学习(CL)[2],[42] 进行比较。通过采用一个简单的区分性思想,即将同一图像的不同视图的表示拉近,并将不同图像的表示推开,CL 方法自然地赋予了预训练模型强大的实例可区分性。与 CL 相反,MIM 更多地关注学习输入图像中的局部关系以完成重建任务,而不是像对比学习那样对不同图像之间的关系进行建模[35]。因此,人们怀疑 MIM 在学习区分性表示方面效率较低。这个问题在[30],[56] 的实验结果中已经显现。基于以上分析,自然而然地会提出这样一个问题:我们能否利用对比学习来进一步加强 MIM 方法学习到的表示?或者换句话说,MIM 方法能否从对比学习中受益?沿着这个方向,一些当代的工作尝试通过简单地结合对比学习和 MIM 学习目标来训练视觉表示模型[46],[65]。但它们与 MIM 方法相比只显示出边际性能提升。这些结果表明,充分利用两个框架的优势并非易事。这些挑战归因于它们之间的各种区别,包括输入增强、训练目标、模型架构等。
为了克服这些挑战并为下游任务学习更好的图像表示,我们的目标是探索一种可能的方法,通过统一框架中的对比学习来提升 MIM。通过一系列仔细的研究,我们发现输入视图增强和潜在特征对齐在协调 MIM 和对比学习中起着重要作用。因此,我们致力于这些组件的开发我们的方法。
我们提出的方法的概述如图 1 所示。更具体地说,我们的方法引入了一个对比 MAE(CMAE)框架来进行表示学习。它采用孪生架构[5]。一个分支是一个在线更新的不对称编码器-解码器,它从几个可见补丁中学习潜在表示以重建掩码图像,类似于 MAE。另一个分支是一个提供对比学习监督的动量编码器。为了利用对比学习来提高编码器输出的特征质量,我们在在线分支中引入了一个辅助特征解码器,其输出特征用于与动量编码器输出进行对比学习。
我们精心设计了每个 CMAE 组件,以使对比学习能够使 MIM 受益。与仅包含可见补丁的在线编码器不同,CMAE 动量编码器接收整个图像补丁集。这种设计确保了其输出特征的语义完整性,以指导在线编码器。另一个值得注意的设计选择是:我们的方法使用两个解码器,一个用于预测图像像素并执行 MIM 任务;另一个用于恢复掩码标记的特征。由于每个补丁的语义不完整且含糊不清,直接使用补丁的特征进行对比学习是有问题的。使用辅助特征解码器可以解决这个问题,从而有利于在线分支内的潜在表示学习。此外,与现有方法不同,这些方法对输入使用强烈的空间数据增强,我们提出了一种像素移动增强方法,用于在对比学习中生成更合理的正视图。这种简单的增强方法被证明对通过对比学习改进 MIM 是有效的。有了上述新颖设计,我们 CMAE 方法的在线编码器可以学习到更具区分性的全息信息特征,并在各种预训练和转移学习视觉任务上实现最先进的性能。
我们的贡献总结如下。1)我们提出了一种新的 CMAE 方法,探索如何通过使用对比学习来改进 MIM 的表示。其学习到的表示不仅保留了局部上下文敏感特征,还模拟了不同图像之间的实例可区分性。2)为了将对比学习应用于 MIM,我们提出了一个特征解码器来补充掩码特征,并提出了一种弱空间移动增强方法来生成合理的对比视图,这两种方法都有效地提高了编码器的特征质量。3)如图 2 所示,我们的方法显著提高了 MIM 的学习表示,并设定了新的最先进的性能。特别是,与先前的艺术作品相比,CMAE 在 ImageNet-1k 分类验证集上实现了 0.7% 的绝对增益,在 ADE20K 语义分割验证数据集上实现了 1.8% 的 mIoU,在 CoCo 验证集上实现了 0.4% 的 APb 和 0.5% 的 APm。
III. 方法
A. 框架
我们的方法的整体框架如图3所示,由三个组件组成。在线编码器和解码器学习从掩蔽观察中重建输入图像。与现有的MIM方法(例如,MAE [30] 和 SimMIM [56])不同,我们的方法通过空间移动裁剪操作进一步处理输入图像。更重要的是,我们的解码器增加了一个额外的特征解码器,用于预测输入图像的特征。动量编码器将输入图像的增强视图转换为特征嵌入,与在线特征解码器预测的特征进行对比学习。这样,在线编码器学习到的表示不仅捕捉输入图像的整体特征,还捕捉区分特征,从而实现更好的泛化性能。我们现在详细阐述这些组件。
让我们将输入图像IsI_sIs表示为被标记化为N个图像块标记序列{xis}i=1N\{x_i^s\}^N_{i=1}{xis}i=1N,其中N是图像块的总数。对于IsI_sIs的掩蔽版本,可见标记集表示为{xvs}\{x_v^s\}{xvs}。类似地,输入图像ItI_tIt被标记化为图像块标记序列{xjt}j=1N\{x_j^t\}^N_{j=1}{xjt}j=1N,以作为动量编码器的输入。
在线编码器:在线编码器FsF_sFs将可见标记xvsx_v^sxvs映射到嵌入特征zvsz_v^szvs。给定标记序列{xis}i=1N\{x_i^s\}^N_{i=1}{xis}i=1N,我们掩蔽出大部分块,并将访问过的块输入到在线编码器。在线编码器采用视觉变换器(ViT)架构[21],遵循MAE [30]。它首先通过线性投影嵌入可见标记xvsx_v^sxvs作为标记嵌入,并添加位置嵌入pvsp_v^spvs。我们将融合的嵌入输入到一系列变换器块中,并得到嵌入特征zvsz_v^szvs。
zvs=Fs(xvs+pvs)(1) z_v^s = F_s(x_v^s + p_v^s) \quad (1) zvs=Fs(xvs+pvs)(1)
预训练后,在线编码器FsF_sFs用于下游任务中提取图像表示。
动量编码器:引入动量编码器是为了为在线编码器提供对比监督,以学习区分表示。与现有的基于连体的方法[10],[65]不同,我们的动量编码器FtF_tFt仅用于对比学习,以及引导在线编码器学习更区分的特征。它与在线编码器FsF_sFs共享相同的架构,但以整个图像为输入,以保留学习到的表示的语义完整性和区分性。对动量编码器使用整个图像作为输入对方法性能很重要,这在第IV-D节中通过实验验证。与NLP中的标记不同,图像标记在其语义含义上几乎是确定的,图像标记在语义含义上是模糊的[65]。为了避免模糊性,我们采用全局表示进行对比学习。动量编码器的均池特征用于其简单性,即,
zt=1N∑j=1NFt(xjt),(2) z_t = \frac{1}{N} \sum_{j=1}^N F_t(x_j^t), \quad (2) zt=N1j=1∑NFt(xjt),(2)
其中xjtx_j^txjt是动量编码器的输入标记,ztjz_t^jztj是动量编码器的输出序列,ztz_tzt是在输出序列ztjz_t^jztj上执行均池操作后获得的特征,用于表示输入图像。
与在线编码器不同,我们通过指数移动平均(EMA)更新动量编码器的参数。也就是说,将FsF_sFs和FtF_tFt的参数分别表示为θs\theta_sθs和θt\theta_tθt,参数通过以下方式更新:
θt←μθt+(1−μ)θs \theta_t \leftarrow \mu \theta_t + (1 - \mu) \theta_s θt←μθt+(1−μ)θs
这里μ\muμ在实验中固定为0.996。动量更新之所以被使用,是因为它通过促进平滑的特征变化来稳定训练,如在MoCo [29]和BYOL [26]中发现的。
在线解码器:解码器的目标是将潜在特征zvsz_v^szvs和MASK标记特征映射到动量编码器和原始图像的特征空间。具体来说,解码器接收编码的可见标记zvsz_v^szvs和MASK标记zmsz_m^szms。
类似于MAE [30],位置嵌入被添加到输入标记中。由于不同的映射目标,我们的在线解码器有两个分支的解码器结构,一个是像素解码器,另一个是特征解码器。像素解码器GpG_pGp学习重建掩蔽块的像素。我们使用完整的标记集,其中包含zvsz_v^szvs和zmsz_m^szms,来预测块的像素ymy_mym。这个模块可以促进模型学习每个图像块的整体表示。我们将像素解码器设置为堆叠的变换器块:
ym′=I⋅Gp(zvs,zms)(3) y'_m = I \cdot G_p(z_v^s, z_m^s) \quad (3) ym′=I⋅Gp(zvs,zms)(3)
其中I是仅选择输出序列ypy_pyp中对应掩蔽标记的预测的指示器,ym′y'_mym′是掩蔽块的输出预测。
为了与动量编码器的输出对齐,特征解码器GfG_fGf被应用以恢复掩蔽标记的特征。特征解码器具有与像素解码器相同的结构,但参数不共享,以服务于不同的学习目标。这种设计选择的突出性将在第III-D节的架构部分中讨论。给定编码的可见标记zvsz_v^szvs,我们添加掩蔽标记zmsz_m^szms,并使用这个完整集来预测掩蔽标记的特征。类似于在动量编码器中所做的,我们在特征解码器的输出yfy_fyf上应用均池操作,将整个图像表示为ysy_sys,然后使用这个特征进行对比学习。
ys=1N∑i=1NGf(zvs,zms)(4) y_s = \frac{1}{N} \sum_{i=1}^N G_f(z_v^s, z_m^s) \quad (4) ys=N1i=1∑NGf(zvs,zms)(4)
其中N是完整集中标记的总数。
B. 视图增强
通常,掩蔽图像建模预训练任务只使用输入图像的一个视图,其中只包含访问过的块。但是对比学习通常采用两个不同的增强视图。为了使MIM和对比学习彼此兼容,我们的方法还生成两个不同的视图,并将它们分别输入到其在线和动量分支。
在对比学习中,最常用的视图增强可以分为两种类型:空间转移(例如,随机调整大小的裁剪,翻转)和颜色转移(例如,颜色抖动和随机灰度化)。对于MIM任务,颜色增强会降低结果[30],因此我们不对在线分支的输入应用它们。空间和颜色数据增强应用于动量分支输入,以避免平凡解。
我们首先考虑两个分支使用两个不同的随机裁剪,遵循对比学习中的常见做法。然而,我们观察到这种方法对模型性能有负面影响(见第IV-D节)。我们推测这个问题与在线/动量编码器的输入之间的差异有关,当随机裁剪的区域相距甚远或语义相关性很少时。与通常对比方法中使用成对视图不同,MIM中掩蔽操作,利用大掩蔽比率(例如,在[30]中为75%),可能放大这种差异,因此创建了假阳性视图。因此,在这些错位的正对上进行对比学习实际上会带来噪声,并阻碍学习区分和有意义的表示。
为了解决上述问题,我们提出了一种名为像素移动的弱增强方法,用于生成在线/动量编码器的输入。核心思想是首先通过原始图像的调整大小随机裁剪获得主图像。然后两个分支共享相同的主图像,并通过在主图像上轻微移动裁剪位置来生成各自的视图。更详细地说,我们将主图像表示为III。III的形状是(w+p,h+p,3)(w + p, h + p, 3)(w+p,h+p,3),其中w,h是我们模型的目标输入大小的宽度和高度,p是允许的最大移动范围。对于在线分支,我们使用[0:w, 0:h, :]区域作为我们的输入图像IsI_sIs。对于动量分支,我们使用[rw: rw + w, rh: rh + h, :]区域作为我们的输入图像ItI_tIt。rw和rh是在[0, p)范围内的独立随机值。随后,我们仅对在线编码器的输入图像ItI_tIt应用掩蔽操作而不应用颜色增强,这与MAE一致。对于动量编码器的输入图像,我们使用颜色增强但不应用掩蔽操作。每个编码器输入的不同增强生成同一图像的不同视图,以促进在线和动量编码器之间的对比学习。
C. 训练目标
重建损失:遵循[30],我们使用归一化的像素作为重建任务的目标。我们采用均方误差(MSE)作为损失函数,并仅在像素解码器预测和原始图像之间的掩蔽块上计算损失。数学公式为:
Lr=1Nm∑(ym′−ym)2(5) L_r = \frac{1}{N_m} \sum_{(y'_m - y_m)^2} \quad (5) Lr=Nm1(ym′−ym)2∑(5)
其中NmN_mNm是图像中掩蔽块的数量,LrL_rLr是损失值。
对比损失 :为了清晰起见,我们从损失函数和水头结构两个方面来描述我们的方法的对比损失设计。我们考虑了两个广泛使用的损失函数,即InfoNCE [12]、[29]损失和BYOLstyle [10]、[26]损失。前者旨在同时拉近来自同一样本的正视图表示并推开负样本,而后者仅最大化正视图之间的相似性。尽管一些近期的工作发现它们可能本质上是统一的,但由于它们在表示学习中的不同影响,我们仍然分别分析它们。在我们的方法中,我们观察到使用InfoNCE损失可以获得更好的性能,因此我们默认使用它。详细信息请参阅第IV-D节。对于头部结构,我们采用了广泛使用的“投影-预测”结构,遵循[14],[26]。具体来说,我们将“投影-预测”和“投影”头部分别附加到特征解码器和动量编码器。动量编码器的投影头部HtH_tHt也通过指数移动平均进行更新。由于在线/动量编码器生成输入的差异(参见第III-B节),我们使用不对称的对比损失,这与以前的方法[14],[26]不同。特征解码器的表示ysy_sys通过“投影-预测”结构转换为ypsy_p^syps。类似地,对于动量编码器的表示ztz_tzt,我们应用投影头部并得到zptz_p^tzpt。然后我们计算它们之间的余弦相似度ρ\rhoρ:
ρ=yps⋅zpt∥yps∥2∥zpt∥2(6) \rho = \frac{y_p^s \cdot z_p^t}{\|y_p^s\|_2 \|z_p^t\|_2} \quad (6) ρ=∥yps∥2∥zpt∥2yps⋅zpt(6)
我们表示ρ+\rho^+ρ+为正对余弦相似度,由同一图像的ypsy_p^syps和zptz_p^tzpt构成。ρj−\rho^-_jρj−表示第j个负对的余弦相似度。我们使用不同图像中的zptz_p^tzpt构建负对。InfoNCE损失的损失函数是:
Lc=−logexp(ρ+/τ)exp(ρ+/τ)+∑j=1K−1exp(ρj−/τ)(7) L_c = -\log \frac{\exp(\rho^+/\tau)}{\exp(\rho^+/\tau) + \sum_{j=1}^{K-1} \exp(\rho^-_j / \tau)} \quad (7) Lc=−logexp(ρ+/τ)+∑j=1K−1exp(ρj−/τ)exp(ρ+/τ)(7)
其中τ\tauτ是温度常数,设置为0.07。K是批量大小。
整体学习目标是重建损失LrL_rLr和对比损失LcL_cLc的加权组合:
L=Lr+λcLc(8) L = L_r + \lambda_c L_c \quad (8) L=Lr+λcLc(8)
D. 连接和分析
为了阐明CMAE与先前方法之间的联系和区别,我们从训练目标、输入和架构等多个角度进行了比较评估。结果展示在表I中。我们的主要关注点是那些在MIM中使用对比信息或掩蔽图像输入的方法。仅使用MIM或对比学习的方法不在本讨论范围内,因为它们显然与我们的方法不同。
训练目标:CMAE在优化过程中利用了重建损失和对比损失。从(7)中推断,CMAE中的对比损失包括了视图内匹配和图像间对比。因此,生成的表示被鼓励展现出实例区分和空间敏感性的期望特征。相比之下,MSN [1] 和 ExtreMA [54] 等方法由于与我们的目标不同,忽略了重建损失,并使用掩蔽输入进行正则化或数据增强。iBot [65] 仅采用正视图之间的蒸馏损失,通过最大化视图内匹配分数,忽略了与负样本的对比学习。此外,CMAE仅采用不对称损失进行对比学习,这比iBot的计算成本更低。尽管SIM [46] 也利用了两种损失,但它在重建目标上与CMAE不同。而CMAE恢复同一视图的掩蔽内容,SIM重建另一个视图的特征。我们的实验结果表明,CMAE不仅更简单,而且在表示学习方面也更有效,如优越的性能所示。
输入:大多数先前的对比学习方法[12],[29] 实施了强大的增强技术(例如,随机裁剪,随机缩放)来从同一图像生成正视图。这些操作也常用于掩蔽图像建模场景下的对比学习模型,如iBot。然而,考虑到掩蔽操作,这利用了大掩蔽比率(例如,在[30]中为75%),已经显著降低了输入,进一步应用这些增强可能会生成无效的正视图,从而阻碍对比学习。相比之下,我们提出了一种新颖的、适度的数据增强操作,称为像素移动,以实现更好的正视图对齐。与ExtreMA [54]相比,后者在两个连体分支中使用完全相同的视图,像素移动引入了适度的输入变化,这被证明对对比学习是有益的(参见表IV(a))。
架构:在CMAE中,在线编码器后附加了一个轻量级特征解码器,以补充掩蔽特征。这是与其他方法(如SIM和iBot)的显著区别,后者直接使用可见部分的表示与未掩蔽视图的表示进行匹配。我们认为,在对比学习中对比掩蔽部分的特征和输入图像的特征是不切实际的,考虑到它们展示了不同层次的抽象和语义覆盖。特征解码器有望通过减少对比特征之间的分布差距来促进优化。特征解码器的有效性在表IV©中得到了实证验证。值得注意的是,CMAE的设计是非侵入式的,允许将其直接应用于现有的MIM模型,如MAE和ConvMAE,而无需对MIM模型进行重大修改。
IV. 实验
A. 实现细节
预训练: 我们遵循 MAE [30] 的设置来预训练我们的模型。我们默认采用 AdamW [40] 优化器,动量设置为 β1=0.9,β2=0.95\beta_1 = 0.9, \beta_2 = 0.95β1=0.9,β2=0.95。此外,权重衰减设置为 0.05。我们使用线性缩放规则 [25]:lr=baselr×batchsize256lr = base_lr \times \frac{batch_size}{256}lr=baselr×256batchsize 来设置学习率。基础学习率为 1.5×10−41.5 \times 10^{-4}1.5×10−4,批量大小为 4096。采用余弦学习率调度 [39],预热期为 40 个周期。所有预训练实验均在 32 个 NVIDIA A100 GPU 上进行。
编码器结构: 我们使用 ViT [21] 基础模型作为我们的默认设置。为了进一步验证我们提出的模型的可扩展性,我们将 ViT 替换为也用于 ConvMAE [23] 的混合卷积 ViT。在混合 ViT 中,使用多层卷积网络 [34] 作为标记投影。注意,混合 ViT 的模型大小与 ViT 版本相同,以进行公平比较。我们还尝试使用扩展的编码器来评估我们方法的可扩展性。
B. 在 ImageNet 上的结果
遵循现有工作 [3], [14], [30], [56],我们使用 ImageNet-1K [18],该数据集包含 1.3M 张 1k 类别的图像,作为预训练和微调的数据集。数据集包含两个子集:训练集和验证集。我们仅使用训练集来预训练 CMAE。预训练后,CMAE 在 ImageNet-1k 训练集上进行微调 100 个周期。对于预训练 300 个周期的模型,我们采用 5.e−45.e^{-4}5.e−4 作为基础学习率。由于更长的预训练计划(1600 个周期)使模型能够学习更好的初始化权重以进行微调 [57],我们设置一个较小的基础学习率为 2.5e−42.5e^{-4}2.5e−4。此外,我们遵循常见的微调实践,使用 mixup [61]、cutmix [59]、drop path [33] 等方法来规范模型。
在表 II 中,我们将 CMAE 与竞争方法在 ImageNet 上的微调分类准确率进行了比较。CMAE 达到了 84.7% 的 top-1 准确率,比 MAE [30] 高出 1.1%。在所有使用 ViT 架构的模型中,CMAE 达到了最佳性能。与基于对比学习的 MoCo-v3 [14] 和 DINO [10] 相比,我们的模型可以显著优于它们,分别提高了 1.5% 和 1.9%。与也使用对比目标的 iBOT 和 SIM 相比,我们的 CMAE 达到了更高的性能,分别提高了 0.7% 和 0.8%。以上结果强烈证明了 CMAE 的优越性。
当我们将普通的 ViT 编码器替换为混合卷积 ViT 时,如 ConvMAE [23] 所做,CMAE 进一步提高到 85.0% 和 85.3%,分别在 800 个周期和 1600 个周期的预训练中。这些结果在相同的预训练设置下超过了 ConvMAE 的 0.4% 和 0.7%,验证了 CMAE 对各种网络结构的卓越可扩展性。
值得注意的是,CMAE 可以通过延长训练计划(从 800 个周期到 1600 个周期)获得显著的改进,而 ConvMAE 在 800 个周期时观察到饱和。这一结果表明,CMAE 在学习更好表示方面具有更强的能力。
C. 迁移学习
为了进一步验证 CMAE 的可转移性,我们遵循之前的方法在语义分割数据集 ADE20K [64]、目标检测数据集 COCO2017 [37] 和分类数据集上评估预训练模型。
语义分割: ADE20K [64] 包含 25,562 张 150 细粒度类别的图像。我们采用 Upernet [55] 作为此任务的默认模型,遵循比较方法的设置。骨干网络 ViT-B 从预训练初始化,而其他模块使用 Xavier [24] 初始化。模型在 ADE20K 训练集上进行微调,并在标准验证分割上进行测试。
按照之前的作品,我们在表 III(a) 中报告了 CMAE 的平均交并比 (mIoU) 性能。我们注意到 CMAE 显著超过了 MAE 的 2.9%,这验证了 CMAE 的更强可转移性。此外,与 iBOT [65] 和 CAE [15] 相比,CMAE 也分别提高了 1.0% 和 0.8%。使用相同的混合 ViT 骨干,CMAE 显著优于 ConvMAE 的 1.8%。值得注意的是,CMAE 通过以 52.5% 的成绩超越所有竞争方法,创下了新的最高记录。
目标检测和分割: 我们采用广泛使用的目标检测和实例分割框架 Mask-RCNN [28], [36] 用于在这项任务上的基准测试。ViT-B 用作骨干网络,并用我们的预训练模型初始化。遵循 MAE,我们在 COCO train2017 分割上对模型进行微调,并在 val2017 分割上报告目标检测的 box AP 和实例分割的 mask AP。我们对模型进行 100 个周期的微调。基础学习率为 1.e−41.e^{-4}1.e−4,采用余弦退火调度,权重衰减设置为 0.1。
与其他自监督学习方法的比较结果如表 III(b) 所示。可以看到,CMAE 在 APb 上从 MAE 的 51.7 提高到 52.4,在 APm 上从 45.9 提高到 46.5。使用混合 ViT 结构,CMAE 一致地超过了竞争方法 ConvMAE:APb 从 52.5 提高到 52.9,APm 从 46.5 提高到 47.0。以上有希望的结果再次验证了我们方法的有效性。
分类任务: 为了进一步研究迁移学习在分类任务上的应用,我们在 iNaturalists [48] 和 Places [63] 上验证了我们的模型。在这些数据集上的四个分类任务的实验结果如表 V 所示,与 MAE [30] 相比,top-1 准确率一致提高了 1.0% 到 1.7%。这些结果为提高预训练模型的区分能力提供了进一步的证据。
D. 方法分析
为了理解关键组件的影响并验证我们在 CMAE 中采用的设计选择,我们进行了一系列消融实验。除非另有说明,否则我们在本小节中报告的模型性能是在 300 个预训练周期下的性能。消融结果列在表 IV 中。以下,我们首先验证我们主要设计思想的有效性,然后分别对每个组件进行消融实验。
组件的消融: 在表 IV(a) 中,我们展示了每个组件,即对比学习、像素偏移数据增强和特征解码器对模型性能的影响。我们从 MAE 上对比学习的简单实现开始。具体来说,按照对比方法中输入生成方法,将随机裁剪区域与掩蔽输入到在线/动量编码器。优化第 III-C 节介绍的输出在线编码器和动量编码器之间的对比目标。从表 IV(a) 可以看出,这种直观的方法只带来了边际改进(0.2%)。显然,由于忽略了其与 MIM 的兼容性,对比学习的力量并未完全释放。通过使用我们提出的适度数据增强,即像素偏移,结果可以从 83.1% 增加到 83.6%,这证明了像素偏移的优势。此外,应用特征解码器进一步提高了模型的学习能力,将性能提高到 83.8%,证明了其在我们方法中的有效性。
对比损失: 为了探索 CMAE 中对比损失的影响,我们实验了各种损失权重,即公式 (8) 中的 λc\lambda_cλc。结果如表 IV(d) 所示。注意,当损失权重为 0 时,CMAE 退化为基线 MAE。当增加权重从 0 到 1 时,模型的性能相应增加,这验证了对比学习在增强学习表示中的重要性。当对比学习的权重大于 MIM 时,我们观察到训练不平衡的现象,这会对最终性能产生不利影响。这个实验表明,对比损失和重建损失对于学习有能力的表示都是至关重要的。因此,我们在所有实验中将 λc\lambda_cλc 设置为 1。
我们还进行了控制实验,比较了不同对比损失形式对预训练的影响。在相同的配置下,我们观察到使用 InfoNCE 损失训练的模型比 BYOL 风格损失(83.8% 对 83.4%)表现更好。这一结果表明,InfoNCE 在我们的方法中更有效地利用负样本。
像素偏移增强: 在本节中,我们消融了数据增强的重要性。与对比学习中应用大量数据增强的常见做法不同,我们发现适度的数据增强在对齐对比学习和 MIM 方面更有效。我们将数据增强方法分为两种:空间转移和颜色转移,并分别评估它们的效果。对于空间转移,我们将我们提出的像素偏移与常用的随机调整大小裁剪进行了比较。对于颜色转移,我们比较了两种情况,即动量分支是否使用颜色抖动。
从表 IV(b) 可以看出,像素偏移显著超过了随机裁剪(83.4% 对 83.0%)。像素偏移的优越性应归因于其生成更合理正样本视图的能力。如第 III-B 节所介绍,这一属性有助于对比学习在我们的框架内更好地与 MIM 协作。通过使用颜色转移,结果进一步提高到 83.8%,表明颜色转移对我们的方法是互补的。
我们研究了不同像素偏移范围的影响,通过改变最大允许偏移量。直观地说,更大的偏移范围会在两个增强视图之间引入更大的不对齐。如表 VI 所示,过度的偏移严重降低了模型性能,符合我们关于严重不对齐的正样本对对比学习有害的假设。结果表明,在视图多样性和对齐之间存在最佳平衡。基于不同偏移范围的观察性能,我们选择最大像素偏移为 31 作为默认参数设置,以最大化对比学习,同时保持足够的对齐。
特征解码器: 与现有工作不同,我们在进行对比学习时引入了一个特征解码器来恢复掩蔽补丁的特征。为了研究其有效性,我们在以下两种设置下进行了实验:特征解码器和像素解码器之间共享权重或不共享权重,以及改变特征解码器的深度。
在表 IV© 中,数字 “0” 表示不使用特征解码器,即仅使用在线编码器的输出,其中只包含可见标记的特征,用于对比学习。在这种设置下,我们的方法比使用轻量级两层特征解码器的性能更差。当增加特征解码器的深度时,对性能没有显著影响。然而,当深度增加到 8 时,我们得到了一个平凡解,可能是由于更深结构引起的优化困难。为了在效率和效果之间取得平衡,我们将深度设置为 2。此外,当特征解码器与像素解码器共享权重时,方法表现最差。一个可能的解释是,这两个分支有不同的目标,因此应该采用独立的权重。
动量分支的掩蔽比率: 在这个实验中,我们研究了动量分支掩蔽部分图像补丁是否会影响模型性能。按照之前的作品,我们选择了一组掩蔽比率,包括 {0, 0.25, 0.5, 0.65, 0.75} 用于动量分支,并查看性能如何变化。如图 IV(e) 所示,可以观察到使用完整的图像标记集获得最佳结果。可能的原因是:由于添加动量分支的目的是为模型提供对比监督,因此更倾向于包含图像的全部语义信息。否则,具有退化语义信息的掩蔽输入可能导致对比学习中的次优解。基于这一观察,我们模型中的动量分支在所有实验中都使用整个图像作为输入。
收敛速度: 为了进一步展示我们方法的有效性,我们比较了 CMAE 和 MAE 在 ImageNet-1k 上微调时的收敛行为。使用 1600 个周期预训练的权重作为初始化。如图 5 所示,我们观察到 CMAE 比 MAE 收敛得更快:仅用 55 个微调周期,CMAE 已经超越了 MAE 的最终性能。这一结果表明,CMAE 学习到的表示可以更容易地适应特定任务,这是一个与自监督预训练目的一致的吸引人的特性。
E. 部分微调和线性探测
在任务特定训练的背景下,部分微调 [30], [41], [58], [62] 和线性探测方法都保留了大部分模型组件的冻结状态。然而,关键区别在于被调整的头部的性质:部分微调使用非线性头部,而线性探测使用线性头部。正如 [30] 所强调的,鉴于线性探测与迁移学习性能的相关性最小,部分微调成为评估非线性但更强大的表示的优越协议。鉴于这些观察,我们的研究也强调了部分微调指标。
具体来说,我们按照 [30] 的实验设置来消融预训练 1600 个周期的 CMAE 基础模型。如图 4(a) 所示,我们模型的性能在所有测试设置中始终优于 MAE,例如,当微调一个块时,我们比 MAE 提高了 2.5%。以上结果表明,我们的模型可以有效地提高基线方法的表示质量。注意,当微调的块数为 “0” 时,它退化为线性探测。在这种情况下,我们的模型比 MAE 取得了显著的改进(5.9%)。这些结果表明,我们的方法能够在两种评估指标下提高表示质量。此外,与典型的对比模型 MoCo-V3 [14] 相比,MoCo-V3 在线性探测设置中表现出更优越的性能。然而,在部分微调设置下,CMAE 在所有方面都超过了 MoCo-V3,特别是在仅微调一个块时,CMAE 提高了 3.8%。这也证明了我们模型学习到的特征质量更高。
F. 模型扩展
为了研究我们方法在不同大小模型的可扩展性,我们采用 ViT-small, ViT-base, 和 ViT-large 作为编码器,并报告它们在 ImageNet-1k 微调上的性能。如图 4(b) 所示,CMAE 可以一致地提升 MAE 在所有规模上的性能。这些结果清楚地证明了 CMAE 的卓越可扩展性。
G. 特征分析
为了更有效地审查我们的模型获得的特征,我们使用 ViT-base 模型进行我们的研究。在预训练完成后,我们从 ImageNet-1k 验证集提取特征,并计算以下指标:
- 平均类内距离:该指标测量同一类别内所有图像对之间的平均距离。
- 类内距离的标准差:该指标测量同一类别内图像之间距离的变化。
- 平均类间距离:该指标测量不同类别中心之间的平均距离。
- 类间距离的标准差:该指标测量不同类别中心之间距离的变化。
当我们计算平均类内距离时,CMAE 实现了比 MAE 的 0.03800.03800.0380 更低的平均类内距离 0.03770.03770.0377。此外,如图 6(a) 所示,CMAE 的类内距离标准差小于 MAE(0.01890.01890.0189 对比 0.03710.03710.0371)。这些结果表明,CMAE 提取的特征在潜在空间中更紧密地聚集。关于类间距离,我们计算每个类别中心到其他类别中心的平均距离和距离的标准差。如图 6(b) 和 © 所示,CMAE 展示了更大的平均类间距离(0.03400.03400.0340 对比 0.03090.03090.0309)和更小的类间距离标准差(0.01520.01520.0152 对比 0.03100.03100.0310)。这表明,CMAE 提取的特征在潜在空间中每个类别的分布更均匀,并且类间距离更大。
总而言之,上述结果提供了证据,表明我们的模型能够学习到与 MAE 相比具有增强区分能力的优秀视觉表示。
V. 结论
本文介绍了一种名为对比掩码自编码器(CMAE)的新型自监督学习框架,旨在通过利用对比学习提高 MIM 的表示质量。在 CMAE 中,我们从输入生成和架构的角度提出了两种新设计,以协调 MIM 和对比学习。通过广泛的实验,证明了 CMAE 可以显著提高预训练中学习到的表示的质量。值得注意的是,在三个成熟的下游任务,即图像分类/分割/检测中,CMAE 实现了最先进的性能。未来,我们将探索将 CMAE 扩展到更大的数据集,并基于 CMAE 将图像密集的标题作为对比学习训练的另一个视图。