PaperlayerView exact revision on arXiv

Source-linked AI summary

ConvMAE: Masked Convolution Meets Masked Autoencoders

Peng Gao, Teli Ma, Hongsheng Li, Ziyi Lin, Jifeng Dai, Yu Qiao

arXiv:2205.03892v2cs.CV

TL;DR

ConvMAE addresses whether multi-scale hybrid convolution-transformer backbones can improve masked auto-encoding without sacrificing efficiency. It introduces hierarchical masking, masked convolutions, and multi-scale supervision, and reports stronger ImageNet, detection, and segmentation performance than MAE.

  • Problem

    Applying standard masked auto-encoding to convolution-transformer encoders creates excessive computation and a pretraining-finetuning discrepancy.

  • Method

    ConvMAE combines a multi-scale convolution-transformer encoder with block-wise masking, masked convolutions, and direct multi-scale feature supervision.

  • Results

    ConvMAE improves multi-task visual representation performance over MAE, including 1.4% higher ImageNet-1K finetuning accuracy for ConvMAE-Base than MAE-Base.

  • Takeaways & Limitations

    ConvMAE preserves MAE’s computational efficiency while generating discriminative multi-scale representations for downstream vision tasks.

Abstract

from arXiv · show

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.

1 Introduction

ConvMAE extends masked auto-encoding with a hybrid convolution-transformer encoder to learn discriminative multi-scale representations while preserving efficient training. It addresses the computational and pretraining-finetuning issues of applying standard masking to convolution-transformer networks and reports gains across vision tasks.

  • Motivation and contribution: ConvMAE combines hybrid convolution-transformer architectures with masked convolution in a self-supervised masked auto-encoding framework.The framework introduces minimal modifications to MAE while targeting scalable visual representations.
  • Motivation and contribution: Hybrid encoders use convolution blocks for high-resolution local content and transformer blocks for low-resolution global context, producing multi-scale features.The encoder progressively abstracts images into multi-scale token embeddings.
  • Motivation and contribution: Standard masking is inefficient because later transformer stages must retain all tokens, and visible-token-only pretraining creates a pretraining-finetuning discrepancy.These constraints undermine MAE’s efficiency advantage when applied directly to convolution-transformer encoders.
  • Proposed solution: Block-wise masking progressively upsamples a late-stage mask, while masked convolutions prevent masked and visible regions from mixing in early stages.This preserves computation efficiency in transformer stages and prevents information leakage during training.
  • Results: ConvMAE naturally provides multi-scale features for detection and segmentation, improving performance over MAE and even surpassing supervised-pretrained Swin and MViT.The supplied passage states this cross-task outcome without specifying the corresponding numerical results.
  • Results: 1.4% higher ImageNet-1K finetuning accuracy is reported for ConvMAE-Base than MAE-Base, alongside stronger COCO and ADE20K results.On COCO, ConvMAE-Base reaches 53.2% AP box and 47.1% AP mask with 25 epochs, versus MAE-Base’s 50.3% and 44.9% with 100 epochs; on ADE20K it reaches 51.7% versus 48.1% mIoU.

2 Approach

ConvMAE is an MAE derivative with a multi-scale convolution-transformer encoder, hierarchical masking, masked convolutions, and a multi-scale decoder. These components generate hierarchical features for downstream detection and segmentation while retaining efficient masked-token processing.

  • 2.1 A Brief Revisit of MAE: MAE reconstructs masked RGB patches from visible patches using a transformer encoder-decoder, processing only visible patches in the encoder.This design reduces the computation needed to pretrain large vision models.
  • 2.2 ConvMAE: ConvMAE modifies MAE’s encoder and masking strategy to learn discriminative multi-scale visual representations and reduce pretraining-finetuning discrepancy.The approach is presented as a simple derivative of MAE.
  • Block-wise Masking with Masked Convolutions: Hierarchical masking generates a late-stage mask and upsamples it for earlier convolutional stages so transformer layers receive only visible tokens.Direct input masking would otherwise leave every late-stage token partially visible and force all transformer tokens to be retained.
  • Hybrid Convolution-transformer Encoder: The three-stage encoder uses convolution blocks at early resolutions and transformer blocks at the late coarse resolution to combine local and global information.The first two stages use local convolutions, while the third uses self-attention to aggregate broader context.
  • Block-wise Masking with Masked Convolutions: Masked convolutions prevent early convolutional receptive fields from mixing masked and visible regions, avoiding information leakage during reconstruction.This is necessary because the 5 × 5 depthwise convolutions have receptive fields larger than masked patches.
  • The Multi-scale Decoder and Loss: The multi-scale decoder combines aligned features from encoder stages E1, E2, and E3 before reconstructing masked image patches.E1 and E2 are downsampled to E3’s size, and the objective considers reconstruction of masked patches only.
  • Downstream Finetuning: ConvMAE’s pretrained encoder naturally supplies multi-scale feature maps for existing object detection and semantic segmentation heads.For detection, stage-3 processing is adapted with shifted-window local attention in most layers to reduce computational cost.
  • Video Extension: VideoConvMAE extends the framework with cube embeddings, masked 3D convolutions in early stages, and transformer-based spatial-temporal fusion in stage 3.The extension is described as requiring simple modifications to the image framework.

