← All tutorials

Masked Image Modeling

On this page

1 Masking image patches works like masking words 2 A Vision Transformer reads an image as a row of patches 3 Images are redundant, so most patches must be hidden 4 A big encoder sees only the visible quarter 5 The loss scores only the hidden patches 6 Probing and fine-tuning disagree on masked models 7 MAE in practice 8 References

Masked image modeling hides most of an image’s patches (small square tiles of pixels) and trains an encoder (the network turning an image into a vector) to help rebuild it. Every step draws a fresh random mask, and the loss is the pixel error on the hidden parts. You keep the encoder and drop the small decoder. The method here is MAE (He et al. 2022).

patches of one image, 75% hiddenbig encoder, a ViTsees the visible patches only[M][M][M][M][M][M]encoded patches + mask tokens [M] in placesmall decoderdroppedafterwardspixels for every patch; red outline = scored
Figure 1. Only visible patches enter the big encoder; mask tokens (placeholders for hidden patches) join after it.

1 Masking image patches works like masking words

The text model BERT hides 15% of the words, mostly behind a mask token (one learned vector meaning "missing"), and predicts them. MAE does the same with patches; as an autoencoder (encode the input, decode it back) the image is its own label. Unlike Contrastive Learning and Non-Contrastive Self-Supervised Learning, MAE never compares two views of an image, so there is no collapse (every image mapped to the same output) to prevent.

text, few words hiddenthecat[M]onthematBERT hides about 15%and predicts the wordimage, most hiddenMAE hides 75%
Figure 2. BERT hides about 15% of words; MAE hides 75% of patches.

2 A Vision Transformer reads an image as a row of patches

A Vision Transformer (ViT) cuts an image into 16 by 16 pixel patches, turns each into a vector, a token, plus a positional embedding saying where it sat (as in Positional Encoding and Masking), and mixes the tokens with attention.

224 × 224 pixels196 patches of 16 × 16one vector per patch...196 tokensTransformer
Figure 3. To the Transformer an image is a row of 196 tokens, one per patch.

3 Images are redundant, so most patches must be hidden

Images are full of repeats: a missing patch “can be recovered from neighboring patches with little high-level understanding” (the MAE paper), so MAE hides most patches at random; at a 75% mask ratio, 49 of 196 stay visible.

MAE Figure 2, two ImageNet validation images: a cheetah and a glass cabinet, each shown with 80% of patches hidden, the MAE reconstruction, and the original
Figure 4. From the paper’s Figure 2: “… The masking ratio is 80%, leaving only 39 out of 196 patches.”

MAE picked the ratio by pretraining a large ViT for 800 epochs on ImageNet-1K at each ratio, then scoring a linear probe (one linear layer on frozen features) and fine-tuning of every weight.

Top-1 accuracy against mask ratio506070809010%30%50%75%90%share of patches hiddenfine-tunelinear probepeak 73.5grid 66.0
Figure 5. The probe peaks at 75%; the square is grid masking at 75% (MAE Fig. 5, Table 1f).

A hidden patch next to visible ones can be filled by copying a neighbor; one far from any visible patch must be inferred from the rest of the scene. Grid masking (keep 1 patch in 4) leaves more hidden patches beside visible ones, and its probe drops (the square in Figure 5).

Grid always hides 75%.

Figure 6. Orange marks hidden patches beside 2 or more visible ones; at 75%, Grid has over twice as many as Random.

visiblehidden, beside 2+ visiblehidden, beside 0 or 1 visible

4 A big encoder sees only the visible quarter

The encoder, a big ViT, reads only the visible patches and never sees a mask token. A much smaller decoder then adds a mask token in every hidden spot, outputs pixels for every patch, and is dropped afterwards.

Real images never contain mask tokens, so the encoder should never see them. BEiT (Bao et al. 2022), the first BERT-style masked image model for ViTs, did feed them to its encoder; it also hid about 40% of patches in blocks and predicted discrete visual tokens (codes from a learned image tokenizer) instead of pixels.

Linear probe top-1 (%), ViT-Lmask tokens kept out of the encoder73.5mask tokens fed to the encoder59.6050100
Figure 7. Mask tokens in the encoder, as in BEiT, cost about 14 probe points (Table 1c).

The masking gives each patch a random number, keeps the 49 smallest, and encodes only those.

