PaperlayerView exact revision on arXiv

Source-linked AI summary

SimMIM: A Simple Framework for Masked Image Modeling

Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, Han Hu

arXiv:2111.09886v2cs.CV

TL;DR

Masked image modeling in vision must contend with strong pixel locality and low-level continuous signals, motivating the question of whether specialized designs are necessary for transferable representations. SimMIM uses random patch masking, raw-pixel regression, and a lightweight prediction head, achieving strong ImageNet-1K results and scaling to a 3B model evaluated across four vision benchmarks.

  • Problem

    Vision masked modeling faces modality-specific challenges because nearby pixels are correlated and visual signals are raw, unlike high-level text tokens.

  • Method

    SimMIM randomly masks image patches, regresses masked raw RGB pixels with an ℓ1 loss, and uses a prediction head as light as a linear layer.

  • Results

    87.1% top-1 accuracy on ImageNet-1K is achieved by SwinV2-H using ImageNet-1K data only, while a 3B SwinV2-G uses ∼40× smaller data than JFT-3B across four benchmarks.

  • Takeaways & Limitations

    Simple masked image modeling components can provide strong representation learning and facilitate scaling to larger vision models across multiple tasks.

Abstract

from arXiv · show

This paper presents SimMIM, a simple framework for masked image modeling. We simplify recently proposed related approaches without special designs such as block-wise masking and tokenization via discrete VAE or clustering. To study what let the masked image modeling task learn good representations, we systematically study the major components in our framework, and find that simple designs of each component have revealed very strong representation learning performance: 1) random masking of the input image with a moderately large masked patch size (e.g., 32) makes a strong pre-text task; 2) predicting raw pixels of RGB values by direct regression performs no worse than the patch classification approaches with complex designs; 3) the prediction head can be as light as a linear layer, with no worse performance than heavier ones. Using ViT-B, our approach achieves 83.8% top-1 fine-tuning accuracy on ImageNet-1K by pre-training also on this dataset, surpassing previous best approach by +0.6%. When applied on a larger model of about 650 million parameters, SwinV2-H, it achieves 87.1% top-1 accuracy on ImageNet-1K using only ImageNet-1K data. We also leverage this approach to facilitate the training of a 3B model (SwinV2-G), that by $40\times$ less data than that in previous practice, we achieve the state-of-the-art on four representative vision benchmarks. The code and models will be publicly available at https://github.com/microsoft/SimMIM.

1. Introduction

SimMIM asks whether masked image modeling can learn transferable visual representations without the specialized masking, tokenization, and prediction designs used by earlier approaches. It finds that random masking, raw-pixel regression, and a lightweight head provide strong representation learning and scale effectively.

  • Motivation: Visual locality makes masked prediction difficult because nearby pixels can be duplicated rather than inferred through semantic reasoning.Images contain strongly correlated nearby pixels, while visual signals are raw and low-level rather than human-generated concepts.
  • Prior approaches: Earlier masked-image methods introduced color clustering, additional-network patch tokenization, or block-wise masking to address modality-specific obstacles.These designs produced representations transferable to several visual recognition tasks.
  • SimMIM: SimMIM randomly masks input patches and uses a linear layer to regress masked raw pixels with an ℓ1 loss.The framework is designed to align with the continuous nature of visual signals without complex tokenization or masking mechanisms.
  • Design findings: A masked patch size of 32 supports competitive performance across masking ratios from 10%-70%, whereas an 8-sized mask patch requires an 80% ratio to perform well.Larger patches or higher masking ratios reduce the chance of finding visible nearby pixels.
  • Design findings: Heavier prediction heads or higher target resolutions can improve generation capability without necessarily improving downstream fine-tuning.A lightweight head also brings a remarkable pre-training speedup.

2. Related Work

