Masked Siamese Networks for Label-Efficient Learning

In plain words

Masked Siamese Networks teach a vision model to recognise a picture even when parts of it are hidden. The model never tries to draw the missing pixels. Instead it learns to give a partly hidden picture the same internal summary as the full one. The summary is measured against a set of learned reference points called prototypes. A second copy of the model, updated slowly, supplies the target summary. Training uses only unlabelled photos.

Because the hidden parts are dropped from the input, the masked version has fewer pieces to process, which keeps training cost down. The authors focus on the case where very few labelled examples are available: one to five per class. The representation they learn is meant to be adapted to that task with little labelled data.

TL;DR

MSN is a self-supervised pre-training method for Vision Transformers that does not reconstruct pixels. A masked view is trained to match the prototype assignments of an unmasked view, which comes from an EMA teacher. Unlike DINO, it needs no unmasked anchor view, so only the masked view's activations are kept for backpropagation. Author-reported: 72.4% top-1 with 5 labels per class and 75.7% with 1% of ImageNet labels, both for ViT-B/4. The paper and its code disagree on several details, listed under Paper vs code.

Problem

Self-supervised pre-training should give representations that work with few labels §1. Masked auto-encoders learn by reconstructing masked pixels or tokens. That objective needs low-level detail that classification does not need, so the features often have to be fine-tuned, and this can overfit in the low-shot regime §1. Joint-embedding methods avoid reconstruction and tend to give more semantic features, but they need views whose invariances are designed in advance and a mechanism to stop collapse §1, §2, §6. DINO-style methods also require an unmasked anchor view, which costs memory and compute §4. The paper's stated assumptions are a large unlabelled image set, a small labelled set that may overlap with it, and invariances defined by the chosen augmentations, which may depend on the dataset §2, §6.

Method

MSN builds on a siamese setup with a prototype-based objective §3. For each image, one target view is left unmasked. The anchor views are masked. The objective is to assign the anchor representation to the same prototypes as the target representation §3, Training Objective fig. 3.

1. Views. Each image yields one target view and M ≥ 1 anchor views. The default is one large (224×224) anchor with a random mask and ten small focal anchors (96×96) §3, Input Views §5 app. A.1.
2. Random masking. The masking drops a fraction of the global anchor's patches, so the encoder sees fewer tokens §3, Patchify and Mask app. A.1. The code keeps the CLS token and drops patches only in the first crop group deit.py.
3. Focal masking. A focal anchor is a small crop. It is not patch-dropped. The authors use both random and focal views together, and combining them gives the best low-shot accuracy §5.4, Combining Random and Focal Masking tab. 7.
4. Encoders. The anchor encoder is a ViT whose trainable weights are θ. The target encoder has the same architecture and its weights are an exponential moving average of the anchor's. The representation is the CLS token output §3, Encoder §2, Vision Transformer.
5. Anchor predictions. Each anchor representation is L2-normalised and compared by cosine similarity with K learnable prototypes (K = 1024, dimension d = 256). The softmax temperature is τ = 0.1 §3, Similarity Metric and Predictions app. A.1.
6. Target predictions. The same operation on the target view uses a sharper temperature τ⁺ = 0.025. Targets are computed without gradient. The paper requires τ⁺ < τ and states that training without sharpening can collapse §3, Similarity Metric and Predictions app. A.1.
7. Sinkhorn normalisation. The targets are normalised with Sinkhorn over the batch. This is on by default, and the paper says it removes the need to tune the ME-max weight app. A.1 app. C.1 tab. 11.
8. Loss. The loss is the average cross-entropy between each target distribution and each anchor distribution, minus λ times the entropy of the mean anchor prediction (ME-max). λ = 1.0 by default §3, Training Objective eq. 1 app. A.1.
9. Gradients and updates. Gradients flow only through the anchor predictions, never through the targets §3, Training Objective. The anchor encoder and prototypes are trained with AdamW. The target encoder is updated by EMA with momentum m, which the paper says rises from 0.996 to 1.0 over training §3, Encoder app. A.1.
10. Evaluation. The projection head is discarded at evaluation, and the target encoder's trunk representation is used app. A.1. See Pitfalls for what the code actually evaluates.

The paper gives a guarantee that the representations cannot collapse to a constant, under the assumption that the sharpened target is not uniform. It is stated in Proposition 1 app. B.

