MAE: Masked Autoencoders Are Scalable Vision Learners

Paper Overview

MAE — Masked Autoencoders Are Scalable Vision Learners — introduces a masked autoencoder framework for Vision Transformers that learns powerful visual representations without any labeled data. The method is strikingly simple: randomly mask 75% of image patches, feed only the visible 25% through a large encoder, then use a lightweight decoder to reconstruct the missing pixels. Published at CVPR 2022 by Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick at Facebook AI Research (FAIR), MAE demonstrated that masked image modeling can match the impact that masked language modeling (BERT) had on NLP.

The asymmetric encoder-decoder design is MAE’s computational breakthrough. The encoder sees 49 visible patches plus CLS rather than all 196 patches plus CLS. This sharply reduces attention work, while the smaller decoder (8 blocks, 512 dimensions) reconstructs the full grid. In Table 2, the default design is 2.8× faster for ViT-L and 3.5× faster for ViT-H than encoding mask tokens too. Raw pixels provide a simple target; per-patch normalization improves the main recognition results. No tokenizer, contrastive pairs, or momentum encoder is needed.

Key results speak to MAE’s effectiveness at scale. ViT-H (632M parameters) reaches 86.9% top-1 accuracy on ImageNet-1K with fine-tuning, and 87.8% when fine-tuned at 448×448 resolution — surpassing supervised ViT at every model scale. Linear probing with ViT-L achieves 75.8% top-1, confirming strong representation quality even without fine-tuning. On transfer tasks, MAE pre-trained features achieve 53.3 AP^box on COCO object detection and 53.6 mIoU on ADE20K semantic segmentation, demonstrating that pixel reconstruction learns spatially rich representations that generalize beyond classification.

Why 75% Masking?

Language has high information density — BERT’s 15% masking creates a sufficiently challenging task because missing a single word requires understanding syntax, semantics, and broader context to predict correctly. But images have massive spatial redundancy. Neighboring patches share textures, edges, and colors. At 15% masking, a model can trivially reconstruct missing patches by interpolating from nearby visible patches without learning any high-level understanding of objects, scenes, or spatial relationships. The pretext task becomes too easy to drive meaningful representation learning.

MAE’s 75% masking ratio creates genuine information scarcity. When 3 out of every 4 patches are removed, the remaining patches are too sparse for local interpolation to work. The model must understand objects, scenes, and spatial relationships to reconstruct the missing regions — it cannot simply copy neighboring textures. The paper finds 75% a good choice for both evaluation protocols, while fine-tuning works well across a broad 40–80% range. This is an empirical result for the studied recipe, not a universal optimum for every image model.

The MAE Pipeline

The input image is divided into non-overlapping 16×16 patches (a 224×224 image produces 14×14 = 196 patches). 75% of these patches are randomly selected for masking. The key insight: masked patches are simply removed — they do not enter the encoder at all. Only the visible 25% (49 patches) are fed to the encoder as tokens, with positional embeddings added so the encoder knows each patch’s spatial location within the original image grid.

The encoder output contains 49 encoded patch tokens plus CLS. The decoder receives 197 tokens: CLS, 49 projected visible patches, and 147 copies of a shared learned mask token, each with positional embeddings so the decoder knows which spatial positions are masked and which carry encoded information. The decoder is deliberately lightweight (8 Transformer blocks with 512 dimensions, compared to the encoder’s 24 blocks with 1024 dimensions). It reconstructs per-patch normalized pixel values, and the MSE loss is computed only on masked patches — visible patches that the encoder already saw are excluded from the loss, forcing the model to focus entirely on predicting the unknown regions.

Asymmetric Encoder-Decoder

The encoder is a standard Vision Transformer (ViT-L: 24 blocks, 1024 dimensions, roughly 307M parameters), applied to only the visible patches and CLS. The attention matrix shrinks from 197×197 to 50×50 entries per head. Ignoring CLS gives the familiar 16× reduction in attention entries, but this is not a 16× reduction in total encoder computation: token projections and feed-forward layers scale linearly with sequence length.

The decoder processes the full grid at lower width and depth. Table 1c reports a 3.3× reduction in total training FLOPs relative to putting mask tokens in the encoder. Table 2 measures a 2.8× wall-clock speedup for ViT-L and 3.5× for ViT-H with the default 8-block decoder. These are measured training configurations, not consequences of the attention formula alone.

