[Paper Review] Improving Pixel-based MIM by Reducing Wasted Modeling Capability

Yuan Liu, Songyang Zhang+4 more|arXiv (Cornell University)|Aug 1, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Multi-level Feature Fusion (MFF) to reduce wasted modeling capacity in pixel-based Masked Image Modeling (MIM) by explicitly incorporating low-level features from shallow layers of Vision Transformers. By fusing features from multiple encoder layers using a learnable, dynamic weight-average strategy, MFF improves training efficiency and downstream performance, achieving state-of-the-art results on ViT-S with gains of 2.8% on linear probing, 2.6% on semantic segmentation, and 1.2% on fine-tuning.

ABSTRACT

There has been significant progress in Masked Image Modeling (MIM). Existing MIM methods can be broadly categorized into two groups based on the reconstruction target: pixel-based and tokenizer-based approaches. The former offers a simpler pipeline and lower computational cost, but it is known to be biased toward high-frequency details. In this paper, we provide a set of empirical studies to confirm this limitation of pixel-based MIM and propose a new method that explicitly utilizes low-level features from shallow layers to aid pixel reconstruction. By incorporating this design into our base method, MAE, we reduce the wasted modeling capability of pixel-based MIM, improving its convergence and achieving non-trivial improvements across various downstream tasks. To the best of our knowledge, we are the first to systematically investigate multi-level feature fusion for isotropic architectures like the standard Vision Transformer (ViT). Notably, when applied to a smaller model (e.g., ViT-S), our method yields significant performance gains, such as 1.2\% on fine-tuning, 2.8\% on linear probing, and 2.6\% on semantic segmentation. Code and models are available at https://github.com/open-mmlab/mmpretrain.

Motivation & Objective

• To address the inherent bias of pixel-based MIM methods like MAE toward high-frequency details and low-level features.
• To reduce wasted modeling capacity in pixel reconstruction by leveraging multi-level features from shallow layers.
• To systematically investigate multi-level feature fusion in isotropic architectures such as Vision Transformers.
• To improve downstream performance and training efficiency without increasing model complexity.

Proposed method

• Proposes Multi-level Feature Fusion (MFF), a drop-in solution that fuses features from the output layer and selected shallow/intermediate layers of a Vision Transformer encoder.
• Uses a learnable, dynamic weight-average pooling strategy to combine features, where layer weights are updated during pre-training based on reconstruction importance.
• Applies a simple linear projection layer before fusion to align feature distributions across layers, avoiding domain gap issues.
• Employs a lightweight fusion mechanism that adds minimal computational overhead compared to standard MAE.
• Selects shallow layers (e.g., Layer 0) as primary contributors due to their high-frequency content, validated through frequency analysis.
• Extends the method to PixMIM, demonstrating generalization across different pixel-based MIM frameworks.

Experimental results

Research questions

• RQ1 Why do pixel-based MIM methods like MAE underutilize high-level semantic representations?
• RQ2 How does multi-level feature fusion affect the learning dynamics and loss landscape in Vision Transformers?
• RQ3 What is the optimal set of layers and fusion strategy for improving pixel reconstruction and downstream performance?
• RQ4 Can fusing shallow features reduce the model's over-reliance on low-level details without sacrificing reconstruction fidelity?
• RQ5 Does multi-level fusion improve generalization and robustness on out-of-distribution data?

Key findings

• Fusing features from shallow layers (e.g., Layer 0) with the output layer improves linear probing accuracy by 2.8% on ViT-S, compared to standard MAE.
• Semantic segmentation accuracy improves by 2.6% on ViT-S, and fine-tuning accuracy increases by 1.2% with the same model.
• Pre-training with MFF achieves performance equivalent to 1600 epochs of standard MAE in only 300 epochs, reducing training time by ~5x.
• The method reduces high-frequency information in the learned features and flattens the loss landscape, improving optimization stability.
• Out-of-distribution robustness improves, with MFFMAE achieving 37.8% top-1 accuracy on ImageNet-S, a 2.3% gain over MAE.
• Ablation studies confirm that linear projection and weight-average pooling yield the best trade-off between performance and efficiency, outperforming nonlinear projections and attention-based fusion.