import torch
import torch.nn as nn

MASK_RATIO = 0.75
N, NUM_PATCHES, D = 8, 196, 768
x = torch.randn(N, NUM_PATCHES, D)
encoder = nn.Linear(D, D)                        # stand-in ViT

n_keep = int(NUM_PATCHES * (1 - MASK_RATIO))     # 49
noise = torch.rand(N, NUM_PATCHES)               # fresh every step
keep_ids = noise.argsort(dim=1)[:, :n_keep]
idx = keep_ids.unsqueeze(-1).expand(-1, -1, D)
visible = torch.gather(x, 1, idx)                # [N, 49, D]
latent = encoder(visible)
  1. 1Toy patch tokens and a stand-in encoder, the big encoder of Figure 1.
  2. 2Keeps a random 49 of 196 patches, a fresh mask each step (Figure 6).
  3. 3Encodes the visible patches only, the encoder in Figure 1.

5 The loss scores only the hidden patches

Only hidden patches are graded, as in BERT. The target is each patch's normalized pixels (minus its mean, over its standard deviation).

The decoder guesses every patch’s pixels, and the loss is the average squared pixel error over the hidden patches only.
ei = 1P ‖x̂i − xi‖2
ei is patch i’s mean squared pixel error. x̂i is the decoder’s guess. xi is the real patch, normalized. ‖ ‖2 squares each pixel’s difference and adds them up. P is the number of pixel values in a patch (16 × 16 × 3 colors = 768), so dividing by it gives a per-pixel average. Squaring punishes large misses most, so a guess that is close everywhere costs less than one that is badly off in a few places. Normalizing each patch removes its overall brightness and contrast, so edges and texture count as much in a dark patch as in a bright one.
L = 1|H| Σi ∈ H ei
The average runs over the hidden set H only. Σi ∈ H adds the error of every patch i in H (∈ reads "in"), and dividing by |H|, the number of hidden patches, turns the sum into an average. Visible patches never enter the sum, so every bit of training signal comes from what the encoder could not see, and the average keeps the loss on one scale at any mask ratio.

The loss on toy pixels hides exactly 147 patches with the same argsort as Section 4; hidden is 1 on H, and the last line averages ei over H only.

import torch

MASK_RATIO, EPS = 0.75, 1e-6
N, NUM_PATCHES, P = 8, 196, 768
x = torch.rand(N, NUM_PATCHES, P)
x_hat = torch.rand(N, NUM_PATCHES, P)            # decoder guess
n_hide = int(NUM_PATCHES * MASK_RATIO)           # 147
hide_ids = torch.rand(N, NUM_PATCHES).argsort(dim=1)[:, :n_hide]
hidden = torch.zeros(N, NUM_PATCHES).scatter(1, hide_ids, 1.0)

mean = x.mean(-1, keepdim=True)
std = (x.var(-1, keepdim=True) + EPS).sqrt()
x_norm = (x - mean) / std                        # per patch
e = ((x_hat - x_norm) ** 2).mean(dim=-1)         # e_i, [N, 196]
L = (e * hidden).sum() / hidden.sum()
  1. 1Toy targets, guesses and the hidden set H of Section 5 eq line 2.
  2. 2Normalized pixels per patch, the real patch xi of Section 5 eq line 1.
  3. 3Error ei, Section 5 eq line 1.
  4. 4Average over H only, Section 5 eq line 2.
Fine-tune top-1 (%), ViT-Lraw pixels as the target84.9normalized pixels as the target85.4808590
Figure 8. Normalized targets add half a point; bars start at 80% (MAE Table 1d).

6 Probing and fine-tuning disagree on masked models

For MAE, probe and fine-tuning results are “largely uncorrelated”. The encoder is a stack of 24 Transformer blocks (attention plus a small MLP), and tuning only the last one closes most of the gap. MAE features are not linearly separable until a few layers adapt, so a linear probe undersells them.

Same MAE encoder, three ways to evaluatelinear probe, encoder frozen73.5fine-tune the last block only81.0fine-tune all 24 blocks84.9050100
Figure 9. ImageNet-1K top-1 (%), ViT-L/16 after 800 epochs of MAE.

7 MAE in practice

One step joins both snippets; the mask token fills the hidden slots only after the encoder, and the position embedding pos stays fixed, as MAE’s fixed table of sine and cosine values does.

