SimMIM: A Simple Framework for Masked Image Modeling

In plain words

Language models such as BERT learn by filling in blanked-out words. SimMIM asks whether the same trick works for images: hide parts of a photo, ask a network to paint the missing parts back in, and hope that doing so teaches it what objects look like. No labels are needed, so the network can learn from plain images before it is fine-tuned for a real task such as classification.

Earlier attempts added extra machinery: a separate network that turns image pieces into "visual words", special block-shaped masks, or colour palettes. SimMIM drops all of it. It cuts the image into a grid of large squares, hides a random subset of them, and asks the network to output the original colour values of the hidden pixels. The part that turns the network's features back into pixels is a single linear layer, and the score is just how far the predicted colours are from the true ones, measured only on the hidden squares.

The paper's main lesson is that the masks must be big enough that the network cannot cheat by copying neighbouring pixels. Hiding small squares, or rebuilding the visible parts as well, leads to worse features. A heavier decoder paints nicer pictures but does not give a better network for the downstream task. This plain recipe matches or beats the more complicated methods it is compared against and was used to help train a 3-billion-parameter model.

TL;DR

SimMIM is masked image modeling stripped to four simple choices: patch-aligned random masking with a large masked patch (32×32, ratio 0.6), a learnable mask token replacing masked patch embeddings, a linear (1×1 conv) head that predicts all raw RGB pixels of each output feature's receptive patch, and an \ell_1 loss computed only on masked pixels §3 fig. 1. Ablations on Swin-B show heavier heads, classification targets (iGPT clusters, BEiT dVAE tokens, colour bins) and full-image reconstruction do no better or worse tab. 2 tab. 4 tab. 5. Author-reported ImageNet-1K fine-tuning: 83.8% top-1 for ViT-B (BEiT 83.2) tab. 6 and 87.1% for SwinV2-H at 512² tab. 7. Relevant if you need a cheap, tokenizer-free MIM pre-training recipe that works for hierarchical backbones (Swin) as well as ViT; it is not for linear-probe use (56.7% linear eval on ViT-B) tab. 6.

Problem

