[论文解读] Masked Image Residual Learning for Scaling Deeper Vision Transformers

Guoxi Huang, Hongtao Fu+1 more|arXiv (Cornell University)|Sep 25, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出掩码图像残差学习（MIRL），一种自监督框架，通过训练深层网络重建图像残差而非完整图像，重新构想了视觉变换器（ViTs）的预训练方式，缓解了掩码图像建模（MIM）中的性能退化问题。MIRL使更深的ViTs（如ViT-S-54和ViT-B-48）在ImageNet上达到与更大模型（如ViT-Large）相当或更优的性能，Top-1准确率达到86.2%，并在下游任务中展现出更优的泛化能力。

ABSTRACT

Deeper Vision Transformers (ViTs) are more challenging to train. We expose a degradation problem in deeper layers of ViT when using masked image modeling (MIM) for pre-training. To ease the training of deeper ViTs, we introduce a self-supervised learning framework called Masked Image Residual Learning (MIRL), which significantly alleviates the degradation problem, making scaling ViT along depth a promising direction for performance upgrade. We reformulate the pre-training objective for deeper layers of ViT as learning to recover the residual of the masked image. We provide extensive empirical evidence showing that deeper ViTs can be effectively optimized using MIRL and easily gain accuracy from increased depth. With the same level of computational complexity as ViT-Base and ViT-Large, we instantiate 4.5$ imes$ and 2$ imes$ deeper ViTs, dubbed ViT-S-54 and ViT-B-48. The deeper ViT-S-54, costing 3$ imes$ less than ViT-Large, achieves performance on par with ViT-Large. ViT-B-48 achieves 86.2% top-1 accuracy on ImageNet. On one hand, deeper ViTs pre-trained with MIRL exhibit excellent generalization capabilities on downstream tasks, such as object detection and semantic segmentation. On the other hand, MIRL demonstrates high pre-training efficiency. With less pre-training time, MIRL yields competitive performance compared to other approaches.

研究动机与目标

• 解决在使用掩码图像建模（MIM）进行预训练时，视觉变换器（ViTs）深层网络性能下降的问题。
• 识别出MIM在深层ViT模块中施加了负面优化效应，阻碍了深度扩展。
• 开发一种新颖的自监督学习框架，将深层网络的预训练目标从图像重建转变为残差重建。
• 实现深层ViTs的有效训练，提升泛化能力与效率，以更低的计算成本达到最先进性能。
• 证明MIRL与现有特征级损失方法正交且互补，可进一步提升性能。

提出的方法

• 将ViT编码器划分为浅层与深层两部分，每部分配备专用解码器以实现多解码。
• 重新构想深层网络的预训练目标，使其预测被掩码图像块的残差，而非完整图像块。
• 采用多解码架构，其中浅层部分重建图像主要成分，深层部分学习残差。
• 在浅层部分应用标准的掩码图像建模（MIM）目标，在深层部分应用残差重建目标。
• 通过将编码器划分为更多深度层级，将框架扩展至多段，支持极深ViTs的训练。
• 集成额外的特征级目标（如VGG损失）以进一步提升性能，表明其与MIRL具有正交性。

实验结果

研究问题

• RQ1 为何在NLP与视觉任务中均取得成功的标准掩码图像建模（MIM）方法，在预训练深层ViT时仍导致其性能下降？
• RQ2 是否可通过将深层ViT层的预训练目标从图像重建改为残差重建，缓解性能退化问题？
• RQ3 MIRL在保持或提升性能与效率的前提下，能在多大程度上实现ViTs的深度扩展？
• RQ4 与现有基于MIM的方法及特征级损失方法相比，MIRL在预训练效率与下游迁移性能方面表现如何？
• RQ5 MIRL带来的性能提升是否与其它正则化技术（如特征级损失）正交？

主要发现

• MIRL有效缓解了标准MIM在深层ViT层中引发的性能退化问题，实现了深层架构的稳定训练。
• ViT-S-54（54层）比ViT-Base深4.5倍，且比ViT-Large便宜3倍，其在ImageNet上的性能与ViT-Large相当。
• ViT-B-48（48层）比ViT-Base深2倍，在ImageNet上达到86.2%的Top-1准确率，优于在相似计算预算下的标准ViT-Large。
• MIRL预训练模型在下游任务中泛化能力更强，在ADE20K语义分割任务上达到53.2%的mIoU，与使用更多数据或更长训练周期的模型性能相当。
• MIRL框架极为高效：其性能与基线方法（如MAE和BEiT）相比，显著减少了预训练时间。
• 梯度分析表明，MIRL增强了深层网络中的梯度流动，其梯度范数高于原始MIM方法，表明训练动态得到改善。