At 75% masking with 196 patches, this ratio is 50^2/197^2\approx 0.0644. The extra token is CLS in both sequences.

Decoder Design Matters (Sometimes)

Decoder depth affects the two evaluation protocols differently. In Table 1a, increasing the decoder from one block to eight changes fine-tuning accuracy from 84.8% to 84.9%, but linear probing from 65.5% to 73.5%. These are 800-epoch, unnormalized-pixel ablations, not the main 1600-epoch benchmark.

The paper suggests that a deeper decoder can handle pixel-level reconstruction specialization, leaving the encoder representation more abstract. Fine-tuning can adapt the encoder’s final layers to recognition, so it depends less on decoder depth. The default decoder uses eight blocks and width 512.

Reconstruction Target

MAE can reconstruct raw pixels or per-patch normalized pixels. The normalized variant centers each target patch and divides by its standard deviation; the released code uses sample variance and a small numerical stabilizer. Table 1d improves ViT-L fine-tuning from 84.9% to 85.4% under the 800-epoch ablation recipe. The loss is mean squared error computed only on masked patches — visible patches that the encoder already saw are excluded from the loss, ensuring the model is evaluated purely on its ability to predict unseen content.

The simplicity of the pixel target is a feature, not a limitation. Unlike BEiT (which requires a separate dVAE tokenizer to convert patches into discrete visual tokens) or methods that introduce negative-pair sampling or a momentum encoder, MAE’s pixel reconstruction requires no additional components beyond the encoder and decoder themselves. This simplicity enables clean scaling — there are no auxiliary models to tune, no negative sampling strategies to balance, and no tokenizer vocabulary to optimize — making MAE straightforward to implement, fast to train, and reliable to scale to larger architectures.

Implementation

The loss above becomes concrete when the same patch order is preserved through masking, encoding, restoration, and reconstruction. The executable walkthrough follows that path, then checks the full pretrained visualization model on real photographs. The benchmark results that follow use the paper’s larger training runs.

This native PyTorch program runs a reduced random-weight training step offline on CPU. The reconstruction gallery below separately uses the official pretrained ViT-L visualization checkpoint. Shapes in the walkthrough distinguish the small smoke configuration from that full model.

Download the Python program · Pinned requirements · Artifact exporter

python -m pip install -r requirements-mae.txt
python mae_training.py --smoke

1. Give each patch a reversible pixel order

Original paper

Math and tensor shapes

L=(H/P)^2

x_{patch}\in\mathbb R^{B\times L\times 3P^2}

A reshape alone cannot move channel-first RGB into pixel-first targets. The permutation makes each consecutive triple a pixel’s RGB values. The roundtrip test checks both exact recovery and ordering, which must match the decoder’s learned output channels.

Smoke: [2, 3, 32, 32] → [2, 16, 192]; ViT-L: [B, 196, 768]

Python implementation

python

def patchify(self, images):
    """Maps BCHW into raster-ordered patches with pixel-then-channel order."""
    p = self.patch_embed.patch_size[0]
    if images.ndim != 4 or images.shape[1] != 3:
        raise ValueError('Expected [B, 3, H, H] RGB images.')
    if images.shape[2:] != (self.image_size, self.image_size):
        raise ValueError('Image dimensions must match the configured square size.')
    b, _, h, _ = images.shape
    grid = h // p
    return images.reshape(b, 3, grid, p, grid, p).permute(
        0, 2, 4, 3, 5, 1).reshape(b, grid * grid, p * p * 3)

def unpatchify(self, patches):
    """Inverts patchify exactly, preserving pixel/channel order."""
    p = self.patch_embed.patch_size[0]
    grid = self.image_size // p
    if patches.ndim != 3 or patches.shape[1:] != (grid * grid, p * p * 3):
        raise ValueError('Patch count or patch vector width is invalid.')
    b = patches.shape[0]
    return patches.reshape(b, grid, grid, p, p, 3).permute(
        0, 5, 1, 3, 2, 4).reshape(b, 3, grid * p, grid * p)

2. Shuffle each image independently, then keep a prefix

Original paper

Math and tensor shapes

K=\lfloor L(1-r)\rfloor

\pi=\operatorname{argsort}(u),\quad \pi^{-1}=\operatorname{argsort}(\pi)

