探究掩码图像建模中的数据可扩展性.
掩码图像建模 (masked image modeling, MIM) 被怀疑无法从更大的数据中获益。这项工作中作者通过大量实验打破了这种误解,表明该方法不仅仅是模型可扩展的,也是数据可扩展的。数据规模从 ImageNet-1K 的 $10\%$ 到完整 ImageNet-22K,模型大小从4900万到10亿,训练长度从 125K 个 iteration 到 500K 个 iteration。作者的研究表明:
- 掩码图像建模也需要更大的数据:作者观察到,非常大的模型在相对较小的数据下往往会过拟合。
- 训练时间的问题:通过掩码图像建模训练的大模型,可以通过更长的训练,从更多的数据中受益。
- 预训练中的验证集损失是衡量模型在多个任务上的微调效果的良好指标:这种观察能够对预训练的模型进行廉价的预评估,而不必每次都在下游任务上微调评估。
掩码图像建模以被破坏的图像为输入,训练的目标是被掩码的区域。在本研究中,作者使用 SimMIM 作为默认的掩码图像建模方法,因为它简单且对视觉编码器的架构没有限制。SimMIM 由一个视觉编码器和一个极其轻量级的预测头 (一个线性层) 组成,通过L1回归损失预测损坏图像的原始像素。为了便于 ViT 的实现,SimMIM 的掩码采用了 patch-wise mask 的策略,掩码的 patch 大小为 32×32,掩码比为0.6。
本研究使用 Swin Transformer V2 作为视觉编码器。由于其通用性和可扩展性,作者在多个下游任务上评估了一系列模型大小的 SwinV2 模型 (参数数量范围从 ∼50M 到 ∼1B, FLOPs 范围从 ∼9G 到 ∼190G)。详细的模型规格如下图所示。作者使用了一种新的变体 SwinV2-g (giant),其参数数量介于 SwinV2-L 和30亿参数的 SwinV2-G (Giant) 之间。
为了研究数据大小对掩码图像建模的影响,作者构建了不同大小的数据集,使用 ImageNet-1K 和 ImageNet-22K 的训练集作为两个大规模的数据集,并随机抽取 ImageNet-1K 训练集中 10%、20%、50% 的图像作为较小的数据集。本研究中使用的所有预训练数据集的细节和统计如下图所示。
下图是不同模型和数据规模下的训练损失、验证损失与 ImageNet-1K 微调精度和训练长度之间的关系。随着训练成本的增加,部分模型的训练损失显著下降,但是验证损失显著上升,即使使用 ImageNet-1K 的 50% 的图像,也存在过拟合现象。
下图给出了每个模型的最佳微调性能。可以发现,当使用小数据集进行训练时,大型模型的表现甚至不如小型模型。使用更多的数据可以获得更好的性能。这些观察结果表明 MIM 的预训练方法也是需要大数据集的。
通过比较不同数据大小预训练的模型的性能,发现与较小的模型相比,大模型的微调性能随着数据大小的增加而饱和得更慢。而且在较短的训练长度下,使用更多数据对预训练的改善并不显著。这一观察表明,虽然更大的模型可以从更多的数据中受益,但训练长度也必须同时增加。
除了 ImageNet-1K 图像分类,作者还评估了 MIM 预训练的 SwinV2-S, SwinV2-B 和 SwinV2-L 在 iNaturalist 2018 细粒度图像分类,ADE20K 语义分割和 COCO 目标检测和实例分割上的效果。如下图所示,随着训练成本的增加,一些模型显示出过拟合的迹象。随着数据量的增加,较小的模型迅速达到饱和,而较大的模型在经过足够的训练后,可以从更多的数据中持续受益。这些结果表明,在 ImageNet-1K 上得出的结论广泛适用于其他视觉任务。
作者还探索了 MIM 训练中的预训练的损失是否是其微调性能的良好指标。对于非过拟合模型 (绿圈),训练损失和验证损失与下游任务的微调性能之间的相关性都为负;对于过拟合模型 (红圈),训练损失与下游任务的微调性能之间的相关性为正,验证损失与下游任务的微调性能之间的相关性为负。以上结果说明在多个任务上,预训练中的验证损失是衡量模型在多个任务上的微调效果的良好指标。