[Paper Review] Contrastive Masked Autoencoders are Stronger Vision Learners

Zhicheng Huang, Xiaojie Jin+6 more|arXiv (Cornell University)|Jul 27, 2022

Domain Adaptation and Few-Shot Learning62 references8 citations

TL;DR

This paper proposes Contrastive Masked Autoencoders (CMAE), a self-supervised vision pre-training framework that unifies masked image modeling (MIM) and contrastive learning (CL) through a siamese encoder-decoder architecture with a momentum encoder and a feature decoder. CMAE achieves state-of-the-art performance, attaining 85.3% top-1 accuracy on ImageNet-1k and 52.5% mIoU on ADE20k, surpassing prior methods by 0.7% and 1.8% respectively.

ABSTRACT

Masked image modeling (MIM) has achieved promising results on various vision tasks. However, the limited discriminability of learned representation manifests there is still plenty to go for making a stronger vision learner. Towards this goal, we propose Contrastive Masked Autoencoders (CMAE), a new self-supervised pre-training method for learning more comprehensive and capable vision representations. By elaboratively unifying contrastive learning (CL) and masked image model (MIM) through novel designs, CMAE leverages their respective advantages and learns representations with both strong instance discriminability and local perceptibility. Specifically, CMAE consists of two branches where the online branch is an asymmetric encoder-decoder and the momentum branch is a momentum updated encoder. During training, the online encoder reconstructs original images from latent representations of masked images to learn holistic features. The momentum encoder, fed with the full images, enhances the feature discriminability via contrastive learning with its online counterpart. To make CL compatible with MIM, CMAE introduces two new components, i.e. pixel shifting for generating plausible positive views and feature decoder for complementing features of contrastive pairs. Thanks to these novel designs, CMAE effectively improves the representation quality and transfer performance over its MIM counterpart. CMAE achieves the state-of-the-art performance on highly competitive benchmarks of image classification, semantic segmentation and object detection. Notably, CMAE-Base achieves $85.3\%$ top-1 accuracy on ImageNet and $52.5\%$ mIoU on ADE20k, surpassing previous best results by $0.7\%$ and $1.8\%$ respectively. The source code is publicly accessible at \url{https://github.com/ZhichengHuang/CMAE}.

Motivation & Objective

• To address the limited discriminability in masked image modeling (MIM) representations, which hinders performance on downstream vision tasks.
• To explore how contrastive learning (CL) can be effectively integrated with MIM to enhance feature quality and instance discriminability.
• To design a unified framework that harmonizes the strengths of MIM (local perceptibility) and CL (instance discriminability) without compromising either.
• To improve transfer performance across image classification, semantic segmentation, and object detection through better pre-training.

Proposed method

• CMAE employs a siamese architecture with an online asymmetric encoder-decoder branch for masked image reconstruction and a momentum-updated encoder branch for contrastive learning supervision.
• A novel feature decoder is introduced in the online branch to reconstruct masked token features, enabling meaningful contrastive learning from incomplete representations.
• Pixel shifting augmentation is applied to generate plausible positive views for contrastive learning, ensuring spatial consistency and compatibility with MIM.
• The momentum encoder processes full images, producing semantically intact features to guide the online encoder and enhance feature discriminability.
• The framework jointly optimizes masked reconstruction loss and contrastive contrastive loss, aligning features across views and improving representation quality.
• The method is scalable across ViT architectures (small, base, large), demonstrating consistent gains over MAE across model sizes.

Experimental results

Research questions

• RQ1 Can contrastive learning be effectively integrated with masked image modeling to improve representation quality?
• RQ2 How can contrastive learning be made compatible with the masked reconstruction objective in MIM, given differences in input augmentation and feature alignment?
• RQ3 Does combining MIM and CL lead to better downstream transfer performance on vision benchmarks?
• RQ4 What architectural and augmentation designs are necessary to ensure that contrastive learning enhances, rather than disrupts, the MIM objective?

Key findings

• CMAE-Base achieves 85.3% top-1 accuracy on ImageNet-1k, surpassing the previous state-of-the-art by 0.7%.
• On the ADE20k semantic segmentation benchmark, CMAE-Base achieves 52.5% mIoU, a 1.8% improvement over the prior best result.
• In object detection on COCO, CMAE-Base achieves 45.7% AP, demonstrating strong generalization to dense prediction tasks.
• Under partial fine-tuning, CMAE achieves a 3.8% performance gain over MoCo-V3 when fine-tuning only one transformer block, indicating higher-quality features.
• Linear probing results show a 5.9% improvement over MAE, confirming enhanced representation quality across evaluation protocols.
• Feature analysis confirms that CMAE learns more compact intra-class clusters and larger, more uniform inter-class distances, indicating superior discriminability.