Masked image modeling developed alongside masked language modeling but remained less central in vision, where visual locality and low-level signals complicate direct transfer of the language-domain approach. SimMIM contrasts with prior vision methods by avoiding clustering, tokenization, and other specialized components.

  • Masked language modeling: Masked language modeling became a dominant NLP self-supervised approach by predicting invisible tokens from visible context and supporting large language models.The approach generalizes across language understanding and generation tasks using large unlabeled datasets.
  • Masked image modeling: Masked image modeling progressed in parallel but remained non-mainstream in computer vision for a long period.Recent vision Transformer methods renewed interest by adding specialized designs to masked prediction.
  • Prior masked-image methods: Prior masked-image methods used pixel clustering, mean-color prediction, or dVAE tokenization with block-wise masking.BEiT additionally requires pre-training a discrete VAE network to create classification targets.
  • Related reconstruction methods: Reconstruction-based methods recover visible original signals and study regularization or architectural bottlenecks, unlike masked prediction of invisible signals.The paper positions these methods as following a different learning philosophy.
  • Image inpainting: Image inpainting seeks to improve masked-image reconstruction quality, whereas this paper advocates masked image modeling as a self-supervised pretext task for representation learning.The paper states that stronger inpainting capability does not necessarily produce stronger downstream fine-tuning performance.
  • Other self-supervised approaches: Other visual self-supervised tasks include colorization, jigsaw solving, split-brain auto-encoding, rotation prediction, and clustering.Some also predict invisible signal components, such as color channels.

3. Approach

SimMIM masks image signals, encodes the transformed image, and predicts the original masked content using simple components. The approach emphasizes patch-aligned random masking, raw-pixel regression with ℓ1 loss, and a prediction head as light as a linear layer, while evaluating representations mainly through fine-tuning.

  • Framework: SimMIM masks part of an input image and predicts the original signals at the masked area through an encoder-based framework.Its major components include masking strategy, encoder architecture, prediction head, and prediction target.
  • Masking strategy: Patch-aligned random masking makes each image-processing patch either fully visible or fully masked.The default masked patch size is 32×32 for both Swin Transformer and ViT.
  • Masking strategy: Figure 2 compares square masking, block-wise masking on 16-sized patches, and random masking with patch sizes from 4 to 32 at a 0.6 mask ratio.The comparison visualizes the masked areas generated by these strategies.
  • Prediction head: The prediction head can be as lightweight as a linear layer, alongside heavier alternatives such as a 2-layer MLP or inverse Swin models.The head maps encoder outputs to the prediction target.
  • Prediction target: Raw-pixel regression predicts continuous RGB values for the masked area rather than classifying clustered or tokenized targets.Feature vectors are mapped back to the input resolution to predict pixel values.
  • Prediction target: An ℓ1 loss averages absolute differences between predicted and input RGB values over masked pixels.The paper also evaluates ℓ2 and smooth-ℓ1 losses, which perform similarly well, while adopting ℓ1 by default.
  • Evaluation protocols: Representations are evaluated mainly by ImageNet-1K fine-tuning because it is considered a more usable practical scenario.Linear probing is reported for system-level comparison but is not the primary metric.

4. Experiments

Experiments show that SimMIM’s simple design choices produce strong transferable representations across masking, prediction, and scaling settings. Performance remains strong with random masking, raw-pixel regression, lightweight heads, moderate target resolutions, and masked-area prediction.

  • Masking Strategy: 83.0% top-1 accuracy is achieved by random masking, exceeding block-wise and other specialized strategies by +0.3%.
  • Masking Strategy: Random masking with masked patch size 32 performs stably across masking ratios from 10% to 70%, while patch size 64 degrades accuracy.Smaller patch sizes require higher masking ratios, but remain below the overall accuracy achieved with patch size 32.
  • Masking Strategy: Fine-tuning accuracy is highest at a moderate AvgDist range of 10 to 20, suggesting that prediction distance should be neither too small nor too large.Small distances may emphasize short connections, whereas large distances may be difficult to learn.
  • Prediction Head: A linear prediction head transfers better than heavier heads despite their lower training losses, while also reducing pre-training cost.The results indicate that stronger inpainting capability does not necessarily improve downstream performance.
  • Prediction Loss Area: 82.8% accuracy from masked-area prediction exceeds 81.7% from reconstructing all image pixels.The reconstruction task can produce better-looking images, but capacity spent on unmasked-area recovery may not benefit fine-tuning.
  • Scaling Experiments: 87.1% top-1 accuracy is reached on ImageNet-1K with SwinV2-H using ImageNet-1K data only, while a 3B SwinV2-G model uses approximately 40× less data than JFT-3B across four benchmarks.The 3B model reports 84.0% on ImageNet-V2, 63.1/54.4 box/mask mAP on COCO, 59.9 mIoU on ADE20K, and 86.8% on Kinetics-400.

5. Conclusion