Prediction and objective equations, using the paper's notation. The prediction is not numbered in the paper. The objective is eq. 1:

p_{i,m} = \mathrm{softmax}\left(\frac{z_{i,m}\,\mathbf{q}}{\tau}\right), \qquad \bar{p} = \frac{1}{MB}\sum_{i=1}^{B}\sum_{m=1}^{M} p_{i,m}

Both are unnumbered in the paper; they are cited to §3, Similarity Metric and Predictions and §3, Training Objective.

\frac{1}{MB}\sum_{i=1}^{B}\sum_{m=1}^{M} H\left(p_i^{+},\, p_{i,m}\right) - \lambda\, H(\bar{p})

This is the only numbered equation, eq. 1 §3, Training Objective.

Pseudocode, written from the paper's description. The paper has no algorithm box. Where the official code differs, see Paper vs code.

# MSN pre-training step  [§3] [app. A.1]
# Input:  images (B, 3, 224, 224); prototypes (K, d), K = 1024, d = 256
# Output: scalar loss; updates to the anchor encoder and prototypes

def msn_step(images, anchor_encoder, target_encoder, prototypes):
    # 1. One unmasked target and M anchor views per image  [§3, Input Views]
    #    Default: one global anchor and ten focal anchors  [§5] [app. A.1]
    target_view = augment(images, size=224)           # (B, 3, 224, 224)
    global_anchor = augment(images, size=224)         # (B, 3, 224, 224)
    focal_anchors = [augment(images, size=96) for _ in range(10)]
    # 2. Random mask on the global anchor only; drop a fraction of patches
    #    (0.15 ViT-S/16, 0.3 ViT-B/16 and B/8, 0.7 ViT-L/7 and B/4)  [§3]
    masked_global = drop_patches(global_anchor, ratio=mask_ratio)
    # 3. Focal anchors are small crops, not patch-dropped  [§3] [app. A.1]
    # 4. Encode anchors with the trainable encoder; take CLS output  [§3]
    anchor_feats = anchor_encoder([masked_global] + focal_anchors)
    # 5. Encode the unmasked target with the EMA encoder, without gradient  [§3]
    with no_grad:
        target_feats = target_encoder(target_view)    # (B, d)
    # 6. Anchor predictions: cosine to prototypes, softmax at tau = 0.1
    #    [§3, Similarity Metric and Predictions] [app. A.1]
    anchor_preds = softmax(
        normalize(anchor_feats) @ normalize(prototypes).T / 0.1)  # (B*M, K)
    # 7. Target predictions: same similarity, sharper tau+ = 0.025, no gradient
    #    [§3] [app. A.1]
    with no_grad:
        target_preds = softmax(
            normalize(target_feats) @ normalize(prototypes).T / 0.025)
        # 8. Sinkhorn normalisation of the targets over the batch (default on)
        #    [app. A.1] [app. C.1]
        target_preds = sinkhorn(target_preds)                     # (B, K)
        # 9. Pair every anchor with its image's target; M = 11 by default  [§3]
        target_preds = tile(target_preds, times=M)                # (B*M, K)
    # 10. Cross-entropy between target and anchor distributions, averaged
    #     over anchors and images  [§3, Training Objective] [eq. 1]
    loss_ce = mean(-sum(target_preds * log(anchor_preds), dim=1))
    # 11. ME-max: entropy of the mean anchor prediction, weight lambda = 1.0
    #     (subtracted; rewards using all prototypes)  [§3] [app. A.1]
    mean_preds = mean(anchor_preds, dim=0)                        # (K,)
    entropy = -sum(mean_preds * log(mean_preds))
    loss = loss_ce - 1.0 * entropy
    # 12. AdamW update of anchor encoder and prototypes; no grad to targets
    #     [§3, Training Objective] [app. A.1]
    adamw_step(loss, anchor_encoder.parameters() + [prototypes])
    # 13. EMA update of the target encoder; momentum m rises from 0.996 to 1.0
    #     [§3, Encoder] [app. A.1]
    for target_param, anchor_param in zip(target_encoder, anchor_encoder):
        target_param.data = m * target_param.data + (1 - m) * anchor_param.data
    return loss

Hyperparameters

Pre-training, with the official config for ViT-S/16 msn_vits16.yaml. Values not in the paper say so.

