Reproduce “MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers” on OpenResearch

About this paper

Jihao Liu, Xin Huang, Jinliang Zheng, Yu Liu, Hongsheng Li

MixMAE introduces a self-supervised pretraining method for hierarchical Vision Transformers that uses mixed image inputs for the encoder and a dual reconstruction strategy for the decoder. This approach enables more efficient pretraining, achieving 85.1% top-1 accuracy on ImageNet-1K with Swin-B in fewer epochs and GPU hours than existing methods, while demonstrating strong transferability to downstream tasks.

Key insights

• Providing meaningful visual context by filling masked regions with patches from another image (the "Mix" strategy) is more effective for pretraining hierarchical ViTs than using placeholder `[MASK]` tokens or zeros.
• The dual reconstruction mechanism effectively leverages information from both mixed images to achieve substantial performance gains (e.g., +2.6 mIoU on ADE20K) with minimal additional computational overhead.
• Restricting self-attention within tokens from the same original image group ("masked self-attention") during encoding is crucial for faster convergence and improved performance in the complex mixed-input setup.

Summary by alphaXiv.