import torch
import torch.nn as nn

N, NUM_PATCHES, P, D = 8, 196, 768, 256
MASK_RATIO, EPS = 0.75, 1e-6
embed, to_pixels = nn.Linear(P, D), nn.Linear(D, P)
pos = torch.randn(1, NUM_PATCHES, D)             # fixed; MAE uses sin-cos
mask_token = nn.Parameter(torch.zeros(1, 1, D))
enc_layer = nn.TransformerEncoderLayer(D, nhead=8, batch_first=True)
encoder = nn.TransformerEncoder(enc_layer, num_layers=4)
dec_layer = nn.TransformerEncoderLayer(D, nhead=8, batch_first=True)
decoder = nn.TransformerEncoder(dec_layer, num_layers=1)
model = nn.ModuleList([embed, to_pixels, encoder, decoder])
params = [*model.parameters(), mask_token]
decay = [p for p in params if p.ndim >= 2 and p is not mask_token]
no_decay = [p for p in params if p.ndim < 2 or p is mask_token]
groups = [dict(params=decay, weight_decay=0.05), dict(params=no_decay, weight_decay=0.0)]
opt = torch.optim.AdamW(groups, lr=2.4e-3, betas=(0.9, 0.95))

patches = torch.rand(N, NUM_PATCHES, P)          # one batch
n_keep = int(NUM_PATCHES * (1 - MASK_RATIO))
keep = torch.rand(N, NUM_PATCHES).argsort(dim=1)[:, :n_keep]
idx = keep.unsqueeze(-1).expand(-1, -1, D)
latent = encoder(torch.gather(embed(patches) + pos, 1, idx))

full = mask_token.expand(N, NUM_PATCHES, D)
full = torch.scatter(full, 1, idx, latent)
x_hat = to_pixels(decoder(full + pos))           # [N, 196, P]

hidden = torch.ones(N, NUM_PATCHES).scatter(1, keep, 0.0)
mean = patches.mean(-1, keepdim=True)
std = (patches.var(-1, keepdim=True) + EPS).sqrt()
e = ((x_hat - (patches - mean) / std) ** 2).mean(dim=-1)
L = (e * hidden).sum() / hidden.sum()
opt.zero_grad()
L.backward()
opt.step()
  1. 1The parts of Figure 1, and AdamW with no weight decay on biases, norms or the mask token.
  2. 2Masks, then encodes the visible patches, the encoder in Figure 1.
  3. 3Mask tokens fill the hidden slots and the decoder outputs pixels, the decoder in Figure 1.
  4. 4Normalized pixel loss on H, eq lines 1 and 2 of Section 5.
  5. 5One update of every learned part in Figure 1.
Actsevery step; a fresh mask per image
Keepsonly the weights and the mask token; the mask is redrawn every step
Plugs inpretraining; then drop the decoder and fine-tune (50 epochs, ViT-L) or probe (90)
Settingsbatch 4096, AdamW, base lr 1.5×10−4 (times batch/256), weight decay 0.05, 40 warm-up epochs, 800 or 1600 epochs
CostViT-L, 1600 epochs in 31 hours on 128 TPU v3 cores; skipping the hidden 75% trains 2.8 times faster (sec. 4; speedup, Table 2)
Codefacebookresearch/mae
Gotchasmask tokens stay out of the encoder; the official script needs --norm_pix_loss (off by default); no weight decay on biases, norms or position embeddings
Knob, defaultRaise itLower it
Mask ratio, the share of patches in H; 75%fewer hidden patches next to visible ones; past 75% the probe falls (Fig. 5); at 100% nothing is left to seemore hidden patches that copying a neighbor can fill, and a weaker probe; at 0 nothing is left to predict
Decoder depth, Transformer blocks in the decoder; 8the probe stops gaining (73.3% at 12 against 73.5% at 8); fine-tuning stays near 84.9% (Table 1a)the probe falls to 65.5% at 1 block, the fewest that still pass visible tokens to mask tokens, while fine-tuning holds at 84.8% (Section 6)

From the MAE paper, sec. 3 and 4, Tables 8 to 10; main_pretrain.py in the official code.

Self-Supervised Learning Failure Modes and Checks comes next, on shortcuts through overlapping views and on detecting collapse.

8 References