Masked Autoencoders Are Scalable Vision Learners

In plain words

Masked autoencoders learn to see by playing a fill-in-the-blanks game with pictures. The image is cut into small squares, and most of them, about three quarters, are hidden at random. A large network looks only at the squares that are left. A small second network then tries to redraw the hidden squares from what the first one produced, plus a placeholder marker for each gap.

The big network never sees the blanks during its work, so training is cheap, and hiding that many squares makes the game hard enough to teach useful things. No labels are needed, because the picture itself is the answer. This is the image counterpart of the masked-word prediction used in language models. After the practice, the small network is thrown away and the big one is used for recognition.

TL;DR

• Self-supervised pre-training for ViT: mask 75% of patches, encode only the 25% that remain, and reconstruct the masked pixels with a small decoder §3.
• The encoder never sees mask tokens, which the paper reports both improves accuracy and cuts training compute §4.1 Mask token.
• Author-reported ImageNet-1K results use ImageNet-1K data only: ViT-Huge at 224 px reaches 86.9% and at 448 px 87.8% after end-to-end fine-tuning tab. 3.
• Trade-off: linear probing is weak (ViT-L 75.8%) while fine-tuning and partial fine-tuning are strong app. B §4.3.

Problem

Vision models of growing size need large labelled datasets, and self-supervised pre-training has worked well for language but lagged in vision §1. The paper names three gaps between the two fields §1 (i)-(iii):

• Architecture. Convolutions do not take mask tokens or positional embeddings easily; the Vision Transformer (ViT) removes this obstacle §1 (i).
• Information density. Words are dense in meaning, whereas image patches are highly redundant, so a missing patch can often be copied from its neighbours. A high masking ratio is needed to make the task non-trivial §1 (ii) §3 Masking.
• Decoder role. In vision the decoder reconstructs pixels, which are lower-level than the recognition target. The decoder design therefore changes the semantic level of the learned features §1 (iii) §4.1 Decoder design.

Contrastive methods depend heavily on data augmentation; the paper argues that masking makes augmentation less critical §4.1 Data augmentation.

Method

The model is an asymmetric encoder-decoder §3. Images are split into regular non-overlapping patches, following ViT §3 Masking app. A.1. The steps below follow the paper's order.

1. Masking. Sample a uniformly random subset of patches without replacement and remove the rest. The default ratio is 75% §3 Masking tab. 1 caption. The paper calls this "random sampling" and contrasts it with block-wise and grid-wise sampling §4.1 Mask sampling strategy.
2. Encoder. A standard ViT applied only to the visible patches. Each patch gets a linear projection plus a positional embedding, and the encoder runs its Transformer blocks on this set. No mask tokens are used §3 MAE encoder. The paper appends an auxiliary dummy token that plays the role of the class token for later classification app. A.1 ViT architecture.
3. Decoder input. Project the encoder output to the decoder width. Append one shared, learned mask token for each masked position, then put every token back in its original image position §3 MAE decoder [§3 Simple implementation].
4. Decoder. Add positional embeddings to all tokens, run the decoder Transformer blocks, and apply a final linear layer that outputs the pixel values of each patch §3 MAE decoder §3 Reconstruction target. The default decoder has 8 blocks at width 512 tab. 1 caption. Its cost is below 10% of the encoder's computation per token §3 MAE decoder.
5. Loss. Mean squared error between predicted and original pixels, computed only on masked patches §3 Reconstruction target. The paper says that including visible patches slightly lowers accuracy (about 0.5%) §3 footnote 1.
6. Target normalization (optional). Normalize each masked patch by its own pixel mean and standard deviation before computing the loss [§3 Reconstruction target]. The main ImageNet models use this variant tab. 3 caption; the ablations use unnormalized pixels by default tab. 1 caption.
7. After pre-training. Discard the decoder. The encoder is applied to full, uncorrupted images for recognition fig. 1 caption.

The paper's loss and masking are not written as numbered equations; the statements above are the complete description in the source §3.

# MAE pre-training step  [§3 Simple implementation]
# Input:  images (N, 3, H, W); mask_ratio (default 0.75)
# Output: scalar loss