The first K shuffled positions enter the encoder. A binary mask is created in shuffled order, then restored to the original spatial order. Here 0 means visible and 1 means removed. The inverse permutation is saved for the decoder; it is not a second random draw.

Smoke at r = 0.75: 16 → 4 visible patches; mask [2, 16]

Python implementation

python

def random_masking(self, tokens, mask_ratio):
    """Samples independent permutations, keeping only visible tokens."""
    if not math.isfinite(mask_ratio) or not 0 < mask_ratio < 1:
        raise ValueError('Mask ratio must be finite and strictly between zero and one.')
    b, length, width = tokens.shape
    keep = int(length * (1 - mask_ratio))
    if not 0 < keep < length:
        raise ValueError('Mask ratio must leave at least one visible and one masked patch.')
    order = torch.rand(b, length, device=tokens.device).argsort(dim=1)
    restore = order.argsort(dim=1)
    visible = tokens.gather(1, order[:, :keep, None].expand(-1, -1, width))
    mask = tokens.new_ones(b, length)
    mask[:, :keep] = 0
    return visible, mask.gather(1, restore), restore

3. Encode visible patches and one CLS token

Original paper

Math and tensor shapes

V=\operatorname{keep}(XW+P_{enc})

Z=E([c+p_0;V])

An 8 × 8 convolution is a shared linear projection of each smoke-test patch; the pretrained model uses 16 × 16. Fixed two-dimensional sine/cosine positions preserve each visible patch’s original location despite shuffling. Every encoder block sees only K + 1 tokens.

Smoke [2, 5, 32]; real ViT-L [B, 50, 1024] including CLS

Python implementation

python

def forward_encoder(self, images, mask_ratio):
    self.patchify(images)  # Validate the input before convolution.
    tokens = self.patch_embed(images) + self.pos_embed[:, 1:]
    visible, mask, restore = self.random_masking(tokens, mask_ratio)
    cls = (self.cls_token + self.pos_embed[:, :1]).expand(images.shape[0], -1, -1)
    latent = torch.cat([cls, visible], dim=1)
    for block in self.blocks:
        latent = block(latent)
    return self.norm(latent), mask, restore

4. Project the latent, append masks, and undo the shuffle

Original paper

Math and tensor shapes

U=ZW_d+b_d

R=\operatorname{unshuffle}([U_{vis};m;\ldots;m])

T=[U_{cls};R]

The projected visible patches retain the shuffled keep order. Appending exactly L − K mask tokens reconstructs a length-L shuffled list; gathering with the saved inverse puts the visible values and missing slots back into raster order.

Smoke: [2, 5, 32] → [2, 17, 24]; 12 shared mask tokens per image

Python implementation

python

def restore_tokens(self, latent, restore):
    projected = self.decoder_embed(latent)
    missing = restore.shape[1] + 1 - projected.shape[1]
    masks = self.mask_token.expand(projected.shape[0], missing, -1)
    shuffled = torch.cat([projected[:, 1:], masks], dim=1)
    patches = shuffled.gather(1, restore[..., None].expand(-1, -1, projected.shape[2]))
    return torch.cat([projected[:, :1], patches], dim=1)

5. Add decoder positions and predict every patch

Original paper

Math and tensor shapes

\hat X=\operatorname{dropCLS}(D(T+P_{dec})W_{pix}+b_{pix})

Identical learned mask vectors become distinguishable through their decoder positions. The smoke decoder has one block; the genuine visualization model has eight blocks of width 512. The decoder predicts visible positions too, although their outputs do not contribute to the training loss.

Smoke: [2, 17, 24] → [2, 16, 192]; ViT-L: [B, 196, 768]

Python implementation

python

def forward_decoder(self, latent, restore):
    tokens = self.restore_tokens(latent, restore) + self.decoder_pos_embed
    for block in self.decoder_blocks:
        tokens = block(tokens)
    return self.decoder_pred(self.decoder_norm(tokens))[:, 1:]

6. Average pixel error only over removed patches

Original paper

Math and tensor shapes

t_i=(x_i-\mu_i)/(s_i^2+10^{-6})^{1/2}

e_i=\operatorname{mean}_c(\hat x_{ic}-t_{ic})^2

\mathcal L=\frac{\sum_i m_i e_i}{\sum_i m_i}

When norm_pix_loss is enabled, the variance uses correction=1, matching torch.var in the author code. Without that option, t is the original ImageNet-normalized RGB patch. The smoke enables patch normalization; the official visualization release uses raw pixel targets. These two target spaces must not be confused when displaying images.

