ConvMAE: Masked Convolution Meets Masked Autoencoders

TL;DR

ConvMAE introduces a multi-scale hybrid convolution-transformer architecture with a novel masking strategy for efficient masked auto-encoding, leading to improved performance on vision tasks.

Contribution

The paper proposes a new ConvMAE framework that combines masked convolution with a block-wise masking strategy to enhance representation learning and efficiency.

Findings

ConvMAE-Base improves ImageNet-1K accuracy by 1.4%.

ConvMAE-Base outperforms MAE-Base in object detection after fewer epochs.

The proposed methods reduce computational cost and improve multi-scale feature supervision.

Abstract

Vision Transformers (ViT) become widely-adopted architectures for various vision tasks. Masked auto-encoding for feature pretraining and multi-scale hybrid convolution-transformer architectures can further unleash the potentials of ViT, leading to state-of-the-art performances on image classification, detection and semantic segmentation. In this paper, our ConvMAE framework demonstrates that multi-scale hybrid convolution-transformer can learn more discriminative representations via the mask auto-encoding scheme. However, directly using the original masking strategy leads to the heavy computational cost and pretraining-finetuning discrepancy. To tackle the issue, we adopt the masked convolution to prevent information leakage in the convolution blocks. A simple block-wise masking strategy is proposed to ensure computational efficiency. We also propose to more directly supervise the multi-scale features of the encoder to boost multi-scale features. Based on our pretrained ConvMAE models, ConvMAE-Base improves ImageNet-1K finetuning accuracy by 1.4% compared with MAE-Base. On object detection, ConvMAE-Base finetuned for only 25 epochs surpasses MAE-Base fined-tuned for 100 epochs by 2.9% box AP and 2.2% mask AP respectively. Code and pretrained models are available at https://github.com/Alpha-VL/ConvMAE.