PaperlayerView exact revision on arXiv

Source-linked AI summary

Masked Autoencoders Are Scalable Vision Learners

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, Ross Girshick

arXiv:2111.06377v3cs.CV

TL;DR

Learning useful visual representations is challenging because images contain substantial spatial redundancy and large models increasingly depend on labeled data. This paper introduces masked autoencoders with asymmetric encoder-decoders and high-ratio masking, achieving improved transfer performance and 87.8% ImageNet-1K accuracy with ViT-Huge.

  • Problem

    The paper addresses how to learn useful visual representations from images despite spatial redundancy and growing dependence on large labeled datasets.

  • Method

    MAE masks random image patches, encodes only visible patches, and uses a lightweight decoder to reconstruct the missing pixels from latent representations and mask tokens.

  • Results

    MAE pre-training improves transfer performance over supervised pre-training across detection and segmentation tasks, while ViT-Huge reaches 87.8% accuracy on ImageNet-1K.

  • Takeaways & Limitations

    The study supports scalable benefits from simple self-supervised autoencoding in vision and a trajectory resembling scalable self-supervised learning in NLP.

  • Takeaways & Limitations

    The method reflects biases in its training data and may generate nonexistent content, requiring further research and consideration for image generation.

Abstract

from arXiv · show

This paper shows that masked autoencoders (MAE) are scalable self-supervised learners for computer vision. Our MAE approach is simple: we mask random patches of the input image and reconstruct the missing pixels. It is based on two core designs. First, we develop an asymmetric encoder-decoder architecture, with an encoder that operates only on the visible subset of patches (without mask tokens), along with a lightweight decoder that reconstructs the original image from the latent representation and mask tokens. Second, we find that masking a high proportion of the input image, e.g., 75%, yields a nontrivial and meaningful self-supervisory task. Coupling these two designs enables us to train large models efficiently and effectively: we accelerate training (by 3x or more) and improve accuracy. Our scalable approach allows for learning high-capacity models that generalize well: e.g., a vanilla ViT-Huge model achieves the best accuracy (87.8%) among methods that use only ImageNet-1K data. Transfer performance in downstream tasks outperforms supervised pre-training and shows promising scaling behavior.

1. Introduction

MAE addresses the gap between masked autoencoding in language and vision with high-ratio patch masking and an asymmetric encoder-decoder. This design reduces redundancy, makes reconstruction challenging, and enables efficient, scalable visual representation learning.

  • Motivation: MAE is motivated by the slower progress of vision autoencoding despite the success of masked autoencoding in NLP.The paper attributes this gap to differences in architectures, information density, and decoder roles between language and vision.
  • Motivation: Vision’s spatial redundancy makes predicting a few missing patches insufficient, so MAE masks a very high portion of random patches to require holistic understanding.The strategy reduces redundancy and creates a challenging self-supervisory task beyond low-level image statistics.
  • Method: MAE encodes only visible patches without mask tokens, then uses a lightweight decoder with mask tokens to reconstruct the input in pixel space.The decoder reconstructs pixels from the latent representation and mask tokens, while the encoder processes only the visible subset.
  • Efficiency: 75% masking enables the encoder to process 25% of patches while optimizing accuracy and reducing overall pre-training time by 3× or more.Shifting mask tokens to the small decoder produces a large reduction in computation.
  • Results: 87.8% accuracy is achieved by a vanilla ViT-Huge model fine-tuned on ImageNet-1K, outperforming previous results using only ImageNet-1K data.The paper presents this as evidence that MAE pre-training supports high-capacity models with improved generalization.

2. Related Work

Related work spans masked prediction in NLP, classical and denoising autoencoding, masked image encoding, and contrastive self-supervised learning. These approaches differ in their corruption, prediction target, architecture, and reliance on data augmentation.

  • Masked Prediction in NLP: Masked language models such as BERT and GPT hold out input tokens, predict missing content, scale effectively, and transfer representations across downstream tasks.These methods are established as successful NLP pre-training approaches.
  • Autoencoding: Autoencoders learn representations by encoding inputs into latent representations and decoding them, while denoising autoencoders reconstruct clean signals from corrupted inputs.Examples of corruption include masking pixels and other transformations.
  • Masked Image Encoding: Masked image encoding learns visual representations from corrupted images, progressing from masked-pixel denoising and convolutional inpainting to Transformer-based pixel and patch prediction.Representative methods include Context Encoder, iGPT, and masked patch prediction studied in ViT.
  • Contrastive Learning: Computer-vision self-supervised learning includes contrastive methods that model similarity or dissimilarity across multiple views and strongly depend on data augmentation.Autoencoding represents a conceptually different direction from contrastive learning.