3 Experiments

ConvMAE is evaluated across image classification, object detection, semantic segmentation, video understanding, and convergence studies. Across these tasks, it improves accuracy, transfer performance, or convergence speed over MAE and other pretrained backbones.

  • ImageNet-1K classification: ConvMAE improves ImageNet-1K finetuning accuracy by 1.4% over MAE with the same 1,600 pretraining epochs.It also surpasses SimMIM by 1.0% in finetuning accuracy.
  • Semantic segmentation: ConvMAE reaches 51.7% mIoU on ADE20K and outperforms MAE by 3.6% mIoU after 1,600-epoch pretraining.The evaluation uses UperNet with ADE20K validation data.
  • Video understanding: VideoConvMAE pretrained for 800 epochs with a multi-scale decoder surpasses VideoMAE pretrained for 1,600 epochs by more than 2.9% on Kinetics-400.On Something-Something-v2, the 800-epoch model slightly exceeds VideoMAE at 2,400 epochs, reducing pretraining by 3x.
  • Convergence and ablations: Removing masked convolution decreases ImageNet-1K finetuning accuracy from 84.6% to 81.5%.The ablation attributes this difference to information leakage from overlapping convolution windows.
  • Convergence and ablations: ConvMAE surpasses MAE’s final ImageNet-1K finetuning performance at 58 epochs and its COCO detection performance at 16 epochs.The COCO result corresponds to 6.6x faster convergence speed.

4 Related Work

Related work combines locality and hierarchy with ViTs, while masked autoencoding learns strong representations through masked patch reconstruction.

  • Hybrid convolution-transformer architectures reintroduce locality into ViTs and achieve strong performance across vision tasks.
  • Masked autoencoding learns representations by reconstructing masked patches, following approaches motivated by BERT.

5 Conclusion

ConvMAE combines hybrid local-global blocks with masked autoencoding to generate discriminative multi-scale features while retaining MAE’s efficiency and low pretraining-finetuning gap.

  • ConvMAE uses hybrid local-global blocks to generate discriminative multi-scale features.
  • ConvMAE maintains the computational efficiency and low pretraining-finetuning gap of original MAE.
  • The authors report significantly improved performance across various vision tasks and plan to study improved reconstruction targets.
  • The paper reports no anticipated negative societal impact from the proposed work.

6 Appendix

ConvMAE uses convolutional stages for high-resolution local processing and a transformer stage for low-resolution global reasoning, with model variants scaled to match MAE counterparts.

  • Architecture Details of ConvMAE Encoder: Stages 1 and 2 process progressively downsampled high-resolution embeddings with convolutional blocks for local information fusion.
  • Architecture Details of ConvMAE Encoder: Stage 3 projects features into lower-resolution tokens and applies transformer blocks for efficient global reasoning and an enlarged field of view.
  • Architecture Details of ConvMAE Encoder: ConvMAE small, base, large, and huge use parameter scales matched to MAE-small, MAE-base, MAE-large, and MAE-huge.
  • Model Architecture Details: Table 7 reports channel dimensions, layer counts, spatial resolutions, MLP ratios, and attention heads for the three encoder stages.
  • Model Scaling up and down: ConvMAE variants consistently outperform their MAE counterparts on ImageNet finetuning across parameter scales.
Loading 2205.03892v2…