Prediction [B, L, 3P²], mask [B, L] → scalar loss

Python implementation

python

def forward_loss(self, images, prediction, mask):
    target = self.patchify(images)
    if prediction.shape != target.shape or mask.shape != target.shape[:2]:
        raise ValueError('Prediction and mask dimensions must match target patches.')
    if not torch.all((mask == 0) | (mask == 1)) or mask.sum() == 0:
        raise ValueError('Mask must be binary with at least one removed patch.')
    if self.norm_pix_loss:
        mean = target.mean(dim=-1, keepdim=True)
        variance = target.var(dim=-1, keepdim=True, correction=1)
        target = (target - mean) / (variance + 1e-6).sqrt()
    patch_error = (prediction - target).square().mean(dim=-1)
    return (patch_error * mask).sum() / mask.sum()

7. Backpropagate through both halves of the autoencoder

Original paper

Math and tensor shapes

\theta\leftarrow\operatorname{AdamW}(\theta,\nabla_\theta\mathcal L)

Gradients reach the patch projection, attention blocks, learned mask token, and pixel prediction head. The fixed sine/cosine positions remain frozen. This minimal AdamW step omits the paper’s long schedule, distributed batches, learning-rate schedule, and mixed precision; it checks the mechanism rather than reproducing a benchmark.

One CPU smoke update: 2 images, 12 masked patches per image

Python implementation

python

def train_step(model, images, optimizer):
    """Takes one reconstruction update; both encoder and decoder learn."""
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss, prediction, mask = model(images, mask_ratio=.75)
    loss.backward()
    optimizer.step()
    return loss.detach(), prediction.detach(), mask.detach()

8. Show predictions and a visible-patch composite separately

Original paper

Math and tensor shapes

C=(1-m)\odot x+m\odot\hat x

x_{display}=\operatorname{unpatchify}(C)

Raw-pixel targets make direct display possible after undoing ImageNet channel normalization. The composite retains the observed pixels exactly; the full prediction exposes what the model emits at unsupervised visible positions. This function rejects normalized-patch models because their outputs do not directly encode the original brightness and contrast.

ViT-L prediction [1, 196, 768] → RGB [1, 3, 224, 224]

Python implementation

python

def reconstruct(self, images, prediction, mask):
    """Produces a visible-patch composite only for raw-pixel target models."""
    if self.norm_pix_loss:
        raise ValueError('Display requires a model trained on raw pixel targets.')
    expanded_mask = mask[..., None]
    composite = self.patchify(images) * (1 - expanded_mask) + prediction * expanded_mask
    return self.unpatchify(composite)

Inspect genuine pretrained reconstructions

The released ViT-L/16 visualization model predicts these photographs with 147 of 196 patches masked. Seed 7 fixes the mask. This release uses raw pixel targets without an extra GAN loss; it is distinct from the normalized-patch model used for the main recognition benchmarks.

Original

25% visible

Full prediction

Prediction + visible

Original

25% visible

Full prediction

Prediction + visible

All 398 checkpoint tensors loaded strictly. Maximum absolute prediction differences from the author forward pass were 4.95e-6 and 6.20e-6. The parity run preserves the author’s masking, positions, encoder/decoder sequence, and loss, while replacing timm with native blocks and checking explicit softmax attention against SDPA.

Download masks, numerical results, hashes, and provenance · Official visualization weights · Official model license

Reproduce this export

Use Python 3.13.13, PyTorch 2.12.0+cpu, NumPy 2.1.2, and Pillow 12.2.0. Save models_mae.py and util/pos_embed.py from revision efb2a806 in an official-source directory, naming the helper util_pos_embed.py. The adapter replaces the removed NumPy np.float alias with the equivalent builtin float.

From the repository root, with the checkpoint downloaded locally:

python examples/mae_export.py --checkpoint mae_visualize_vit_large.pth --official-dir official-source --output public/examples/mae-demo

The exporter verifies the licensed source photographs in public/examples/vit-demo. It directly resizes RGB images to 224 × 224 with Pillow bicubic interpolation, scales by 1/255, then applies ImageNet mean (0.485, 0.456, 0.406) and standard deviation (0.229, 0.224, 0.225). Outputs are de-normalized, clipped, and saved as PNG. The masked preview fills missing normalized pixels with zero, which displays as the ImageNet mean color.

