MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers

In this paper, we propose Mixed and Masked AutoEncoder (MixMAE), a simple but efficient pretraining method that is applicable to various hierarchical Vision Transformers. Existing masked image modeling (MIM) methods for hierarchical Vision Transformers replace a random subset of input tokens with a special [MASK] symbol and aim at reconstructing original image tokens from the corrupted image. However, we find that using the [MASK] symbol greatly slows down the training and causes pretraining-finetuning inconsistency, due to the large masking ratio (e.g., 60% in SimMIM). On the other hand, MAE does not introduce [MASK] tokens at its encoder at all but is not applicable for hierarchical Vision Transformers. To solve the issue and accelerate the pretraining of hierarchical models, we replace the masked tokens of one image with visible tokens of another image, i.e., creating a mixed image. We then conduct dual reconstruction to reconstruct the two original images from the mixed input, which significantly improves efficiency. While MixMAE can be applied to various hierarchical Transformers, this paper explores using Swin Transformer with a large window size and scales up to huge model size (to reach 600M parameters). Empirical results demonstrate that MixMAE can learn high-quality visual representations efficiently. Notably, MixMAE with Swin-B/W14 achieves 85.1% top-1 accuracy on ImageNet-1K by pretraining for 600 epochs. Besides, its transfer performances on the other 6 datasets show that MixMAE has better FLOPs / performance tradeoff than previous popular MIM methods. Code is available at https://github.com/Sense-X/MixMIM.

PDF Abstract CVPR 2023 PDF CVPR 2023 Abstract

Results from the Paper


Task Dataset Model Metric Name Metric Value Global Rank Uses Extra
Training Data
Result Benchmark
Semantic Segmentation ADE20K val MixMIM-L mIoU 53.8 # 35
Semantic Segmentation ADE20K val MixMIM-B mIoU 50.3 # 47
Object Detection COCO 2017 MixMIM-B mAP 52.2 # 7
Object Detection COCO 2017 MixMIM-L mAP 54.1 # 5
Image Classification ImageNet MixMIM-B Top 1 Accuracy 85.1% # 262
Number of params 88M # 866
GFLOPs 16.3 # 376
Image Classification iNaturalist 2018 MixMIM-B Top-1 Accuracy 77.5% # 23
Image Classification iNaturalist 2018 MixMIM-L Top-1 Accuracy 80.3% # 17
Image Classification iNaturalist 2019 MixMIM-L Top-1 Accuracy 83.9 # 4
Image Classification Places205 MixMIM-B Top 1 Accuracy 68.3 # 4
Image Classification Places205 MixMIM-L Top 1 Accuracy 69.3 # 2
Image Classification Places365 MixMIM-L(ViT-L) Top 1 Accuracy 60.3 # 4
Image Classification Places365 MixMIM-B (ViT) Top 1 Accuracy 58.9 # 7

Methods