SimMIM is a simple yet effective self-supervised framework for masked image modeling. It combines random masking, raw-pixel regression, and a lightweight prediction head.

  • SimMIM uses random masking with moderately large masked patches to predict raw RGB pixels by direct regression.
  • A linear-layer prediction head is sufficient for the framework.

A. Detailed Architectures

The detailed architecture specifications use 192 × 192 inputs for pre-training and 224 × 224 inputs for fine-tuning.

  • Pre-training uses input images sized 192 × 192, while fine-tuning uses 224 × 224 inputs.

B. The Effect of Learning Rate Schedulers

The experiments compare cosine and step learning-rate schedulers under otherwise similar settings. The step scheduler performs marginally better in fine-tuning.

  • The ablation study uses a cosine learning-rate scheduler, while scaling experiments use a step scheduler to reduce experimental overhead.
  • The step scheduler performs marginally better than the cosine scheduler in fine-tuning accuracy.

C. Results on Downstream Tasks

The downstream evaluation covers iNaturalist 2018 classification, COCO object detection, and ADE20K semantic segmentation. These experiments use task-specific frameworks and training settings, with ADE20K initialized from supervised ImageNet-1K fine-tuning weights.

  • The additional downstream evaluations include iNaturalist 2018 classification, COCO object detection, and ADE20K semantic segmentation.
  • iNaturalist 2018 is a long-tail classification dataset with more than 8,000 categories and 437,513 training images.
  • COCO object detection uses a Mask-RCNN framework with a 3× training schedule of 36 epochs.
  • ADE20K semantic segmentation uses a UPerNet framework and trains models for 80K iterations at 512×512 input resolution.
  • ADE20K segmentation models are initialized from supervised ImageNet-1K fine-tuned weights because they outperform direct initialization from self-supervised weights.

C.2. Ablation Studies

Ablations show that lightweight prediction heads and raw-pixel regression remain competitive across downstream tasks, while masking effectiveness is associated with AvgDist.

  • Prediction target: Raw-pixel regression achieves on-par or better performance than classification-based prediction targets in the ablations.The prediction-target comparison covers iNat-2018, COCO, and ADE20K.
  • Masking strategy: An AvgDist of 15 is empirically good for masked image modeling, and AvgDist is a good indicator of learning effectiveness.These observations are reported across ImageNet, iNat-2018, COCO, and ADE20K.
  • Scaling experiments: Scaling experiments evaluate SimMIM on COCO object detection and ADE20K semantic segmentation using Swin architectures.The experiments use Swin-B, Swin-L, and SwinV2-H.

C.3. Scaling Experiments

Scaling experiments show that SimMIM improves detection and segmentation over supervised counterparts across Swin model sizes, with larger models benefiting more.

  • Scaling results: +2.1 / +2.9 / +4.2 mAPbox and +2.4 / +3.5 / +4.4 mIoU are achieved on Swin-B, Swin-L, and SwinV2-H, respectively.The gains are reported for COCO object detection and ADE20K semantic segmentation.
  • Scaling results: The scaling results indicate broad effectiveness of SimMIM and suggest that larger models benefit more from the approach.

D. More Results on Channel-wise Bin Color Discretization

Channel-wise bin color discretization shows that the best prediction resolution depends on the number of color bins, with moderately fine-grained targets favored.

  • Resolution by bin count: The 2-bin and 4-bin targets achieve their best accuracy at 192^2 prediction resolution.
  • Resolution by bin count: All other evaluated bin numbers achieve their best accuracy at the lower 6^2 prediction resolution.
  • Target granularity: The results imply that a moderately fine-grained prediction target is encouraged for this discretization approach.

E. SimMIM with ConvNets

SimMIM also transfers to a ResNet-50×4 ConvNet setting, where masked tokens replace post-stem features and pretraining improves ImageNet-1K accuracy over supervision.

  • Architecture and setup: SimMIM uses ResNet-50×4 with masked tokens replacing features after a stride-2 3 × 3 convolution and 2×2 max-pooling stem.The overall training setup remains the same as for Swin-Base.
  • ImageNet-1K result: 81.6% top-1 accuracy is achieved on the ImageNet-1K validation set after 300-epoch pretraining and 100-epoch fine-tuning.
  • ImageNet-1K result: +0.9% over the supervised counterpart is achieved, comparing 81.6% with 80.7%.The result is presented as evidence of SimMIM’s generality.
Loading 2111.09886v2…