def mae_step(images, mask_ratio):
    # Cut each image into non-overlapping patches  [§3 Masking]
    patches = patchify(images)                # (N, L, p*p*3)
    # Embed each patch and add its positional embedding  [§3 MAE encoder]
    tokens = linear(patches) + pos_embed      # (N, L, D)
    # Shuffle the patches per image and keep the first share  [§3]
    kept = int(L * (1 - mask_ratio))
    order = argsort(random_noise(N, L))       # new shuffle every step
    visible = gather(tokens, order[:, :kept]) # (N, kept, D)
    # Encoder: ViT blocks on visible patches only  [§3 MAE encoder]
    latent = encoder(visible)                 # (N, kept, D)
    # Project to decoder width and add one shared mask token per gap
    # [§3 MAE decoder]
    latent = decoder_embed(latent)            # (N, kept, D_dec)
    fill = repeat(mask_token, N, L - kept)    # (N, L - kept, D_dec)
    # Put every token back in its image position  [§3 Simple implementation]
    full = unshuffle(concat(latent, fill), order)   # (N, L, D_dec)
    # Positional embeddings on all tokens, then the decoder  [§3]
    full = decoder_blocks(full + decoder_pos_embed)
    # Linear head: pixel values for each patch  [§3 Reconstruction target]
    pred = decoder_head(full)                 # (N, L, p*p*3)
    # Target pixels, optionally normalized per patch  [§3]
    target = patchify(images)                 # (N, L, p*p*3)
    if normalize_targets:
        # Per-patch mean and variance; the code adds eps 1e-6 to the
        # variance  [code: models_mae.py#L205-L208]
        target = (target - mean_per_patch(target)) / sqrt(
            var_per_patch(target) + 1e-6)
    # Squared error averaged over each patch's pixels  [§3]
    per_patch = mean((pred - target) ** 2, over pixels)   # (N, L)
    # Average over hidden patches only, across the batch  [§3]
    # hidden: (N, L), 1 where the patch was masked
    return sum(per_patch * hidden) / sum(hidden)

Class tokens are omitted above. The implementation adds a class token to the encoder and strips it from the decoder output before the loss models_mae.py models_mae.py.

Hyperparameters

Pre-training and fine-tuning values come from the appendix tables app. A.1 unless noted. Where the paper leaves a value out, the table says "not stated".

Results

Author-reported. Setting: ImageNet-1K, top-1 validation accuracy at 224 px unless noted, single crop §4. Ablation numbers are ViT-L/16, 800 epochs, fine-tuning (ft) and linear probing (lin) tab. 1 caption.

Ablations (ViT-L/16, IN1K, 800 epochs) tab. 1

Wall-clock (128 TPU-v3 cores, TensorFlow, 800 epochs) tab. 2

ImageNet-1K fine-tuning tab. 3

The text gives the previous best with IN1K data only as 87.1% (512 px) §4.2.

COCO object detection and instance segmentation (Mask R-CNN) tab. 4

ADE20K semantic segmentation (UperNet, mIoU) tab. 5

Classification transfer (MAE on IN1K, then fine-tuned) tab. 6

Previous-best entries use pre-training on 1 billion images (Places205) or 3.5 billion (Places365) tab. 6 caption.

Other results

• Linear probing (ViT-B / ViT-L / ViT-H): 68.0 / 75.8 / 76.6; MoCo v3 reaches 77.6 with ViT-L app. B.
• Partial fine-tuning, ViT-L: tuning one block raises accuracy from 73.5% to 81.0%; tuning the last MLP sub-block gives 79.1%. Tuning 4 blocks is close to full fine-tuning; the paper reports a 2.6-point gap over MoCo v3 with 4 blocks tuned §4.3 fig. 9.
• Robustness, ViT-H without special fine-tuning: IN-Adversarial 68.2% vs 33.1% for the supervised baseline app. C.
• Scratch ViT-L/H with the paper's supervised recipe: 82.6% and 83.1% (80.9% without EMA) app. A.2.

Paper vs code

Code: facebookresearch/mae at commit efb2a8062c206524e35e47d04501ed4f544c0ae8. Paths are relative to the repository root. The code is licensed CC-BY-NC-4.0 (non-commercial) LICENSE. Nothing is quoted from it; it is cited only by path and line for comparison.