Scaling to Larger Models

MAE’s computational efficiency enables training Vision Transformers at unprecedented scale on ImageNet-1K alone. ViT-B (86M parameters) reaches 83.6% fine-tuning accuracy, already 1.3 points above supervised ViT-B (82.3%). ViT-L (307M parameters) reaches 85.9%, surpassing supervised ViT-L (82.6%) by 3.3 points. The gap widens with scale — larger models benefit more from self-supervised pre-training because MAE’s masking provides implicit data augmentation that combats the overfitting plaguing supervised training at scale.

ViT-H (632M parameters) reaches 86.9% fine-tuning accuracy, and fine-tuning at 448×448 resolution pushes to 87.8% — the best result on ImageNet-1K using only that dataset for both pre-training and evaluation. Supervised ViT-H is extremely difficult to train on ImageNet-1K due to severe overfitting, but MAE’s 75% masking provides an effective regularizer: with three-quarters of the input removed at each step, the model sees many different patch subsets across training. The ability to train such large models with only self-supervision and no external data marks a paradigm shift for vision, demonstrating that scale and self-supervised objectives can substitute for labeled data.

How MAE Compares

MAE’s approach differs fundamentally from contrastive and self-distillation methods:

Key Results

ImageNet Classification

Transfer Learning

MAE pre-trained features transfer effectively to dense prediction tasks that demand spatially aware representations. On COCO object detection, a ViT-L backbone with Mask R-CNN achieves 53.3 AP^box, demonstrating that pixel reconstruction learns features rich enough for precise object localization. On ADE20K semantic segmentation, ViT-L with UperNet reaches 53.6 mIoU, confirming that MAE’s representations capture both local texture and global scene structure. These transfer results are particularly noteworthy because MAE’s pretext task — reconstructing masked pixels — directly encourages the encoder to build spatially coherent internal representations, unlike contrastive objectives that primarily optimize for global image-level similarity.

Why MAE Matters

MAE demonstrated that masked image modeling can learn transferable visual representations with a recipe adapted to images. The studied recipe works well with 75% masking, normalized pixel targets perform competitively with discrete tokens, and asymmetric encoding saves training compute. With the default 8-block decoder, Table 2 reports a 3.5× wall-clock speedup for ViT-H (632M parameters) relative to encoding mask tokens too, making self-supervised pre-training more efficient; the paper benchmarks training on 128 TPU-v3 cores.

MAE’s simplicity is its greatest strength. No contrastive pairs, no momentum encoder, no tokenizer, no negative sampling — just mask, encode visible patches, decode, and reconstruct pixels. This simplicity makes MAE easy to implement, fast to train, and reliable to scale. It demonstrated that the Vision Transformer itself is a powerful learner when given the right pretext task, catalyzing the shift toward masked image modeling as the dominant self-supervised paradigm for vision. BEiT is earlier related work; later approaches such as SimMIM and I-JEPA explore other ways to learn from missing image content, including different targets and architectures.

Key Takeaways

1. 75% random masking works well in this recipe — far beyond BERT’s 15% for language, because images have massive spatial redundancy that requires extreme masking to create a challenging pretext task.

2. Asymmetric encoding saves training compute — Table 1c reports 3.3× fewer FLOPs; Table 2 measures 2.8× wall-clock speedup for ViT-L and 3.5× for ViT-H with an 8-block decoder.

3. Decoder depth affects evaluation protocols differently — one to eight blocks changes fine-tuning by 0.1 points and linear probing by 8.0 points in Table 1a.

4. MAE scales better than supervised training — ViT-H reaches 86.9% with MAE pre-training on ImageNet-1K alone, where supervised training overfits. The gap between MAE and supervised widens with model size.

5. Simple pixel reconstruction works — no tokenizer, no contrastive pairs, no momentum encoder. Per-patch normalized MSE on masked patches produces representations that transfer effectively to detection (53.3 AP^box) and segmentation (53.6 mIoU).

Related Reading

• SimCLR — Contrastive framework that inspired many SSL improvements
• MoCo — Momentum contrast for building large consistent dictionaries
• BYOL — Self-supervised learning without negative pairs
• DINO — Self-distillation with Vision Transformers
• VICReg — Variance-invariance-covariance regularization
• V-JEPA — Joint-embedding predictive architecture for video