Low-shot evaluation app. A.2:

Linear evaluation, 100% labels app. A.3:

Fine-tuning and transfer app. A.4, A.5. These protocols are not in this repository. The README points fine-tuning to the MAE codebase README.md.

Results

Author-reported. Setup: pre-training on ImageNet-1K; ViT backbones (ViT-S, B and L with patch sizes 16, 8, 7 and 4); baselines are DINO, iBOT, MAE, BYOL, SimCLRv2, PAWS and MoCov3 as listed in each table §5. Evaluation protocols are in Hyperparameters.

Headline numbers. ViT-B/4 reaches 72.4% top-1 with 5 labels per class tab. 1 and 75.7% with 1% of labels, about 13 per class tab. 2 §5.1. The paper says the 1% result exceeds an 800M-parameter convolutional network, while the best published 1% number it cites (76.6%) comes from a multi-stage semi-supervised pipeline, not a self-supervised one §5.1.

tab. 1 Extreme low-shot, ImageNet-1K pre-training. Top-1 %, mean ± std over 3 splits, for 1, 2 and 5 labelled images per class.

MAE uses partial fine-tuning except for the 1-image setting and ViT-H/14, which use a linear classifier §5.1 app. A.2.

tab. 2 Low-shot, 1% of ImageNet-1K labels (~13 per class). Top-1 %. † marks numbers the authors computed from public checkpoints.

The footnote in §5.4 says the ViT-S/16 row removes Sinkhorn normalisation. Its value, 67.2, matches the no-Sinkhorn, λ = 5 row of tab. 11 app. C.1.

tab. 3 Linear evaluation, 100% of ImageNet-1K labels. Top-1 %.

tab. 4 End-to-end fine-tuning, ViT-B/16, 100% labels. Top-1 %. The MSN row's pretraining is 600 epochs.

tab. 5 Transfer, fine-tuning from ImageNet-1K ViT-B/16. Top-1 %. Masking ratio 0.3 tab. 5 caption.

tab. 6 Transfer, linear evaluation, ViT-B/16. Top-1 %.

tab. 7 Masking strategy, ViT-B/16, 1% of labels. Top-1 %. The last row uses two anchors, one random and one focal.

tab. 8 Random masking ratio, 1% of labels. Top-1 %. Ratios are fractions of dropped patches.

tab. 9 Augmentation and view invariance, ViT-B/16, 1% of labels. Top-1 %.

The text gives 48.3% for the second row. See Pitfalls.

tab. 10 Compute, ViT-L/7 pre-training. Single p4d-24xlarge (8 A100 GPUs), batch 2 per GPU.

tab. 11 Sinkhorn normalisation, ViT-S/16, masking 0.15, 1% of labels. Top-1 %.

tab. 12 Number of prototypes, ViT-B/16, masking 0.3, 1% of labels. Top-1 %.

tab. 13 MAE low-shot protocols, ImageNet-1K. Top-1 %. The paper marks the ViT-H/14 partial fine-tuning row as "rand". It has no value.

tab. 14 Alternative ImageNet validation sets, fine-tuned ViT-B/16 app. D. Higher is better for A, R and Sketch (top-1 %). Lower is better for C (mean corruption error).

tab. 15 Robustness to missing patches at evaluation, 1% of labels. Top-1 %. app. E

tab. 16 Cosine similarity between representations of masked and unmasked views, ImageNet-1K validation. app. E

Figures are described in the paper, not reproduced here. fig. 1 has two panels: at 1% of labels, frozen-feature and fine-tuning results; and low-shot results comparing MSN ViT-L/7 with DINO ViT-B/8 and MAE ViT-L/16. fig. 2 shows samples from a generative model conditioned on MSN representations §1. fig. 3 is the method schematic and fig. 4 the masking strategies §3. fig. 5–fig. 9 compare MSN and DINO representations under masking app. F.

Paper vs code

Official code: github.com/facebookresearch/msn, at commit 4388dc1eadbe3042b85d3296d41b9b207656e043. It is linked in the abstract, and its owner is facebookresearch, the organisation the authors belong to. Code licence: CC BY-NC 4.0, non-commercial. The paper itself is CC BY 4.0, so its text and tables can be reused with attribution. The code cannot be reused commercially, so this card cites it by path and line only and quotes no code.

