Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SupMAE: Supervised Masked Autoencoders Are Efficient Vision Learners

Feng Liang, Yangguang Li, Diana Marculescu

TL;DR

SupMAE addresses the gap in MAE-style pre-training by introducing a supervised classification branch that operates on a subset of visible patches, enabling global feature learning with improved efficiency. The method jointly optimizes a reconstruction objective and a classification objective, resulting in faster pre-training, greater robustness, and stronger transfer to few-shot and dense downstream tasks such as ADE20K segmentation. Empirical results show competitive ImageNet performance with substantially reduced compute, enhanced robustness across variants, and superior few-shot and segmentation transfer, demonstrating the value of combining supervised signals with masked image modeling. This work highlights the potential of multi-objective pre-training to produce more efficient and generalizable vision representations.

Abstract

Recently, self-supervised Masked Autoencoders (MAE) have attracted unprecedented attention for their impressive representation learning ability. However, the pretext task, Masked Image Modeling (MIM), reconstructs the missing local patches, lacking the global understanding of the image. This paper extends MAE to a fully supervised setting by adding a supervised classification branch, thereby enabling MAE to learn global features from golden labels effectively. The proposed Supervised MAE (SupMAE) only exploits a visible subset of image patches for classification, unlike the standard supervised pre-training where all image patches are used. Through experiments, we demonstrate that SupMAE is not only more training efficient but it also learns more robust and transferable features. Specifically, SupMAE achieves comparable performance with MAE using only 30% of compute when evaluated on ImageNet with the ViT-B/16 model. SupMAE's robustness on ImageNet variants and transfer learning performance outperforms MAE and standard supervised pre-training counterparts. Codes are available at https://github.com/enyac-group/supmae.

SupMAE: Supervised Masked Autoencoders Are Efficient Vision Learners

TL;DR

SupMAE addresses the gap in MAE-style pre-training by introducing a supervised classification branch that operates on a subset of visible patches, enabling global feature learning with improved efficiency. The method jointly optimizes a reconstruction objective and a classification objective, resulting in faster pre-training, greater robustness, and stronger transfer to few-shot and dense downstream tasks such as ADE20K segmentation. Empirical results show competitive ImageNet performance with substantially reduced compute, enhanced robustness across variants, and superior few-shot and segmentation transfer, demonstrating the value of combining supervised signals with masked image modeling. This work highlights the potential of multi-objective pre-training to produce more efficient and generalizable vision representations.

Abstract

Recently, self-supervised Masked Autoencoders (MAE) have attracted unprecedented attention for their impressive representation learning ability. However, the pretext task, Masked Image Modeling (MIM), reconstructs the missing local patches, lacking the global understanding of the image. This paper extends MAE to a fully supervised setting by adding a supervised classification branch, thereby enabling MAE to learn global features from golden labels effectively. The proposed Supervised MAE (SupMAE) only exploits a visible subset of image patches for classification, unlike the standard supervised pre-training where all image patches are used. Through experiments, we demonstrate that SupMAE is not only more training efficient but it also learns more robust and transferable features. Specifically, SupMAE achieves comparable performance with MAE using only 30% of compute when evaluated on ImageNet with the ViT-B/16 model. SupMAE's robustness on ImageNet variants and transfer learning performance outperforms MAE and standard supervised pre-training counterparts. Codes are available at https://github.com/enyac-group/supmae.
Paper Structure (19 sections, 1 equation, 3 figures, 12 tables, 1 algorithm)

This paper contains 19 sections, 1 equation, 3 figures, 12 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. SupMAE Method
  4. Experimental Results
  5. Main results on ImageNet-1K
  6. Robustness evaluation on ImageNet variants
  7. Transfer Learning experiments
  8. Few-shot learning on 20 classification datasets
  9. Semantic segmentation in ADE20k
  10. Ablation studies
  11. Discussion
  12. Conclusion
  13. Acknowledgment
  14. Appendix
  15. Implementation details on ImageNet
  16. ...and 4 more sections

Figures (3)

  • Figure 1: Illustration of the proposed SupMAE method. The proposed SupMAE extends MAE by adding a branch for supervised classification in parallel with the existing reconstruction objective. In the pre-training phase, only a subset of the visible patches is processed by a ViT encoder. Their corresponding patch features are used to (1) reconstruct the missing pixels and (2) classify the category. In the fine-tuning phase, the encoder is applied to uncorrupted images for recognition tasks.
  • Figure 2: Performance of different pre-training epochs Comparison between MAE and SupMAE when pre-trained for different epochs. SupMAE is efficient and shows a much faster convergence speed.
  • Figure 3: Comparison between MAE and SupMAE when fine-tuned for 100 epochs on ImageNet-1K. The model architecture is ViT-B/16. Both MAE and SupMAE are pre-trained for 200 epochs. Our SupMAE brings a much better initialization point than its MAE counterpart.