[Paper Review] Masked Image Residual Learning for Scaling Deeper Vision Transformers

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

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Masked Image Residual Learning (MIRL), a self-supervised framework that reformulates pre-training for deeper Vision Transformers (ViTs) by training deeper layers to reconstruct image residuals rather than full images, alleviating a degradation problem in masked image modeling (MIM). MIRL enables deeper ViTs—such as ViT-S-54 and ViT-B-48—to achieve performance on par with or exceeding larger models like ViT-Large, with 86.2% top-1 accuracy on ImageNet and superior downstream generalization.

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.

Motivation & Objective

• Address the performance degradation observed in deeper layers of Vision Transformers (ViTs) when pre-trained with masked image modeling (MIM).
• Identify that MIM imposes negative optimization effects on deeper ViT blocks, hindering depth scaling.
• Develop a novel self-supervised learning framework that shifts the pre-training objective of deeper layers from image reconstruction to residual reconstruction.
• Enable effective training of deeper ViTs with improved generalization and efficiency, achieving state-of-the-art performance at lower computational cost.
• Demonstrate that MIRL is orthogonal to and complementary with existing feature-level loss methods, further boosting performance.

Proposed method

• Segment the ViT encoder into shallow and deep parts, each with a dedicated decoder for multi-decoding.
• Reformulate the pre-training objective for deeper layers to predict the residual of masked image patches, rather than the full patch.
• Use a multi-decoding architecture where the shallow part reconstructs the main image component and the deep part learns the residual.
• Apply a standard masked image modeling (MIM) objective on the shallow segment and a residual reconstruction objective on the deep segment.
• Extend the framework to multiple segments by dividing the encoder into more depth levels, enabling training of very deep ViTs.
• Integrate additional feature-level objectives (e.g., VGG loss) to further improve performance, showing orthogonality to MIRL.

Experimental results

Research questions

• RQ1 Why do deeper layers of ViTs underperform when pre-trained with standard masked image modeling (MIM), despite MIM's success in NLP and vision?
• RQ2 Can reformulating the pre-training objective for deeper ViT layers from image reconstruction to residual reconstruction alleviate performance degradation?
• RQ3 To what extent can MIRL enable scaling ViTs along depth while maintaining or improving performance and efficiency?
• RQ4 How does MIRL compare in pre-training efficiency and downstream transfer performance to existing MIM-based methods and feature-level loss approaches?
• RQ5 Is the improvement from MIRL orthogonal to other regularization techniques, such as feature-level losses?

Key findings

• MIRL effectively alleviates the degradation problem in deeper ViT layers caused by standard MIM, enabling stable training of deeper architectures.
• ViT-S-54, a 54-block model 4.5× deeper than ViT-Base and 3× cheaper than ViT-Large, achieves performance on par with ViT-Large on ImageNet.
• ViT-B-48, 2× deeper than ViT-Base, achieves 86.2% top-1 accuracy on ImageNet, outperforming standard ViT-Large under similar computational budgets.
• MIRL-pretrained models generalize better on downstream tasks, achieving 53.2% mIoU on ADE20K semantic segmentation, comparable to models pre-trained with more data or longer schedules.
• The MIRL framework is highly efficient: it achieves competitive performance with significantly less pre-training time than baseline methods like MAE and BEiT.
• Gradient analysis shows MIRL enhances gradient flow in deeper layers, with higher gradient norms compared to vanilla MIM, suggesting improved training dynamics.