[论文解读] DeepMIM: Deep Supervision for Masked Image Modeling

Sucheng Ren, Fangyun Wei+3 more|arXiv (Cornell University)|Mar 15, 2023

Advanced Neural Network Applications被引用 11

一句话总结

DeepMIM 重新研究遮罩图像建模（MIM）中的深度监督，并展示在中间 ViT 块中添加轻量解码器可提升表示学习、收敛性和下游性能，适用于 MAE 与基于 CLIP 的设置。

ABSTRACT

Deep supervision, which involves extra supervisions to the intermediate features of a neural network, was widely used in image classification in the early deep learning era since it significantly reduces the training difficulty and eases the optimization like avoiding gradient vanish over the vanilla training. Nevertheless, with the emergence of normalization techniques and residual connection, deep supervision in image classification was gradually phased out. In this paper, we revisit deep supervision for masked image modeling (MIM) that pre-trains a Vision Transformer (ViT) via a mask-and-predict scheme. Experimentally, we find that deep supervision drives the shallower layers to learn more meaningful representations, accelerates model convergence, and expands attention diversities. Our approach, called DeepMIM, significantly boosts the representation capability of each layer. In addition, DeepMIM is compatible with many MIM models across a range of reconstruction targets. For instance, using ViT-B, DeepMIM on MAE achieves 84.2 top-1 accuracy on ImageNet, outperforming MAE by +0.6. By combining DeepMIM with a stronger tokenizer CLIP, our model achieves state-of-the-art performance on various downstream tasks, including image classification (85.6 top-1 accuracy on ImageNet-1K, outperforming MAE-CLIP by +0.8), object detection (52.8 APbox on COCO) and semantic segmentation (53.1 mIoU on ADE20K). Code and models are available at https://github.com/OliverRensu/DeepMIM.

研究动机与目标

• 在自监督的遮罩图像建模（MIM）背景下，推动重新审视深度监督。
• 证明深度监督可改善浅层表示、收敛性和注意力多样性。
• 展示 DeepMIM 与多种 MIM 目标及分词器选择的兼容性。
• 用多种预训练配置量化在 ImageNet 分类、对象检测和语义分割上的改进。

提出的方法

• 在 ViT-B 编码器的中间 ViT 块（第 6、8、10 层）附加三个轻量解码器，以在 MIM 预训练期间启用深度监督。
• 可选地使用逐步混合目标生成器，将原始图像与 MAE 重构结果混合，作为中间块的目标。
• 以所有解码器的重构损失之和加上最终解码器损失进行训练，确保在多个深度处获得监督。
• 展示 DeepMIM 在预训练计划（300 与 1600 轮）中与 MAE、MAE 变体以及基于 CLIP 的目标的兼容性。
• 通过重构损失、不同层之间的 CKA 相似性以及注意力头多样性来评估表示。

实验结果

研究问题

• RQ1 在 MIM 预训练中，对中间 ViT 块应用重构监督是否能提升表示学习？
• RQ2 深度监督、混合目标和目标类型如何影响重构损失、特征相似性以及下游迁移性能？
• RQ3 DeepMIM 是否与多样的重构目标（像素、HOG、DINO 特征、CLIP 等）以及 MIM 框架兼容？
• RQ4 DeepMIM 对下游任务的影响如何，例如 ImageNet 分类、COCO 检测、ADE20K 分割和 Kinetics 视频分类？

主要发现

• DeepMIM 在多种目标和计划上持续优于 MAE 的性能（例如，ViT-B/MAE 下在 ImageNet 上提升 +0.8 top-1）。
• DeepMIM 相较于标准 MAE，获得更低的训练和验证重构损失，表明优化更强。
• CKA 分析表明 DeepMIM 的中间块变得更具判别力，并且与最终表示对齐得更好。
• DeepMIM 增加了注意力头多样性，表明浅层块中出现更丰富的多头表示。
• 当与像 CLIP 这样的强分词器结合时，DeepMIM 在 ImageNet 分类、COCO 检测和 ADE20K 分割上达到最先进的结果（例如，使用 DeepMIM-MAE-CLIP 在 ImageNet-1K 得到 85.6 的 top-1）。
• 与 MAE 相比，DeepMIM 提高了在 ImageNet-A、ImageNet-R 和 ImageNet-C 上的鲁棒性与域外性能。