• Loss placement and evaluation feature. The paper says the loss applies to the encoder's CLS output and that the projection head is discarded at evaluation §3, Encoder app. A.1. In the code, the MLP head sits inside the ViT's feature function, so both anchor and target losses use the MLP output deit.py msn_train.py. The code comment calls h "before the head" msn_train.py, but on this path h is already after the MLP. Evaluation uses the CLS output of the last block, with no final LayerNorm and no head deit.py linear_eval.py logistic_eval.py.
• Projection head batch-norm. The paper says batch-norm at "the input and hidden layers" app. A.1. The code has linear → BN → GELU → linear → BN → GELU → linear, with no BN before the first linear. Hidden width is 2048 msn_train.py msn_vits16.yaml.
• Target temperature. The paper gives τ⁺ = 0.025 app. A.1. The code uses τ = 0.1 and raises the softmax to 1/T with T = 0.25 msn_vits16.yaml losses.py. Both give the same target distribution, because softmax(s/0.1)^4 normalises to softmax(s/0.025). The sharpening schedule starts and ends at 0.25, so it is constant msn_train.py.
• Schedules. The paper says the EMA momentum reaches 1.0 at the end of training app. A.1. The code increments it over 1.25 × the total steps and stops at 0.9992 msn_train.py. The learning-rate and weight-decay cosine schedules use the same 1.25 × horizon msn_train.py utils.py utils.py.
• Pre-training batch. The paper's 1024 is the total batch app. A.1. The config has 64 per process msn_vits16.yaml and the README runs 16 processes README.md. These agree only if each process is one GPU.
• Linear-evaluation batch. The paper gives 16,384 app. A.3. The code uses 256 per process linear_eval.py. The README runs 4 × 8 = 32 processes README.md, which gives 8,192 if each is one GPU. The two do not agree.
• Logistic regression λ. The paper gives no value app. A.2. The code default is 0.00025, divided by the number of embeddings logistic_eval.py. The README example uses 0.0025, ten times the default README.md.
• ViT-S/16 pre-training epochs. tab. 1 says 800; tab. 3 says 600 for the same model. The config uses 800 msn_vits16.yaml. The README's checkpoint is named vits16_800ep README.md.
• Repository coverage. The repository has one pre-training config, for ViT-S/16 . There are no configs for ViT-B/16, B/8, B/4 or L/7, though deit_small_p7 exists deit.py. Fine-tuning and transfer are not in this repository README.md. The robustness and diffusion visualisations are not in it either.
• Augmentation details. The code fixes crop scales, colour-jitter strength (0.5), probabilities and blur radius data_manager.py msn_vits16.yaml. The paper gives none of these app. A.1.
• Sinkhorn iterations. The paper does not state the count. The code uses 3 losses.py.
• Patch dropping. Only the first crop group is patch-dropped, which is the 224 global anchor. The focal crops are not deit.py deit.py. This matches app. A.1, but implementers who drop patches on every view will differ from the code.
• Other code-only details. Prototypes start uniform in ±1/√256 msn_train.py. Norm-layer and bias weights get no decay msn_train.py. Gradient norm is clipped at 3.0 msn_train.py. The final learning-rate floor is 1e-6 msn_vits16.yaml. SyncBatchNorm is used across processes msn_train.py. None of these are in the paper.
• Released checkpoints. The README links checkpoints for ViT-S/16 (800 epochs), ViT-B/16 (600), ViT-L/16 (600), ViT-B/4 (300) and ViT-L/7 (200) README.md. The repository states no licence for them beyond the code's.

No discrepancies were found in the target and anchor encoder roles, the 1024 prototypes of dimension 256, the anchor temperature 0.1, the ME-max weight 1.0, the AdamW optimiser, the warm-up of 15 epochs, the weight-decay range 0.04 → 0.4, or the one-random-plus-ten-focal anchor layout.

Pitfalls

