[Paper Review] MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers

Jihao Liu, Xin Huang+3 more|arXiv (Cornell University)|May 26, 2022

Advanced Neural Network Applications7 citations

TL;DR

MixMAE proposes a novel masked autoencoding framework for efficient pretraining of hierarchical Vision Transformers by mixing tokens from two images to replace masked patches, enabling dual reconstruction without [MASK] tokens. It achieves 85.1% top-1 accuracy on ImageNet-1K with only 600 pretraining epochs, outperforming prior MIM methods in FLOPs/performance tradeoff.

ABSTRACT

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.

Motivation & Objective

• To address pretraining inefficiency and pretrain-finetune inconsistency in masked image modeling (MIM) for hierarchical Vision Transformers.
• To eliminate the use of [MASK] tokens in the encoder, which slows training and causes distributional mismatch with fine-tuning.
• To develop a method compatible with hierarchical ViTs like Swin Transformer that maintains high representation quality while improving training efficiency.
• To achieve better FLOPs/performance tradeoff compared to existing MIM methods such as SimMIM and MAE.

Proposed method

• MixMAE creates a mixed input by replacing masked patches in one image with visible patches from another image, avoiding the use of [MASK] tokens in the encoder.
• The hierarchical Vision Transformer encoder processes the mixed input and produces hidden representations for both original images.
• Before decoding, the hidden representations are unmixed and filled with [MASK] tokens to enable dual reconstruction.
• A lightweight decoder reconstructs both original images from the mixed representations, using a dual reconstruction loss to improve feature quality.
• The method supports variable masking ratios by mixing multiple images, with 75% masking ratio (four images) yielding optimal performance.
• The approach uses masked self-attention in the decoder and avoids [MASK] tokens in the encoder, ensuring pretrain-finetune consistency.

Figure 1 : Overview of MixMAE. For pretraining, two images are mixed with a random mixing mask to create a mixed image. MixMAE takes the mixed image as input and reconstructs the two original images. Right before decoding, the token embeddings are unmixed and filled with mask tokens for dual reconst

Experimental results

Research questions

• RQ1 Can a masked autoencoding approach be designed for hierarchical Vision Transformers that avoids [MASK] tokens in the encoder to improve training efficiency and consistency?
• RQ2 How does replacing masked patches with patches from a different image affect representation learning compared to using [MASK] tokens or learnable vectors?
• RQ3 What is the optimal masking ratio and mixing strategy for achieving high performance with minimal computational cost in hierarchical ViTs?
• RQ4 Does dual reconstruction from a mixed input lead to better downstream transfer performance than single-image reconstruction?
• RQ5 Can MixMAE achieve strong performance with fewer pretraining epochs compared to existing MIM methods?

Key findings

• MixMAE achieves 85.1% top-1 accuracy on ImageNet-1K using only 600 pretraining epochs, demonstrating high efficiency and strong performance.
• With 900 pretraining epochs, MixMAE achieves 51.0 mIoU on ADE20K, showing strong generalization to dense prediction tasks.
• The 75% masking ratio achieved by mixing four images outperforms using [MASK] tokens at the same ratio, confirming the effectiveness of the mixing strategy.
• Dual reconstruction significantly improves performance, especially on ADE20K, where it boosts mIoU by 2.6 points compared to single reconstruction.
• MixMAE outperforms both SimMIM and MAE in FLOPs/performance tradeoff across seven downstream benchmarks, including image classification, object detection, and semantic segmentation.
• The method shows superior pretraining efficiency by eliminating [MASK] tokens in the encoder, reducing computation waste and avoiding pretrain-finetune inconsistency.

Figure 2 : Tradeoffs of FLOPs vs. (left) top-1 accuracy on ImageNet-1K, (middle) AP box on COCO, (right) and mIoU on ADE20K. All results are from various self-supervised pretraining methods followed by supervised finetuning. All entries on COCO [ 31 ] use Mask RCNN [ 22 ] framework. All entries on A