Masked signal prediction drives NLP pre-training, but in vision it was overshadowed by contrastive learning §1. The authors list three modality gaps: strong pixel locality lets a model fill holes by copying neighbours instead of reasoning; pixels are low-level signals unlike human-made word tokens; and pixels are continuous whereas MLM classifies discrete tokens §1. Prior MIM methods bridged these with special designs (iGPT colour clustering, BEiT's dVAE tokenizer plus block-wise masking) §1 §2. The paper asks which components actually matter and whether simple choices suffice §3.1.

Method

SimMIM has four components: masking strategy, encoder, prediction head, prediction target §3.1. The encoder is a standard ViT or Swin Transformer §3.1.

Masking §3.2

1. Divide the input into a grid of masked patches of size s \times s (default s=32, which equals the last-stage patch size of Swin; also 32 for ViT) §3.2 §4.1.1. Each masked patch is either fully visible or fully masked §3.2.
2. Choose a random subset of masked patches with ratio r (default 0.6) §4.1.1. Code picks exactly \lceil r \cdot N \rceil of the N=(H/s)^2 cells via a random permutation, then upsamples the mask to the model's patch-embedding grid (4×4 tokens for Swin, 16×16 for ViT) data_simmim.py.
3. After patch embedding, replace every masked token with a single learnable mask-token vector of the embedding dimension §3.2. Code: x = x*(1-w) + mask_token*w, done before absolute position embeddings (Swin) or before prepending the class token and adding position embeddings (ViT); mask token is truncated-normal initialised with std 0.02 simmim.py simmim.py.

Encoder and head §3.3 §3.4

4. Run the encoder on the full token sequence (masked tokens included) and take the final feature map, down-sampled 32× (Swin) or 16× (ViT) §3.4 simmim.py.
5. Prediction head: a 1×1 convolution (a linear layer per feature vector) with output dimension \text{stride}^2 \times 3, e.g. 3072 = 32 \times 32 \times 3 for Swin, so each feature vector predicts the RGB values of its own 32\times32 pixel block §3.4. Code reshapes with PixelShuffle to a full-resolution image simmim.py.

Loss §3.4

6. \ell_1 loss averaged over masked pixels only:

L = \frac{1}{\Omega(\mathbf{x}_M)}\|\mathbf{y}_M - \mathbf{x}_M\|_1 eq. 1

where \mathbf{x},\mathbf{y} \in \mathbb{R}^{3HW\times 1} are input and predicted RGB values, M the set of masked pixels and \Omega(\cdot) the number of elements §3.4. Code: element-wise L1, multiplied by the pixel-level mask, summed, divided by (mask pixel count + 1e-5) and by the 3 channels simmim.py.

Algorithm: SimMIM pre-training step  [§3] [code: models/simmim.py#L114-L121]
Input: batch of images x (B,3,H,W), normalised; mask m (B, H/p, W/p) at model-patch grid p
1. tokens = patch_embed(x)                          # (B, L, C)
2. tokens = tokens*(1-m) + mask_token*m             # per-token replacement
3. z = encoder(tokens)  -> feature map (B, C', H/S, W/S), S = 32 (Swin) or 16 (ViT)
4. y = PixelShuffle_S( Conv1x1(C' -> S*S*3)(z) )    # (B,3,H,W)
5. M = upsample m by p to pixel resolution
6. loss = sum(|x - y| * M) / (sum(M) + 1e-5) / 3
Output: loss

Alternatives studied §3.2 §3.3 §3.4

• Masking: square masking (a centrally placed region from Context Encoder, relaxed to random position) and BEiT block-wise masking at masked-patch sizes 16 and 32 §3.2.
• Heads: 2-layer MLP, inverse Swin-T, inverse Swin-B §3.3.
• Lower target resolutions: downsample the original image by 32×, 16×, 8×, 4×, 2× §3.4.
• Classification targets: iGPT's 512 k-means colour clusters (the cluster centres learned in iGPT), BEiT dVAE tokens, and channel-wise equal-width colour bins (e.g. 8 or 256 bins per channel) §3.4.
• Losses: \ell_2 and smooth-\ell_1 §3.4.
• Reconstruction variant: also regress visible pixels (loss on the whole image) §4.1.5 tab. 4.

None of these alternatives is in the released code; it implements only the default linear head and masked \ell_1 loss simmim.py.

AvgDist diagnostic §4.1.2 fig. 3: the mean Euclidean distance from each masked pixel to the nearest visible pixel. Fine-tuning accuracy follows a ridge over AvgDist, with the best entries roughly in [10, 20] §4.1.2 fig. 3; the appendix calls an AvgDist of 15 empirically good app. C.2 fig. 7. The exact computation procedure (units, sampling) is not stated beyond the definition.

ConvNet variant app. E: for ResNet-50×4, the mask token replaces features after a stem of a 3×3 stride-2 convolution followed by 2×2 max-pooling.

Hyperparameters

Ablation default (Swin-B, ImageNet-1K) §4.1.1. Code LRs are specified per 512 images and scaled linearly by total batch / 512 main_simmim.py main_finetune.py.

ViT-B system comparison §4.2:

Swin scaling §4.3 tab. 8:

Downstream (appendix) app. C.1:

Results

Author-reported. Ablations: Swin-B, 192², 100-epoch pre-training and fine-tuning on ImageNet-1K, top-1 §4.1.1.

Masking strategy tab. 1:

Prediction head tab. 2 tab. 10:

Heavier heads reach slightly lower pre-training loss (0.3722 inverse Swin-B vs 0.3743 linear) yet transfer worse §4.1.3.

Prediction resolution tab. 3:

Prediction target tab. 5 tab. 11:

Channel-wise bin classification, Swin-B, 100 epochs tab. 13:

Loss scope tab. 4: masked area only 82.8 vs full image 81.7.

LR scheduler (Swin-B, both with 10-epoch warm-up; step: 8e-4, ÷10 at 90% and 95%) app. B tab. 9: cosine 82.8 (100 ep) / 83.0 (300 ep); step 82.9 / 83.3.

ViT-B system comparison, ImageNet-1K tab. 6:

Swin scaling, ImageNet-1K, pre-trained at 192² tab. 7:

Swin scaling on COCO and ADE20K tab. 12:

SwinV2-G (3B), trained with SimMIM's help on data the authors describe as ~40× smaller than JFT-3B: 84.0% ImageNet-V2 top-1, 63.1/54.4 COCO box/mask mAP, 59.9 ADE20K mIoU, 86.8% Kinetics-400 top-1; details deferred to the SwinV2 paper §4.3.

ResNet-50×4: 81.6% ImageNet-1K top-1 after 300-epoch pre-training and 100-epoch fine-tuning, vs 80.7% supervised app. E.

Paper vs code

Checked at d3e29bc. The README notes SimMIM was later merged into the Swin Transformer repository README.md; this card uses the standalone repo.

• ViT-B pre-training schedule. Paper: cosine schedule, 20-epoch linear warm-up, 800 epochs §4.2. Code: multistep schedule (×0.1 at epoch 700), 10 warm-up epochs, 1e-4 per 512 (4e-4 at 2048) simmim_pretrain__vit_base__img224__800ep.yaml. This is the Swin scaling recipe of §4.3, not the one §4.2 describes.
• Fine-tuning warm-up. Paper: 10 epochs §4.1.1, also used for ViT-B and Swin scaling ("other settings as ablation") §4.2 §4.3. Code: 20 epochs in every fine-tune config simmim_finetune__swin_base__img224_window7__100ep.yaml simmim_finetune__vit_base__img224__800ep.yaml.
• Swin-L fine-tuning. Paper: layer decay 0.75 and drop path from the ablation (0.1) §4.3 §4.1.1. Code: layer decay 0.7 and drop path 0.2 simmim_finetune__swin_large__img224_window14__800ep.yaml simmim_finetune__swin_large__img224_window14__800ep.yaml.
• Prediction target. Paper says raw RGB pixel values §3.4. Code regresses the ImageNet mean/std-normalised input tensor, so targets are per-channel standardised values, not [0, 255] or [0, 1] pixels data_simmim.py simmim.py.
• ViT details not in the paper. Pre-training uses a shared relative position bias and no absolute position embedding, layer-scale init 0.1 and drop path 0.1; fine-tuning switches to per-layer relative position bias (copied from the shared table) and mean pooling simmim_pretrain__vit_base__img224__800ep.yaml simmim_finetune__vit_base__img224__800ep.yaml config.py utils.py. The paper only says "vanilla ViT" §3.1.
• Swin window change. Pre-training at window 6 and fine-tuning at window 7 needs resizing of the relative position bias tables; code interpolates them utils.py. The paper does not describe this.
• Unstated training details in code: gradient clipping at norm 5.0 config.py, apex AMP O1 by default main_simmim.py, 1e-5 in the loss denominator simmim.py.
• Matches: ablation pre-train LR 8e-4 and fine-tune LR 5e-3 at batch 2048 after linear scaling; Swin 800-epoch step schedule at 7/8; Swin-B layer decay 0.8 and ViT-B 0.65; pre-train augmentation; loss normalisation equals eq. 1.

Pitfalls

• Learning rates are batch-scaled in code. Config BASE_LR is per 512 images and multiplied by total batch / 512 (and by gradient-accumulation steps) main_simmim.py. Copying 8e-4 into a config at batch 2048 gives 4× the intended LR.
• "Masked patch" ≠ model patch. Masks are drawn on a 32×32 grid and then expanded to the 4×4 (Swin) or 16×16 (ViT) token grid; the image size must divide by 32 and 32 by the model patch size data_simmim.py. Using a 16-pixel mask with ViT/16 is the small-patch setting that scored worse tab. 1.
• Mask ratio is rounded up. \lceil 0.6 \times 36 \rceil = 22 cells at 192² and \lceil 0.6 \times 49 \rceil = 30 at 224², so the actual ratio is about 0.61 data_simmim.py.
• Encoder sees the full sequence. Masked positions stay in the sequence as mask tokens §3.2 simmim.py, so every pre-training step runs the encoder over all tokens of the image. The relative cost figures (1.0× vs BEiT 1.5×, DINO 2.0×) are attributed by the authors to the method's simplicity §4.2 tab. 6, not to skipping tokens.
• Loss only on masked pixels. Including visible pixels drops fine-tuning from 82.8 to 81.7 tab. 4; visualisations of the visible area show checkerboard artifacts because it is never trained §4.4 fig. 4. Don't judge the model by full-image reconstructions.
• Lower inpainting loss ≠ better features. Heavier heads lower the loss but transfer worse §4.1.3; select checkpoints by downstream fine-tuning, not pre-training loss.
• Linear probing is weak. 56.7% on ViT-B tab. 6; SimMIM features are meant for fine-tuning, and the authors explicitly do not optimise for linear evaluation §3.5.
• Table 1 caveat. The 0.5-ratio best (83.0) is not the default; 0.6 was chosen for stability §4.1.2. Differences across most settings are a few tenths of a point, and no variance over seeds is reported.
• Classification targets at 192² are N/A for 256 bins and iGPT clusters tab. 5; the reason is not stated.
• Table 6 linear eval values. SimMIM and BEiT are both listed at 56.7% tab. 6; the paper does not comment on the coincidence.
• ADE20K numbers start from the supervised fine-tuned checkpoint, not the raw self-supervised weights, because that worked better app. C.1. Reproductions that initialise from pre-trained weights directly should expect lower numbers.
• SwinV2-G results depend on a private dataset (ImageNet-22K-ext) and the SwinV2 paper's recipe §4.3; they cannot be reproduced from this paper alone.
• Scaling recipe uses a step LR schedule, chosen so runs of different lengths can share the first stage §4.3; step was marginally better than cosine in the appendix tab. 9, so the cosine ablation numbers and step-scheduled scaling numbers are not perfectly comparable.

Lineage

• Builds on: BEiT (arXiv:2106.08254), keeping the learnable mask token but dropping the dVAE tokenizer and block-wise masking §3.2 §3.4; Context Encoder (pixel inpainting with square masks) §2; iGPT colour clustering as an alternative target §3.4. Encoders: ViT and Swin Transformer / SwinV2 §3.1 §4.3.
• Compared against: supervised DeiT/Swin baselines, DINO, MoCo v3, ViT's masked patch prediction, BEiT tab. 6 tab. 7.
• Superseded by: not known from content in this wiki.