Features & Representation

Masked Autoencoder

A self-supervised pre-training method for Vision Transformers that hides 75% of an image's patches, encodes only the visible ones, and trains a light decoder to reconstruct the missing pixels; the encoder is then fine-tuned for recognition, detection and segmentation.

advanced

A masked autoencoder (MAE) pre-trains a Vision Transformer without labels by deleting most of an image’s patches and asking the network to paint them back in. Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár and Ross Girshick of Facebook AI Research posted it to arXiv in November 2021 and presented it at CVPR 2022. It is an encoder–decoder with two deliberate asymmetries: the encoder sees only the visible quarter of the patches, and the decoder that reconstructs the rest is small and is thrown away after pre-training. It belongs to the masked image modeling family surveyed in self-supervised learning. Implementation details, the linear-probing comparison and robustness tests are in the appendix (arXiv v3 and the CVF supplementary material); the main text matches the CVPR paper.

Problem Addressed

Masked prediction had made BERT-style pre-training standard in language, and the idea is older in vision: denoising autoencoders treated masking as one kind of corruption (Vincent et al., 2008), and context encoders inpainted a large missing region with a convolutional network (Pathak et al., 2016). Yet by 2021 autoencoding in vision lagged behind language, and contrastive methods were the popular choice. He et al. pointed to three differences:

  • Architecture. Convolutions have no natural slot for mask tokens or position embeddings; ViT removed that obstacle.
  • Information density. A sentence with a few words missing is a hard puzzle, but an image is spatially redundant: a missing patch can be recovered from its neighbors with little understanding of the scene. The remedy is to mask so much that local interpolation fails.
  • The decoder’s role. BERT’s decoder can be an MLP because missing words are semantic; pixels are not, so the decoder’s design decides how abstract the encoder’s representation becomes.

The practical gap was scale. Supervised ViT-Large and ViT-Huge overfit ImageNet-1K, and BEiT (Bao et al., 2022) needed a separately trained discrete VAE tokenizer and ran its encoder over every patch, mask tokens included.

Architecture

Masking

The image is cut into non-overlapping patches as in ViT. Patches are sampled uniformly at random without replacement, and the default masking ratio is 75%: with 16×1616 \times 16 patches of a 224×224224 \times 224 image, the encoder sees 49 of 196 patches. The uniform distribution avoids a bias toward masking the center. No sparse operations are needed: the token list is shuffled, truncated, and later unshuffled, as the NumPy example below shows.

Encoder

The encoder is a standard ViT applied only to the visible patches: linear patch embedding, added position embeddings (fixed sine–cosine ones, per the appendix), Transformer blocks, and a final layer normalization. Masked patches are removed, not replaced by mask tokens, so the encoder works on a quarter of the sequence and never sees an input that cannot occur at test time. A dummy class token keeps ViT’s classifier interface; average pooling worked about as well.

Decoder

A linear layer maps the encoded tokens to the decoder width. The decoder input is the full token set: the encoded visible patches and, for each masked patch, one shared learned mask token. Position embeddings are added again, since without them a mask token would not know which patch it stands for. A narrower, shallower stack of Transformer blocks follows, ending in a linear layer with one output per pixel value in a patch. The default has 8 blocks of width 512, about 9% of ViT-L’s FLOPs per token. Recognition uses the encoder alone.

Reconstruction Target and Loss

For patch ii with pixel vector xi∈RP2C\mathbf{x}_i \in \mathbb{R}^{P^2 C}, the normalized-pixel variant uses the patch’s own mean μi\mu_i and variance σi2\sigma_i^2 to form the target ti=(xi−μi)/σi2+ϵ\mathbf{t}_i = (\mathbf{x}_i - \mu_i)/\sqrt{\sigma_i^2 + \epsilon}. With M\mathcal{M} the set of masked patches and x^i\hat{\mathbf{x}}_i the decoder output,

L=1∣M∣∑i∈M1P2C∥x^i−ti∥22,\mathcal{L} = \frac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \frac{1}{P^2 C} \left\lVert \hat{\mathbf{x}}_i - \mathbf{t}_i \right\rVert_2^2 ,