3. Approach

MAE reconstructs masked image patches with an asymmetric encoder-decoder: the encoder processes only visible patches, while a lightweight decoder combines encoded patches and mask tokens. Random high-ratio masking creates a sparse, nontrivial reconstruction task and enables efficient pre-training.

  • Architecture: MAE uses an asymmetric autoencoder whose encoder processes only visible patches without mask tokens, while the decoder reconstructs the original signal from encoded patches and mask tokens.The decoder receives the full token set, including encoded visible patches and shared learned mask tokens.
  • MAE encoder: 25% of the full patch set is an example encoder input size, enabling training of very large encoders with only a fraction of compute and memory.The encoder is a ViT that embeds visible patches and processes them with Transformer blocks.
  • MAE decoder: <10% computation per token versus the encoder is used by the default decoder, which runs only during pre-training while the encoder produces recognition representations.The decoder is narrower and shallower than the encoder and can be designed independently of it.
  • Reconstruction target: MAE predicts pixel values for each masked patch and trains with mean squared error; reconstructing normalized patch pixels improves representation quality.The decoder output is reshaped into a reconstructed image.
  • Masking: 75% masking works well for both fine-tuning and linear probing, while high-ratio random masking reduces redundancy and prevents easy extrapolation from neighboring visible patches.Patches are sampled uniformly without replacement, avoiding a potential center bias.

4. ImageNet Experiments

ImageNet experiments show that MAE pretraining substantially improves ViT performance while enabling efficient scaling. Ablations identify high random masking, visible-only encoding, simple augmentation, and extended training as key properties of the approach.

  • ViT-Large baseline: 84.9% accuracy follows MAE fine-tuning, surpassing 82.5% for supervised ViT-L trained from scratch and 76.5% in the original scratch implementation.Fine-tuning used 50 epochs, compared with 200 epochs from scratch.
  • Masking ratio: 75% masking works well for both linear probing and fine-tuning, while fine-tuning remains effective across masking ratios from 40–80%.The linear-probing accuracy gap reaches approximately 20% between 54.6% and 73.5%.
  • Architecture and efficiency: 3.3× lower training FLOPs results from omitting mask tokens in the encoder, producing a 2.8× wall-clock speedup and up to 4.1× with favorable configurations.A single-block decoder still achieves 84.8% fine-tuning accuracy.
  • Masking and augmentation: Random sampling performs best because it supports higher masking ratios, whereas block-wise masking degrades at 75% and produces blurrier reconstructions.Cropping-only augmentation works well; color jittering degrades results, and MAE remains decent without augmentation.
  • Training schedule: 1600 epochs do not saturate linear-probing accuracy, and MAE training ViT-L takes 31 hours versus 36 hours for MoCo v3 under the stated hardware setup.MAE accuracy improves steadily with longer training, unlike the cited MoCo v3 saturation at 300 epochs.
  • Scaling and comparisons: 87.8% accuracy with ViT-H at 448 resolution exceeds the 87.1% previous best using only IN1K data, while MAE also outperforms BEiT with 3.5× faster per-epoch training.MAE’s gain over scratch training increases for higher-capacity models, supporting scaling to larger networks.

5. Transfer Learning Experiments