• Which feature is evaluated. The paper says the trunk output is used app. A.1. The code evaluates the last block's CLS output without the final LayerNorm, and pre-trains the loss on the head's output. An implementation that uses the head output for evaluation, or applies the LayerNorm, will not match the reported numbers. The paper does not say which it used. Read deit.py and msn_train.py before choosing.
• Head batch-norm placement. BN before the first linear layer is in the paper's description app. A.1 and absent from the code msn_train.py. Hidden width 2048 is not stated in the paper.
• Sinkhorn is global. Sinkhorn runs over the batch on every process, with the global batch size, using all-reduce losses.py losses.py. The targets therefore depend on the GPU count. The ME-max average is also all-reduced losses.py.
• Sinkhorn is not always on. The footnote in §5.4 says the ViT-S/16 row of tab. 2 removes Sinkhorn. The 67.2 value matches the no-Sinkhorn, λ = 5 row of tab. 11. With Sinkhorn and λ = 1.0, the same model gives 66.4 app. C.1. The config has Sinkhorn on msn_vits16.yaml.
• Schedules are longer than the run. The momentum, learning-rate and weight-decay schedules are all built over 1.25 × the epochs msn_train.py. Momentum ends at 0.9992, not 1.0. The learning rate is still about 0.1 × peak at the last step for the 800-epoch config. This is my arithmetic from the schedule formula utils.py.
• Entropy term is inert. use_ent is true, but ent_weight is 0.0 msn_vits16.yaml msn_train.py. Do not add an entropy term with non-zero weight unless the paper's objective calls for it. The paper's objective has none eq. 1.
• Target encoder mode and buffers. The EMA update covers parameters only msn_train.py. The target's batch-norm running statistics are not averaged. msn_train.py never calls .eval(), so the target encoder runs in training mode with batch statistics. This is the code's behaviour, not a stated design.
• Collapse. app. B proves the no-collapse property only when the sharpened target is not uniform. The paper says that without sharpening it observed collapse §3, Similarity Metric and Predictions. Keep the sharpening even when the batch looks healthy.
• Text and table disagree on the augmentation ablation. The text says 48.3% for colour jitter alone §5.4. tab. 9 says 48.7.
• Epoch counts disagree. ViT-S/16 is 800 epochs in tab. 1 and 600 in tab. 3. Use the config's 800 unless you are reproducing the linear-evaluation row.
• Masking ratio and stability. The paper's ratios are model-specific: 0.15 for ViT-S/16, 0.3 for ViT-B/16 and ViT-B/8, and 0.7 for ViT-L/7 and ViT-B/4 app. A.1. The ablation shows ViT-L/16 at 0.15 and 0.3 gives NaN tab. 8, and the text says those weak-masking runs are unstable §5.4. Use the model-specific ratio.
• Compute claim. The paper says masking 70% "uses half the computation and memory" §1. tab. 10 shows memory falling from 26 GB to 17 GB and throughput rising from 415 to 600 img/s. That is not a halving, and the table does not measure compute directly.
• Scale claims. Pre-training ViT-L/7 needs 18 p4d.24xlarge machines with 70% masking, against more than 42 without masking §1 §5.4. These are the authors' figures.
• Baseline numbers are re-evaluations. The DINO and iBOT ViT-B rows in tab. 2 are marked † because the authors computed them from public checkpoints tab. 2 caption.
• Citation slip. In the source, the iBOT row of tab. 4 cites the MAE reference. Check the iBOT reference before citing that row.
• Placeholder. The ViT-H/14 partial-fine-tuning row in tab. 13 reads "rand" and has no value.
• Unstated values. Projection-head width, drop path, final learning rate, gradient clipping, weight-decay exclusions, prototype initialisation, Sinkhorn iterations and all augmentation strengths are not in the paper. Taking them from the code is the only way to match the reported setup, and this card shows them with their source.
• Licence. The code is non-commercial LICENSE. The paper is CC BY 4.0. Implement from the paper and this card, and do not copy code from the repository into commercial work.

Lineage

• Builds on: DINO (Caron et al., ICCV 2021, "Emerging properties in self-supervised vision transformers"). MSN extends it with random and focal masking and drops the requirement for an unmasked anchor view §4. There is no arXiv id in the paper's bibliography, so none is given here. The clustering-based SSL line (SwAV, PAWS) gives the prototype-and-Sinkhorn machinery §3 app. A.1. The EMA target follows BYOL §3, Encoder app. A.1.
• Compared against: iBOT, MAE, DINO, BYOL, SimCLRv2, PAWS, MoCov3, BEiT, SimMIM, MaskFeat, Data2Vec (arXiv:2202.03555) tab. 1–4.