• Pre-training base lr. Paper: 1.5e-4 app. A.1 Pre-training table. The script default is --blr 1e-3 main_pretrain.py, but the documented command uses --blr 1.5e-4 PRETRAIN.md.
• Pre-training epochs. Paper: 800 for ablations, 1600 for main results tab. 1 caption tab. 3 caption. Script default is 400 main_pretrain.py; documented example is 800 PRETRAIN.md.
• Normalized targets. The script default is off main_pretrain.py. The documented command turns it on PRETRAIN.md, matching the main-result setting in tab. 3 caption.
• Target normalization formula. Paper: "mean and standard deviation" of the patch §3. Code: subtract the per-patch mean, divide by the square root of the per-patch variance plus 1e-6 models_mae.py. The epsilon is not in the paper.
• Fine-tuning recipe for ViT-B. Paper table: base lr 1e-3, layer decay 0.75 for all sizes app. A.1 End-to-end fine-tuning table. Documented ViT-B command: base lr 5e-4, layer decay 0.65 FINETUNE.md.
• Fine-tuning drop path. Paper: 0.1 for ViT-B/L and 0.2 for ViT-H app. A.1. Documented commands: 0.1 for ViT-B FINETUNE.md, 0.2 for ViT-L FINETUNE.md, 0.3 for ViT-H FINETUNE.md. For ViT-L and ViT-H the repo values differ from the paper.
• Fine-tuning pooling. The paper uses the class token (with average pooling reported as similar) app. A.1 ViT architecture. The repo states it uses --global_pool for fine-tuning FINETUNE.md.
• Mixup, cutmix and label smoothing in the script. Script defaults are mixup 0 and cutmix 0 main_finetune.py; the documented commands set mixup 0.8 and cutmix 1.0 FINETUNE.md.
• Pre-training crop scale. Paper: "RandomResizedCrop" with no scale range app. A.1. Code: scale (0.2, 1.0), bicubic, plus horizontal flip main_pretrain.py.
• Patch embedding. Paper: a linear projection of patches §3 MAE encoder. Code: timm's PatchEmbed, initialized like a linear layer models_mae.py models_mae.py.
• Class token. Paper: an auxiliary dummy token is appended during pre-training app. A.1 ViT architecture. Code: the encoder always carries a class token models_mae.py models_mae.py. Consistent in effect.
• Decoder heads. Not stated in the paper. Code: 16 models_mae.py.
• Learning-rate schedule. Paper: cosine decay with 40 warmup epochs app. A.1. Code: linear warmup, cosine decay to min_lr (default 0), updated every iteration with fractional epochs lr_sched.py engine_pretrain.py.
• Weight decay exclusions. Paper: not stated. Code: no decay on biases and normalization layers main_pretrain.py.
• Mixed precision. Paper: not stated. Code: fp16 autocast with a gradient scaler engine_pretrain.py main_pretrain.py.
• Input normalization. Paper: not stated. Code: ImageNet mean and std are applied before patchify, so pixel targets are in that normalized space main_pretrain.py.
• Scope of the repo. Only ImageNet pre-training, fine-tuning and linear probing are in the repo. The COCO and ADE20K code is not here; the paper points to a separate reference for COCO details app. A.3.
• Reproduction claim. The repo says its PyTorch port gets 85.5% fine-tuned ViT-L after 800 epochs, versus 85.4% in the paper's Table 1d PRETRAIN.md. This is the authors' own statement and was not re-run here.

Pitfalls

• Hidden-patch count. The code keeps int(L * (1 - mask_ratio)) patches models_mae.py. For 224 px with patch 16 (196 patches), that is 49 visible patches at 75%; the rounding is truncation, not rounding.
• Loss is a batch average, not per image. The loss is summed over all hidden patches in the batch and divided by their total count models_mae.py. Because every image has the same hidden count at a fixed ratio, this equals the per-image mean averaged over the batch.
• Visible patches are not in the loss. Reconstructions of visible patches look worse than hidden ones. The paper keeps them this way on purpose fig. 2 caption.
• Normalized targets change the default. The ablations and the script use unnormalized pixels by default, and the main results use normalized pixels tab. 1 caption tab. 3 caption. Check which one a given experiment uses.
• Decoder depth depends on the use. A single decoder block already gives 84.8% fine-tuned, but 65.5% linear probing versus 73.5% with 8 blocks tab. 1a §4.1 Decoder design. Choose the depth for the downstream protocol.
• Mask tokens in the encoder hurt. Using them drops linear probing by about 14 points and costs 3.3× the FLOPs in the paper's setting tab. 1c §4.1 Mask token.
• Fine-tuning instability at ViT-H. The repo notes that the class-token setting can produce NaN losses when fine-tuning ViT-H on GPUs with mixed precision; turning off mixed precision avoids it but is slower FINETUNE.md.
• Position embedding at 448 px. The 448-px result needs the 224-px positional embedding to be resized. The paper does not describe this. The code interpolates the position tokens bicubically, keeps the class token unchanged pos_embed.py main_finetune.py.
• Costs compared across papers. Paper: 31 hours for ViT-L, 1600 epochs, 128 TPU-v3 §4.2. Repo: about 42 hours for 800 epochs on 64 V100 GPUs PRETRAIN.md. Different hardware and frameworks; not comparable.
• Per-epoch work. The encoder sees 25% of patches per epoch, while contrastive methods see 200% or more §4.1 Training schedule. Epoch counts therefore do not show equal compute.
• Linear probing understates the features. Linear probing is weaker than contrastive baselines, and the paper argues it is a poor proxy for transfer app. B §4.3. Partial or full fine-tuning is the protocol the paper favours.
• Relative position bias in segmentation. ADE20K transfer turns it on with zero initialization, but pre-training does not use it app. A.4. Do not copy it into pre-training.
• Linear probing batch size. The paper's table gives an effective batch of 16384 app. A.1 Linear probing table. The script default is 512 per device main_linprobe.py, and the effective batch depends on the launch command, which the repo does not document.
• Broader impacts. The paper flags that reconstructions reflect training-data biases and may produce content that does not exist §6.
• Table and figure numbers. The paper's text and captions cite its own numbering. Table and figure numbers here follow the arXiv v3 LaTeX source order, because the rendered PDF was not read. Appendix tables are numbered continuing from the main text tab. 8–tab. 13.

Lineage

• Builds on: masked language modelling in BERT §1 §2; denoising autoencoders §2; the Vision Transformer §3 app. A.1.
• Compared against: BEiT (token prediction with a dVAE) tab. 3 tab. 4; MoCo v3 and contrastive learning tab. 3 app. B; DINO tab. 3; iGPT app. B.
• Superseded by: not stated in the source.