Transfer experiments show that MAE pre-training consistently improves downstream detection, segmentation, and classification, while pixel reconstruction matches or surpasses token-based alternatives. Gains over supervised pre-training are especially pronounced for larger models, and classification results scale strongly with model size.

  • Object detection and segmentation: MAE exceeds supervised pre-training across all evaluated COCO detection and segmentation configurations.Mask R-CNN is fine-tuned end-to-end with a ViT backbone adapted for FPN; results report box AP and mask AP.
  • Object detection and segmentation: 4.0 points separate MAE from supervised pre-training with ViT-L: 53.3 vs. 49.3, APbox.With ViT-B, the corresponding gain is 2.4 points: 50.3 vs. 47.9, APbox.
  • Object detection and segmentation: Pixel-based MAE is better than or on par with BEiT, while both outperform MoCo v3; MAE is also simpler and faster.MoCo v3 is on par with supervised pre-training.
  • Semantic segmentation: MAE improves ADE20K semantic segmentation over supervised pre-training, including 3.7 points for ViT-L, and outperforms token-based BEiT.The experiments use UperNet, with results reported as mIoU.
  • Classification tasks: MAE classification accuracy scales strongly with larger models and surpasses previous best results on both iNaturalists and Places.The Places comparisons exceed prior results obtained by pre-training on billions of images.
  • Pixels vs. tokens: Normalized-pixel reconstruction is statistically similar to dVAE-token reconstruction across tested cases, showing tokenization is unnecessary for MAE.Using dVAE tokens is better than unnormalized pixels, but the difference from normalized pixels is statistically insignificant.

6. Discussion and Conclusion

The study finds that a simple autoencoder provides scalable self-supervised learning for computer vision, while adapting language-inspired methods to images requires care. The authors also caution that learned dataset biases and generated nonexistent content pose broader-impact concerns.

  • Discussion and Conclusion: A simple autoencoder, analogous to self-supervised methods in NLP, provides scalable computer-vision pre-training on ImageNet and transfer learning.The authors contrast this with the continued dominance of supervised pre-training in computer vision.
  • Discussion and Conclusion: Images lack a semantic decomposition like words, so MAE removes random patches and reconstructs pixels rather than semantic entities.Despite this difference, the method infers complex, holistic content.
  • Broader impacts: Training-data biases, including harmful societal biases, can be reflected by the model, which may also generate nonexistent content.The authors call for further research and consideration when using the method to generate images.

A. Implementation Details · A.1. ImageNet Experiments

The ImageNet implementation follows standard ViT conventions while adapting MAE’s asymmetric encoder-decoder with projection and sinusoidal positional embeddings. Separate training recipes are used for pre-training, fine-tuning, linear probing, and partial fine-tuning, including feature normalization for linear classifiers.

  • A.1. ImageNet Experiments: The encoder follows standard ViT blocks with multi-head self-attention, MLPs, LayerNorm, a final LayerNorm, and a projection layer matching decoder width.Sine-cosine positional embeddings are added to both encoder and decoder inputs.
  • A.1. ImageNet Experiments: An auxiliary dummy token is appended during pre-training to serve as the class token for linear probing and fine-tuning.The method also works similarly without this token by using average pooling.
  • A.1. ImageNet Experiments: Pre-training uses Xavier-uniform initialization and linear learning-rate scaling, while omitting color jittering, drop path, and gradient clipping.The default configuration is specified in Table 8.
  • A.1. ImageNet Experiments: End-to-end fine-tuning follows supervised ViT practice and uses layer-wise learning-rate decay.The default configuration is specified in Table 9.
  • A.1. ImageNet Experiments: Linear probing uses a distinct recipe in which regularization is generally harmful, so mixup, cutmix, drop path, and color jittering are disabled and weight decay is zero.Table 10 specifies LARS with a large batch for faster training; SGD works similarly with a 4096 batch.
  • A.1. ImageNet Experiments: An affine-free BatchNorm layer normalizes encoder features before the linear classifier, preserving the linear property of the probing setup.Pre-computing feature means and standard deviations is an alternative normalization procedure.

A.2. Supervised Training ViT-L/H from Scratch · A.3. Object Detection and Segmentation in COCO · A.4. Semantic Segmentation in ADE20K