the mean squared error over masked patches only. Computing it over all pixels, as classical denoising autoencoders do, cost about 0.5% accuracy. The authors use unnormalized pixels (ti=xi\mathbf{t}_i = \mathbf{x}_i) in their ablations and normalized pixels in their final models, observing that per-patch normalization enhances local contrast.

The NumPy script below runs this bookkeeping on a synthetic 64×6464 \times 64 image cut into 16 patches of 16×1616 \times 16. The masking follows the official random_masking function, and the normalization uses the unbiased variance, as torch.var does in the official loss. A crude “decoder” fills each masked patch with the nearest visible one.

import numpy as np

rng = np.random.default_rng(0)

# A 64x64 grayscale synthetic image: a bright square and a dark disc on a diagonal gradient.
S, P = 64, 16
yy, xx = np.mgrid[0:S, 0:S]
img = 0.2 + 0.5 * (xx + yy) / (2 * (S - 1))
img[6:30, 34:58] = 0.95                                    # square
img[(yy - 44) ** 2 + (xx - 20) ** 2 < 15 ** 2] = 0.05      # disc

def patchify(x, P):
    """(H, W) -> (L, P*P): non-overlapping patches in raster order."""
    h, w = x.shape
    return x.reshape(h // P, P, w // P, P).transpose(0, 2, 1, 3).reshape(-1, P * P)

def unpatchify(p, P, h, w):
    return p.reshape(h // P, w // P, P, P).transpose(0, 2, 1, 3).reshape(h, w)

x = patchify(img, P)                                      # (L, P*P) = (16, 256)
L = len(x)

# Random masking exactly as in MAE's random_masking: argsort of uniform noise.
mask_ratio = 0.75
len_keep = int(L * (1 - mask_ratio))                      # 4 visible patches
noise = rng.random(L)
ids_shuffle = np.argsort(noise)                           # small noise = keep
ids_restore = np.argsort(ids_shuffle)                     # inverse permutation
ids_keep = ids_shuffle[:len_keep]
x_visible = x[ids_keep]                                   # what the encoder sees: (4, 256)
mask = np.ones(L); mask[:len_keep] = 0                    # 0 = keep, 1 = remove (shuffled order)
mask = mask[ids_restore]                                  # back to raster order

print("ids_shuffle:", ids_shuffle)
print("ids_restore:", ids_restore)
print("visible patches (raster index):", np.sort(ids_keep))
print("mask on the 4x4 grid (1 = masked):")
print(mask.reshape(4, 4).astype(int))

# Decoder input: encoded visible tokens + one shared mask token per removed patch, then unshuffle.
# Here the "encoding" is the identity and the mask token is a vector of NaNs, to make placement visible.
mask_token = np.full(P * P, np.nan)
full = np.concatenate([x_visible, np.repeat(mask_token[None], L - len_keep, 0)])
full = full[ids_restore]                                  # align every token with its patch position
vis = ~np.isnan(full[:, 0])
print("unshuffled visible tokens equal the original patches:", np.array_equal(full[vis], x[vis]))
print("mask tokens land exactly on the masked positions:", np.array_equal(~vis, mask.astype(bool)))

# Per-patch normalized target (norm_pix_loss): mean and unbiased variance of each patch.
mean = x.mean(-1, keepdims=True)
var = x.var(-1, ddof=1, keepdims=True)
target = (x - mean) / np.sqrt(var + 1e-6)

# A crude "decoder": fill each masked patch with the visible patch nearest to it on the grid.
gy, gx = np.divmod(np.arange(L), S // P)
pred_pix = x.copy()
for i in np.flatnonzero(mask):
    j = ids_keep[np.argmin((gy[ids_keep] - gy[i]) ** 2 + (gx[ids_keep] - gx[i]) ** 2)]
    pred_pix[i] = x[j]
pred_norm = (pred_pix - pred_pix.mean(-1, keepdims=True)) / np.sqrt(pred_pix.var(-1, ddof=1, keepdims=True) + 1e-6)

def mae_loss(pred, target, mask, masked_only=True):
    per_patch = ((pred - target) ** 2).mean(-1)            # mean over the pixels of each patch
    w = mask if masked_only else np.ones_like(mask)
    return (per_patch * w).sum() / w.sum()

print(f"\nnearest-visible-patch baseline")
print(f"  raw pixels,        masked patches only: {mae_loss(pred_pix, x, mask):.4f}")
print(f"  raw pixels,        all 16 patches:      {mae_loss(pred_pix, x, mask, False):.4f}")
print(f"  normalized pixels, masked patches only: {mae_loss(pred_norm, target, mask):.4f}")
print(f"  predicting 0 (each patch's mean) in normalized space: {mae_loss(np.zeros_like(target), target, mask):.4f}")
print("reassembled image equals the original:", np.array_equal(unpatchify(x, P, S, S), img))
ids_shuffle: [11  3 13  2 15  1  8  6  0  7 14  4 10 12  5  9]
ids_restore: [ 8  5  3  1 11 14  7  9  6 15 12  0 13  2 10  4]
visible patches (raster index): [ 2  3 11 13]
mask on the 4x4 grid (1 = masked):
[[1 1 0 0]
 [1 1 1 1]
 [1 1 1 0]
 [1 0 1 1]]
unshuffled visible tokens equal the original patches: True
mask tokens land exactly on the masked positions: True

nearest-visible-patch baseline
  raw pixels,        masked patches only: 0.1160
  raw pixels,        all 16 patches:      0.0870
  normalized pixels, masked patches only: 1.5014
  predicting 0 (each patch's mean) in normalized space: 0.9956
reassembled image equals the original: True

ids_restore[j] is the shuffled-list position of the token that belongs at patch jj, so indexing with it puts every token back in place. Counting the four trivially copied visible patches lowers the raw-pixel loss from 0.116 to 0.087 with no better prediction. In normalized space, copying a neighbor (1.50) scores worse than predicting a flat patch (0.996, close to the 255/256255/256 that a flat prediction scores on any non-constant 256-pixel patch). In our reading, this is the squared error’s preference for hedged, averaged predictions, which shows up as blur in reconstructions.

Training

Pre-training uses the ImageNet-1K training images without labels. The appendix gives AdamW with β2=0.95\beta_2 = 0.95, weight decay 0.05, batch 4,096, base learning rate 1.5×10−41.5 \times 10^{-4} scaled linearly with batch size, cosine decay with 40 warm-up epochs, and random resized crops (with horizontal flips) as the only augmentation. Ablations use ViT-L/16 for 800 epochs; the final models train for 1,600. Fine-tuning (50 epochs for ViT-L and ViT-H) uses layer-wise learning-rate decay, RandAugment, mixup and CutMix. After it, the baseline MAE ViT-L reached 84.9%, against 82.5% for the authors’ 200-epoch supervised ViT-L (82.6% in the appendix and Table 3).

Ablations

All with ViT-L/16, giving ImageNet-1K top-1 for end-to-end fine-tuning / linear probing (Table 1 and Figure 5 of He et al., 2022):

Factor Setting Fine-tune Linear
Masking ratio 10% 83.2 54.6
75% (default) 84.9 73.5
90% 83.0 66.1
Decoder depth 1 block 84.8 65.5
8 blocks (default) 84.9 73.5
Decoder width 128 84.9 69.1
1,024 84.3 73.1
Mask token in the encoder (3.3× FLOPs) 84.2 59.6
Target pixels, normalized 85.4 73.9
PCA coefficients 84.6 72.3
dVAE tokens (as in BEiT) 85.3 71.6
Augmentation none 84.0 65.7
crop + color jitter 84.3 71.9
Mask sampling block-wise, 75% 82.8 63.9
grid (keep 1 in 4) 84.0 66.0

Fine-tuning tolerates masking ratios from 40% to 80%, far above BERT’s 15% and the 20–50% of earlier vision work. A deep decoder matters mainly for linear probing; the authors’ explanation is that it absorbs the reconstruction-specific layers, leaving the encoder’s output more abstract. Mask tokens in the encoder cost 14 points of linear accuracy, which the authors attribute to the mismatch between masked pre-training inputs and real test images. Normalized pixels beat raw pixels and PCA, suggesting to the authors that high-frequency content helps, and match tokens without a tokenizer. Block-wise masking (BEiT’s) made the task harder and the reconstructions blurrier, grid masking easier and sharper; both learned worse representations. Unlike contrastive methods, MAE needs little augmentation, because random masks already create new samples every iteration. From 100 to 1,600 epochs, ViT-L rose from 82.3% to 85.1% fine-tuned and from 57.3% to 75.1% linear, without saturating.

Speed. On 128 TPU-v3 cores, 800 epochs of ViT-L took 42.4 hours with mask tokens in the encoder and 15.4 without (2.8×); with a one-block decoder, ViT-H, or both, the speedup reached 3.5–4.1×, partly because attention cost is quadratic in sequence length, and memory falls too. ViT-L for 1,600 epochs took 31 hours, against 36 for 300 epochs of the contrastive MoCo v3.

A Tiny MAE

The PyTorch script below (version 2.14.1, one CPU thread, about 30 seconds, identical output on two runs) builds a miniature MAE for 32×3232 \times 32 synthetic squares, discs, triangles and crosses: 64 tokens of 4×44 \times 4 pixels, a three-block encoder of width 64 and a one-block decoder of width 32. It pre-trains for 2,000 steps on 4,000 unlabeled images, then compares the encoder with its own random initialization on 25 labelled images per shape, by a linear probe on frozen average-pooled features and by 300 steps of fine-tuning.

import math
import sys
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

torch.set_num_threads(1)                     # single-threaded CPU run: identical output on every run
print(f"torch {torch.__version__}")

def make_shapes(n, rng, size=32):
    """n grayscale 32x32 images of four classes (square, disc, triangle, cross) with random
    size, brightness and position on a flat gray background."""
    yy, xx = np.mgrid[0:size, 0:size].astype(np.float32)
    X = np.empty((n, 1, size, size), np.float32)
    y = rng.integers(0, 4, n)
    for i in range(n):
        r = rng.uniform(6, 10)                                     # half-width
        cy, cx = rng.uniform(r + 1, size - r - 1, 2)
        dy, dx = yy - cy, xx - cx
        if y[i] == 0:   m = (abs(dy) < r) & (abs(dx) < r)
        elif y[i] == 1: m = dy ** 2 + dx ** 2 < r ** 2
        elif y[i] == 2: m = (abs(dy) < r) & (abs(dx) < (dy + r) / 2)
        else:           m = ((abs(dy) < r) & (abs(dx) < r / 3)) | ((abs(dx) < r) & (abs(dy) < r / 3))
        X[i, 0] = 0.4 + m * rng.uniform(0.3, 0.6)
    return torch.from_numpy(X), torch.from_numpy(y)

def sincos_2d(d, grid):
    """Fixed 2D sine-cosine position embeddings (half the channels for y, half for x)."""
    def one_d(pos, d):
        omega = 1.0 / 10000 ** (torch.arange(d // 2) / (d / 2))
        out = pos[:, None] * omega[None]
        return torch.cat([out.sin(), out.cos()], 1)
    ys, xs = torch.meshgrid(torch.arange(grid).float(), torch.arange(grid).float(), indexing="ij")
    return torch.cat([one_d(ys.flatten(), d // 2), one_d(xs.flatten(), d // 2)], 1)   # (grid*grid, d)

class Block(nn.Module):
    def __init__(self, d, heads=4, mlp=2):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, mlp * d), nn.GELU(), nn.Linear(mlp * d, d))
    def forward(self, z):
        h = self.ln1(z)
        z = z + self.attn(h, h, h, need_weights=False)[0]
        return z + self.mlp(self.ln2(z))

class Encoder(nn.Module):
    """ViT encoder on 4x4 patches of a 32x32 image (64 tokens), no class token."""
    def __init__(self, d=64, depth=3, patch=4, img=32):
        super().__init__()
        self.patch = patch
        self.embed = nn.Linear(patch * patch, d)
        self.register_buffer("pos", sincos_2d(d, img // patch))
        self.blocks = nn.Sequential(*[Block(d) for _ in range(depth)])
        self.norm = nn.LayerNorm(d)
    def forward(self, patches, ids_keep=None):
        t = self.embed(patches) + self.pos                       # (B, L, d)
        if ids_keep is not None:                                 # keep only the visible tokens
            t = torch.gather(t, 1, ids_keep[..., None].expand(-1, -1, t.shape[-1]))
        return self.norm(self.blocks(t))
    def features(self, x):
        return self(patchify(x, self.patch)).mean(1)            # average-pooled, no masking

class MAE(nn.Module):
    def __init__(self, d=64, dec_d=32, patch=4, img=32):
        super().__init__()
        self.encoder = Encoder(d, patch=patch, img=img)
        self.dec_embed = nn.Linear(d, dec_d)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, dec_d))
        self.register_buffer("dec_pos", sincos_2d(dec_d, img // patch))
        self.decoder = Block(dec_d)                              # one decoder block
        self.dec_norm = nn.LayerNorm(dec_d)
        self.pred = nn.Linear(dec_d, patch * patch)              # pixels of one patch
    def forward(self, x, mask_ratio=0.75, zero_pred=False):
        p = patchify(x, 4)
        B, L, _ = p.shape
        keep = int(L * (1 - mask_ratio))
        ids_shuffle = torch.argsort(torch.rand(B, L), 1)
        ids_restore = torch.argsort(ids_shuffle, 1)
        latent = self.encoder(p, ids_shuffle[:, :keep])          # encoder sees 16 of 64 patches
        t = self.dec_embed(latent)
        t = torch.cat([t, self.mask_token.expand(B, L - keep, -1)], 1)
        t = torch.gather(t, 1, ids_restore[..., None].expand(-1, -1, t.shape[-1])) + self.dec_pos
        pred = self.pred(self.dec_norm(self.decoder(t)))
        if zero_pred:                                            # baseline: predict each patch's mean
            pred = torch.zeros_like(pred)
        mask = torch.ones(B, L); mask[:, :keep] = 0
        mask = torch.gather(mask, 1, ids_restore)
        mu, var = p.mean(-1, keepdim=True), p.var(-1, keepdim=True)
        target = (p - mu) / (var + 1e-6).sqrt()                 # per-patch normalized pixels
        return (((pred - target) ** 2).mean(-1) * mask).sum() / mask.sum()

def patchify(x, P):
    B, C, H, W = x.shape
    return x.reshape(B, C, H // P, P, W // P, P).permute(0, 2, 4, 3, 5, 1).reshape(B, -1, P * P * C)

def pretrain(mae, X, steps=2000, bs=32, lr=2e-3, warmup=20):
    opt = torch.optim.AdamW(mae.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.05)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min((s + 1) / warmup,
                0.5 * (1 + math.cos(math.pi * (s - warmup) / (steps - warmup)))))
    for _ in range(steps):
        loss = mae(X[torch.randint(0, len(X), (bs,))])
        opt.zero_grad(); loss.backward(); opt.step(); sched.step()

def linear_probe(enc, Xtr, ytr, Xte, yte, steps=300):
    """Logistic regression on frozen, standardized average-pooled features."""
    with torch.no_grad():
        ftr, fte = enc.features(Xtr), enc.features(Xte)
    mu, sd = ftr.mean(0), ftr.std(0) + 1e-6
    ftr, fte = (ftr - mu) / sd, (fte - mu) / sd
    head = nn.Linear(ftr.shape[1], 4)
    opt = torch.optim.Adam(head.parameters(), lr=0.01, weight_decay=1e-3)
    for _ in range(steps):
        loss = F.cross_entropy(head(ftr), ytr)
        opt.zero_grad(); loss.backward(); opt.step()
    return (head(fte).argmax(1) == yte).float().mean().item()

def fine_tune(enc, Xtr, ytr, Xte, yte, steps=300, bs=32):
    """Train the encoder and a linear head end to end on the few labels."""
    head = nn.Linear(64, 4)
    opt = torch.optim.AdamW(list(enc.parameters()) + list(head.parameters()), lr=1e-3, weight_decay=0.05)
    for _ in range(steps):
        i = torch.randint(0, len(ytr), (bs,))
        loss = F.cross_entropy(head(enc.features(Xtr[i])), ytr[i])
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        return (head(enc.features(Xte)).argmax(1) == yte).float().mean().item()

def run(seed, n_labels=25):
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    X_unlab, _ = make_shapes(4000, rng)                          # labels thrown away
    X_eval, _ = make_shapes(500, rng)
    pool, ypool = make_shapes(400, rng)
    idx = torch.cat([torch.nonzero(ypool == c)[:n_labels, 0] for c in range(4)])
    Xtr, ytr = pool[idx], ypool[idx]                             # 25 labelled images per class
    Xte, yte = make_shapes(1000, np.random.default_rng(1000))   # fixed test set

    mae = MAE()
    random_state = {k: v.clone() for k, v in mae.encoder.state_dict().items()}
    torch.manual_seed(seed + 2)                                  # same masks for all three evaluations
    with torch.no_grad():
        zero = mae(X_eval, zero_pred=True).item()
    torch.manual_seed(seed + 2)
    with torch.no_grad():
        before = mae(X_eval).item()
    pretrain(mae, X_unlab)
    torch.manual_seed(seed + 2)
    with torch.no_grad():
        after = mae(X_eval).item()

    rand_enc = Encoder(); rand_enc.load_state_dict(random_state)  # the same init, never trained
    res = {"loss_zero": zero, "loss_before": before, "loss_after": after,
           "lin_rand": linear_probe(rand_enc, Xtr, ytr, Xte, yte),
           "lin_mae": linear_probe(mae.encoder, Xtr, ytr, Xte, yte)}
    torch.manual_seed(seed + 1)
    res["ft_rand"] = fine_tune(rand_enc, Xtr, ytr, Xte, yte)
    torch.manual_seed(seed + 1)
    res["ft_mae"] = fine_tune(mae.encoder, Xtr, ytr, Xte, yte)
    return res

if __name__ == "__main__":
    seeds = [int(s) for s in sys.argv[1:]] or [6]
    for s in seeds:
        r = run(s)
        print(f"seed {s}: loss on masked patches of 500 held-out images (normalized pixels)")
        print(f"  predict each patch's mean {r['loss_zero']:.3f}, untrained MAE {r['loss_before']:.3f}, "
              f"after 2,000 steps {r['loss_after']:.3f}")
        print(f"  shape accuracy, 100 labels, 1,000 test images (chance 0.25):")
        print(f"    linear probe  random-init encoder {r['lin_rand']:.3f}   MAE encoder {r['lin_mae']:.3f}")
        print(f"    fine-tuning   random-init encoder {r['ft_rand']:.3f}   MAE encoder {r['ft_mae']:.3f}")
torch 2.14.1
seed 6: loss on masked patches of 500 held-out images (normalized pixels)
  predict each patch's mean 0.176, untrained MAE 0.516, after 2,000 steps 0.132
  shape accuracy, 100 labels, 1,000 test images (chance 0.25):
    linear probe  random-init encoder 0.437   MAE encoder 0.567
    fine-tuning   random-init encoder 0.530   MAE encoder 0.674

Over seeds 0–9, the masked-patch loss ended at 0.127–0.134, against 0.171–0.180 for predicting every patch flat; background patches have a zero target, so only shape edges leave room to improve. Linear probing gave 0.437–0.485 for the random encoder and 0.449–0.666 for MAE (means 0.46 and 0.56; MAE lost on one seed). Fine-tuning gave 0.458–0.588 and 0.591–0.732 (means 0.51 and 0.66), MAE winning on every seed; seed 6 is close to the mean gaps.

Inference

The decoder is discarded. The encoder takes the complete, unmasked image and costs exactly a ViT of the same size. For detection and segmentation it becomes the backbone of a task head and is fine-tuned end to end, the usual transfer learning recipe.

Performance

All results from He et al. (2022) unless stated.

ImageNet-1K classification. Fine-tuned at 224×224224 \times 224 after 1,600 epochs with normalized-pixel targets, MAE reached 83.6% with ViT-B/16, 85.9% with ViT-L/16 and 86.9% with ViT-H/14, against 83.2% and 85.2% for BEiT and 83.2% and 84.1% for MoCo v3 (Table 3). Fine-tuned at 448×448448 \times 448, ViT-H reached 87.8%, above the previous best of 87.1% among methods using only ImageNet-1K data. The authors’ supervised ViT-H from scratch reached 83.1%; the gain from MAE grew with model size.

Linear probing versus partial fine-tuning. MAE features are less linearly separable than contrastive ones: linear probing gave 68.0% (ViT-B), 75.8% (ViT-L) and 76.6% (ViT-H), against 77.6% for MoCo v3’s ViT-L (appendix, Table 12). Fine-tuning only the last Transformer block lifted the 800-epoch ViT-L from 73.5% to 81.0%, and tuning four to six blocks came close to full fine-tuning; with one or more blocks tuned, MAE beat MoCo v3 at every depth (Figure 9), so the authors argue that linear separability is not the only measure of representation quality.

Detection and instance segmentation. For COCO, the ViT was adapted to a feature pyramid by splitting its blocks into four groups and resampling their outputs to strides 4–32, inside a Mask R-CNN, with hyperparameters searched for every entry (appendix A.3). Box AP for ViT-B / ViT-L was 50.3 / 53.3 for MAE, against 47.9 / 49.3 for supervised pre-training and 49.8 / 53.3 for BEiT; mask AP, which measures instance segmentation, was 44.9 / 47.2, against 42.9 / 43.9 supervised (Table 4). ViTDet later built a simple pyramid from the last ViT feature map alone; with Mask R-CNN its MAE-pre-trained ViT-L reached 54.6 box AP, against 49.6 with supervised ImageNet-1K pre-training and 50.0 from random initialization (Li et al., 2022, Table 4).

Semantic segmentation. On ADE20K with UperNet, mIoU for ViT-B / ViT-L was 48.1 / 53.6 for MAE, 47.4 / 49.9 supervised and 47.1 / 53.3 for BEiT (Table 5).

Across ImageNet, object detection and semantic segmentation, normalized pixels and BEiT’s dVAE tokens performed the same (Table 7), so the tokenizer is unnecessary.

Strengths

  • Cheap pre-training for large encoders: a quarter of the tokens and 2.8–4.1× wall-clock speedups.
  • Simplicity: pixel targets, no tokenizer, negatives or momentum encoder, and almost no augmentation.
  • Scaling: gains grow with model size and schedule length.
  • Robustness (appendix): ViT-H reached 68.2% on ImageNet-A, against 33.1% for its supervised counterpart.

Limitations

  • Weak frozen features. Linear probing trails contrastive and self-distillation methods. DINOv2 reported 76.6% linear for MAE ViT-H/14 against 86.5% for its ViT-g/14, and 33.3 mIoU for a linear head on frozen MAE features on ADE20K against 49.0 (Oquab et al., 2024, Tables 4 and 10), noting that MAE features need supervised fine-tuning.
  • Blurry reconstructions. The paper presents reconstructions as plausible rather than faithful, blurrier with block masking; in our reasoning, a squared-error loss rewards averaging over plausible completions. MAE is not an image generator, and the authors warn that it may produce nonexistent content.
  • Compute. 1,600 epochs remain costly for large models, and heavy fine-tuning recipes are part of the reported accuracy.

Historical Importance

MAE (2021–2022) showed that the simplest masked image modeling, raw pixels and random masks, scales to ViT-H on ImageNet-1K alone, and its weights went on to initialize backbones such as ViTDet’s and SAM’s. Its asymmetric design made pre-training cheaper rather than costlier. Its ablations tied the best masking ratio to the data’s redundancy, which the video follow-ups supported (90%), and its partial fine-tuning analysis argued for evaluating beyond linear probing.

Variants

  • Video. VideoMAE (Tong et al., 2022) masks the same spatial positions in every frame (“tube masking”) at 90–95%, which the authors attribute to temporal redundancy; it reported 87.4% on Kinetics-400 and 75.4% on Something-Something V2 without extra data. Feichtenhofer et al. (2022) applied MAE to random spacetime patches and found 90% masking optimal, with speedups above 4×.
  • Audio. Audio-MAE (Huang et al., 2022) masks 80% of spectrogram patches and adds local window attention to the decoder.
  • Segmentation backbones. SAM’s image encoder is an MAE-pre-trained ViT-H/16 adapted to 1024×10241024 \times 1024 inputs with windowed attention (Kirillov et al., 2023).
  • Self-distillation. iBOT and DINOv2 predict a teacher network’s outputs for masked tokens instead of pixels.

Implementations

The authors’ PyTorch re-implementation and ViT-B, ViT-L and ViT-H checkpoints, converted from the original TensorFlow/TPU code, are in the facebookresearch/mae repository; its models_mae.py contains the masking, decoder and loss that the NumPy example mirrors. Video and audio variants are in MCG-NJU/VideoMAE, facebookresearch/mae_st and facebookresearch/AudioMAE.

Related

  • Contrastive Learning

    Training an encoder so that embeddings of matching pairs, such as two augmented views of one image, are close and embeddings of non-matching pairs are far apart, through margin losses, the InfoNCE and NT-Xent softmax losses, and recipes such as SimCLR and MoCo.

  • Transfer Learning

    Reusing what a network learned on a large source task, typically by pretraining on many labelled images and then adapting it to a target task with few labels, through a linear probe on frozen features, partial fine-tuning or full fine-tuning.

  • Encoder–Decoder

    A network design in which an encoder compresses the input into a coarse, many-channel representation and a decoder expands it back into the output space, usually at the input's resolution; the standard architecture for segmentation, depth, flow, image translation and pose heatmaps, and, without the spatial grid, for sequence-to-sequence models.

  • SAM

    The Segment Anything Model, a promptable segmentation model that returns object masks for points, boxes or rough masks, trained with a model-in-the-loop data engine on the SA-1B dataset of 1.1 billion masks, and the starting point of a family of faster, sharper, video and medical variants.

  • Mask R-CNN

    An instance segmentation model that adds a small fully convolutional mask branch and RoIAlign to the Faster R-CNN detector, predicting a box, a class and a binary mask for every detected object.

  • Semantic Segmentation

    Labeling every pixel of an image with one of a fixed set of classes, such as road, sky, car or person, without separating individual objects; the dense-prediction counterpart of image classification, also called scene parsing.

References

  1. He, K., Chen, X., Xie, S., Li, Y., Dollár, P. & Girshick, R. (2022). Masked Autoencoders Are Scalable Vision Learners. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 15979–15988.
  2. He, K., Chen, X., Xie, S., Li, Y., Dollár, P. & Girshick, R. (2021). Masked Autoencoders Are Scalable Vision Learners. arXiv:2111.06377 (v3, with appendix).
  3. Vincent, P., Larochelle, H., Bengio, Y. & Manzagol, P.-A. (2008). Extracting and Composing Robust Features with Denoising Autoencoders. Proceedings of the 25th International Conference on Machine Learning (ICML), 1096–1103.
  4. Pathak, D., Krähenbühl, P., Donahue, J., Darrell, T. & Efros, A. A. (2016). Context Encoders: Feature Learning by Inpainting. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2536–2544.
  5. Bao, H., Dong, L., Piao, S. & Wei, F. (2022). BEiT: BERT Pre-Training of Image Transformers. International Conference on Learning Representations (ICLR). arXiv:2106.08254.
  6. Li, Y., Mao, H., Girshick, R. & He, K. (2022). Exploring Plain Vision Transformer Backbones for Object Detection. European Conference on Computer Vision (ECCV), Lecture Notes in Computer Science, 280–296.
  7. Tong, Z., Song, Y., Wang, J. & Wang, L. (2022). VideoMAE: Masked Autoencoders are Data-Efficient Learners for Self-Supervised Video Pre-Training. Advances in Neural Information Processing Systems 35 (NeurIPS).
  8. Feichtenhofer, C., Fan, H., Li, Y. & He, K. (2022). Masked Autoencoders As Spatiotemporal Learners. Advances in Neural Information Processing Systems 35 (NeurIPS).
  9. Huang, P.-Y., Xu, H., Li, J., Baevski, A., Auli, M., Galuba, W., Metze, F. & Feichtenhofer, C. (2022). Masked Autoencoders that Listen. Advances in Neural Information Processing Systems 35 (NeurIPS).
  10. Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A. C., Lo, W.-Y., Dollár, P. & Girshick, R. (2023). Segment Anything. IEEE/CVF International Conference on Computer Vision (ICCV), 3992–4003.
  11. Oquab, M., Darcet, T., Moutakanni, T., Vo, H. V., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El-Nouby, A., et al. (2024). DINOv2: Learning Robust Visual Features without Supervision. Transactions on Machine Learning Research (TMLR). arXiv:2304.07193.