The paper develops a stable recipe for supervised ViT-L/H training from scratch and describes transfer setups for COCO detection/segmentation and ADE20K semantic segmentation. Supervised ViT-L/H reach 82.6% and 83.1% accuracy with EMA, while downstream protocols adapt ViT features and fine-tuning configurations to each task.

  • A.2. Supervised Training ViT-L/H from Scratch: Training supervised ViT-L/H from scratch on ImageNet-1K is unstable, and directly applying prior recipes frequently produces NaN loss.The larger-model recipes were previously unexplored, despite strong publicly implemented baselines for smaller models.
  • A.2. Supervised Training ViT-L/H from Scratch: A stability-oriented recipe uses wd 0.3, batch size 4096, long warmup, β2=0.95, and selected regularization, finishing without NaN loss.The authors state that these choices improve training stability.
  • A.2. Supervised Training ViT-L/H from Scratch: 82.6% accuracy is achieved by ViT-L and 83.1% by ViT-H, compared with 81.5% and 80.9% without EMA, respectively.Both models show an overfitting trend when EMA is not used.
  • A.2. Supervised Training ViT-L/H from Scratch: 82.3% accuracy is achieved by ViT-B, compared with 82.1% without EMA and 81.8% in [53].This result is reported as a by-product of the supervised training recipe.
  • A.3. Object Detection and Segmentation in COCO: COCO comparisons search learning rate, weight decay, drop path rate, and fine-tuning epochs separately for every method, including competitors.The authors state that code and specific configurations will be released.
  • A.4. Semantic Segmentation in ADE20K: ADE20K semantic segmentation uses UperNet with end-to-end fine-tuning for 100 epochs, batch size 16, and an optimal learning rate searched for each entry.The setup follows the semantic segmentation code of [2].
  • A.4. Semantic Segmentation in ADE20K: For ADE20K transfer, relative position bias is enabled only during fine-tuning and initialized as zero, whereas the MAE pre-training does not use it.The BEiT reproduction instead uses relative position bias during both pretraining and fine-tuning.

A.5. Additional Classification Tasks

This section reports additional classification evaluations, including iNaturalist and Places fine-tuning, linear probing, and robustness testing on ImageNet variants. The evaluations adjust fine-tuning settings by dataset and test MAE models without specialized fine-tuning for ImageNet variants.

  • Additional Classification Tasks: iNaturalist and Places fine-tuning follows the Table 9 setting, with learning rate and fine-tuning epochs adjusted for each dataset.These results are reported in Table 6.
  • Additional Classification Tasks: Table 12 reports linear probing results for masked encoding methods, while MAE fine-tuning results appear in Table 3.The table marks one implementation with a dagger (†).
  • Additional Classification Tasks: Robustness evaluation tests the same MAE models on ImageNet validation-set variants without specialized fine-tuning and compares them with previous best system-level results.The reported metric is top-1 accuracy, except ImageNet-C, which uses mean corruption error.

B. Comparison on Linear Probing Results

Linear probing accuracy and fine-tuning accuracy are largely uncorrelated and emphasize different aspects of linear separability. MAE with ViT-L substantially outperforms previous masking-based methods in linear probing, but remains below contrastive methods such as MoCo v3.

  • Comparison on Linear Probing Results: Linear probing and fine-tuning accuracy are largely uncorrelated because they focus on different aspects of linear separability.Existing masked image encoding methods are generally less competitive than contrastive learning in linear probing.
  • Comparison on Linear Probing Results: 75.8% linear probing accuracy: MAE with ViT-L substantially outperforms previous masking-based methods but trails MoCo v3 at 77.6% for ViT-L.The comparison is conducted under the linear probing protocol.

C. Robustness Evaluation on ImageNet

On ImageNet robustness variants, MAE models scale strongly with model size and generally benefit from larger images, outperforming previous specialized systems. MAE pre-training also substantially exceeds supervised training, including a 35-point IN-A advantage for ViT-H.

  • Scaling and image resolution: MAE models show strong scaling across ImageNet validation variants, with larger model sizes producing significant gains.The evaluation reuses models fine-tuned on original ImageNet and performs inference on other validation sets without specialized fine-tuning.
  • Scaling and image resolution: Larger image sizes improve performance on every evaluated validation set except IN-C, while MAE outperforms previous specialized systems.
  • Comparison with supervised training: 35% better on IN-A, MAE pre-training with ViT-H reaches 68.2% versus 33.1% for supervised training.The comparison uses the supervised counterpart described in A.2 and reported at the bottom of Table 13.
Loading 2111.06377v3…