Source-linked AI summary

Masked Siamese Networks for Label-Efficient Learning

Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Michael Rabbat, Nicolas Ballas

arXiv:2204.07141v1cs.LGcs.AIcs.CVeess.IV

TL;DR

Self-supervised vision methods seek representations that support prediction with few labeled examples, but reconstruction-based and joint-embedding approaches have different trade-offs. MSN matches a randomly masked view to an unmasked view at the representation level using a Vision Transformer. It achieves strong low-shot results while improving the computational scalability of joint-embedding architectures.

  • Problem

    Self-supervised learning aims to learn image representations from unlabeled data that transfer effectively to prediction tasks with few labels.

  • Method

    MSN randomly masks patches in one augmented image view and trains its representation to match the unmasked target view without pixel- or token-level reconstruction.

  • Results

    75.7% top-1 accuracy is achieved with 1% of ImageNet-1K labels, while masking 70% of patches uses half the computation and memory of an unmasked joint-embedding baseline.

  • Takeaways & Limitations

    MSNs produce strong off-the-shelf representations for label-efficient learning and improve the scalability of joint-embedding architectures.

Abstract

from arXiv · show

We propose Masked Siamese Networks (MSN), a self-supervised learning framework for learning image representations. Our approach matches the representation of an image view containing randomly masked patches to the representation of the original unmasked image. This self-supervised pre-training strategy is particularly scalable when applied to Vision Transformers since only the unmasked patches are processed by the network. As a result, MSNs improve the scalability of joint-embedding architectures, while producing representations of a high semantic level that perform competitively on low-shot image classification. For instance, on ImageNet-1K, with only 5,000 annotated images, our base MSN model achieves 72.4% top-1 accuracy, and with 1% of ImageNet-1K labels, we achieve 75.7% top-1 accuracy, setting a new state-of-the-art for self-supervised learning on this benchmark. Our code is publicly available.

1 Introduction

Masked Siamese Networks learn high-semantic image representations by matching masked and unmasked views, avoiding pixel- or token-level reconstruction. They achieve strong label-efficient performance while improving joint-embedding scalability.

  • Masked Siamese Networks: MSN matches representations of randomly masked and unmasked image views, performing denoising implicitly at the representation level rather than reconstructing pixels or tokens.The method uses a ViT encoder and trains the masked view to match the unmasked view.
  • Low-shot evaluation: 75.7% top-1 accuracy is achieved on the standard 1% ImageNet low-shot classification task, outperforming the previous 800M-parameter state of the art with nearly 10× fewer parameters.The result uses an MSN-trained ViT-B/4.
  • Low-shot evaluation: 72.1% top-1 accuracy is achieved with only 5 labeled images per class, surpassing DINO by 8% top-1 in the 1–5-images-per-class benchmark.The model is a ViT-L/7 pre-trained with MSN on ImageNet-1K.
  • Scalability: Masking 70% of patches uses half the computation and memory of an unmasked joint-embedding baseline because only unmasked patches are processed by the ViT encoder.A ViT-L/7 can be pre-trained on 18 AWS p4d-24xlarge machines, compared with over 42 without masking.
  • Representation quality: MSN representations retain semantic information about objects despite 70% patch masking while discarding background, pose, and lighting information.The visualization uses generative samples conditioned on representations from a masked ViT-L/7 encoder.
  • Broader evaluation: MSNs are competitive with prior self-supervised methods on benchmarks using more labels, including fine-tuning, linear evaluation, and transfer learning.This extends the reported evaluation beyond low-shot classification.

2 Prerequisites

Self-supervised Siamese networks learn image embeddings that are invariant across views, while Vision Transformers encode images as sequences of patch tokens processed by Transformer layers. Preventing representation collapse is a central challenge for these architectures.

  • Siamese Networks: Siamese networks independently encode two views of an image and train their representations to match despite view differences.The encoder produces anchor and target representations for the two views.
  • Siamese Networks: Representation collapse occurs when a Siamese encoder produces a constant embedding regardless of the input.Existing remedies include contrastive, information-maximization, and asymmetric approaches.
  • Masked Siamese Networks: MSN's schematic applies random augmentations to create anchor and target views, masks the anchor, leaves the target unchanged, and aligns their cluster assignments.The objective uses a standard cross-entropy loss.
  • Vision Transformer: Vision Transformers split images into non-overlapping N × N patches, convert them into patch tokens, add positional embeddings and a [CLS] token, then apply Transformer layers.The [CLS] token aggregates information from the full patch sequence.

3 Masked Siamese Networks

Masked Siamese Networks combine mask denoising with Siamese representation learning: a masked anchor view is matched to an unchanged target view through encoder outputs and prototype-based predictions. The method uses ViTs, masked-patch processing, sharpening, and entropy regularization to support efficient training without reconstructing pixels or tokens.

  • Input Views: MSN generates augmented anchor and target views, masks patches only in the anchor, and trains their representations to agree.The target view remains unchanged while the anchor is patchified and randomly masked.
  • Mask Denoising: MSN applies discriminative mask denoising at the ViT [CLS] output rather than reconstructing masked pixels or tokens.The loss is applied directly to the encoder’s [CLS] representation.
  • Masking Strategies: Random masking drops potentially non-contiguous patches, whereas focal masking preserves a local block and drops patches around it.The framework typically leverages both masking strategies during pre-training.
  • Encoder: The anchor and target encoders produce representations from masked and unmasked views, with the target encoder updated by exponential moving average.Both encoders use ViT trunks, and the [CLS] token supplies the representation.
  • Similarity Metric and Predictions: MSN compares prototype-based prediction distributions for anchor and target representations using cosine similarity and cross-entropy.The predictions use learnable prototypes and separate temperature parameters for anchor and target views.
  • Regularization: Target sharpening and ME-MAX regularization encourage confident, non-collapsed predictions while distributing usage across the prototype set.Training without sharpening has been observed to result in collapsing solutions.

4 Related Work

MSN builds on joint-embedding and masked-image-learning approaches while differing in its global representation objective and masking procedure. Compared with related methods, it focuses on label efficiency and combines random and focal masking without requiring an unmasked anchor view.

  • Joint-Embedding Methods: MSN is related to DINO through Siamese networks, cross-entropy, and a momentum encoder, but generalizes it with random and focal masking.Unlike DINO, MSN does not require an unmasked anchor view.
  • Masked Image Modeling: Earlier masked-image methods commonly remove input content and reconstruct pixels, tokens, channels, or image regions.These approaches frame mask-based learning as prediction or regression of removed content.
  • Joint-Embedding Denoising: Other joint-embedding denoising methods match patch-level outputs to target tokens and may add a global-sequence loss.MSN instead emphasizes reducing labeled data requirements for downstream prediction.
  • Evaluation: Table 1 evaluates label efficiency on ImageNet-1K using extremely small labeled sets and reports mean top-1 accuracy with standard deviation across three splits.This evaluation targets extreme low-shot representation learning.

5 Results

MSN performs strongly in extreme and standard low-shot evaluation, while remaining competitive with prior methods under fuller supervision and transfer learning. Ablations show that masking and view construction affect both performance and computational scaling.

  • Extreme Low-Shot: MSN outperforms other representation-learning approaches across all extreme low-shot supervision levels, with larger gains as labels decrease.Larger models, deeper networks, and smaller patch sizes appear especially beneficial when fewer labels are available.
  • Extreme Low-Shot: Joint-embedding approaches appear more robust than reconstruction-based auto-encoding approaches when downstream supervision is scarce.The text relates this observation to the semantic level of pixel-reconstruction representations and the usefulness of invariance-based pre-training.
  • 1% ImageNet-1K: 75.7% top-1 accuracy on 1% ImageNet-1K labels surpasses the previous 800M-parameter convolutional state of the art with nearly 10× fewer parameters and no fine-tuning.The comparison focuses on self-supervised models; the best literature result overall was 76.6% from a multi-stage semi-supervised pipeline.
  • Standard Evaluation: MSN achieves competitive performance in linear evaluation, full-label fine-tuning, and transfer learning.The best MSN linear-evaluation result is 80.7% top-1 accuracy, while transfer results are comparable to or better than DINO across reported tasks.
  • Ablations: Random masking improves over no masking, focal masking alone degrades performance, and combining random and focal masks yields the strongest results.The ablation highlights the importance of retaining a global view during pre-training.
  • Ablations: Increasing the masking ratio helps larger models improve low-shot performance, although evidence is insufficient to claim universally improved training stability.Weakly masked ViT-L/16 runs were unstable, whereas more aggressively masked runs were stable.
  • Ablations: Sharing image views can cause a shortcut based on color statistics, while independent color and geometric transformations improve 1% ImageNet accuracy from 7% to 52.3%.Using different colors reaches 48.3%, and adding independently applied geometric transformations reaches 52.3%.

6 Conclusion

The paper presents MSN as a self-supervised framework that combines mask-denoising with representation matching and reports label-efficient representations alongside improved joint-embedding scalability. Its view-invariant design requires specifying data transformations, whose optimal choices may depend on the dataset and task.

  • 6 Conclusion: MSN uses mask-denoising without pixel- or token-level reconstruction to learn image representations.The framework matches representations across differently viewed images, including a randomly masked view and an unchanged view.
  • 6 Conclusion: MSN learns strong off-the-shelf representations that excel at label-efficient learning while improving the scalability of joint-embedding architectures.The conclusion states both representation quality and scalability as empirical outcomes.
  • 6 Conclusion: MSN requires specified data transformations, and the optimal transformations and invariances may depend on the dataset and task.Future work proposes more flexible transformation learning and equivariant representations.

A.1 MSN Pre-training

The appendix describes optimization, evaluation, and transfer-learning procedures used for MSN and comparison models. These protocols combine standard Vision Transformer training schedules with frozen-feature linear evaluation or task-specific fine-tuning.

  • A.1 MSN Pre-training: MSN pre-training uses AdamW with batch size 1024, linear learning-rate warmup, and cosine decay.The learning rate warms from 0.0002 to 0.001 over the first 15 epochs.
  • A.1 MSN Pre-training: Low-shot evaluation freezes the pre-trained model and trains an L2-regularized logistic-regression classifier using 1, 2, or 5 labeled samples per class.Each labeled image is center-cropped before representation extraction.
  • A.1 MSN Pre-training: Partial fine-tuning updates the last pre-trained block and a linear head, but large models such as ViT-H/14 can overfit in the low-shot regime.Reported MAE results use partial fine-tuning except in the 1-image-per-class setting and for ViT-H/14, which use linear classification.
  • A.1 MSN Pre-training: Linear evaluation trains a classifier on frozen, L2-normalized representations using a batch size of 16,384 for 100 epochs.The protocol uses basic random-resized-crop and horizontal-flip augmentations.
  • A.1 MSN Pre-training: Full-label ImageNet fine-tuning updates the pretrained ViT for 100 epochs with AdamW, warmup, and cosine learning-rate decay.The setup uses batch size 1024 and learning rate 0.002.
  • A.1 MSN Pre-training: Transfer linear evaluation freezes ImageNet-pretrained weights and fits an L2-regularized logistic-regression classifier on center-cropped downstream images.Images are resized to 256 × 256 pixels before taking 224 × 224 center crops.
  • A.1 MSN Pre-training: Transfer fine-tuning uses 1000 epochs for CIFAR10 and CIFAR100, and 300 epochs for iNat18 and iNat19.The experiments use task-specific batch sizes, learning rates, and RandAugment-based augmentations.

B Theoretical Guarantees

MSN pre-training is theoretically guaranteed to avoid trivial representation collapse under a target-sharpening assumption. The guarantee relies on entropy maximization and target sharpening preventing distinct collapse modes.

  • MSN’s non-collapsing guarantee assumes that target predictions are sharpened away from the uniform distribution.The assumption is stated for mini-batches containing B images and M anchor views per image.
  • If all representations collapse, identical prototype similarities force identical predictions across images and views.The proof uses L2-normalized representations and prototypes with K > 1 learnable prototypes.
  • Uniform collapsed predictions conflict with non-uniform sharpened targets, producing a nonzero gradient that prevents collapse.The proof separates the case of uniform predictions from the case of non-uniform predictions.
  • Non-uniform collapsed predictions are opposed by entropy maximization, which encourages use of the full prototype set.Together with target sharpening, this prevents collapse to a non-uniform distribution.
  • The sharpening mechanism may not satisfy the theoretical assumption without a tiebreaking rule, although practice avoids uniform targets when sharpening starts immediately.A sufficiently small target temperature is important in this practical setting.
  • Omitting Sinkhorn normalization while tuning the ME-MAX weight can improve low-shot performance.The default ME-MAX regularization weight is λ = 1.0, with Sinkhorn normalization applied to targets.

C.2 Number of Prototypes

The section examines how prototype count affects low-shot performance and contrasts alternative evaluation strategies for MAE. Too few prototypes can hurt MSN performance, while larger MAE models benefit from partial fine-tuning.

  • C.2 Number of Prototypes: Using more prototypes has little effect on MSN training, but using too few prototypes can degrade low-shot performance.The experiment uses a ViT-B/16 with a masking ratio of 0.3 and 1% of ImageNet-1K labels.
  • C.3 Masked Auto-Encoder Partial Fine-Tuning: Partial fine-tuning adapts MAE’s last block and a linear head using the available labeled samples.This is contrasted with linear evaluation on frozen pre-trained features.
  • C.3 Masked Auto-Encoder Partial Fine-Tuning: MAE benefits from partial fine-tuning, particularly for sufficiently large models.The supplied passage specifically identifies the ViT-H/14 as an example of a sufficiently large model.

D MSN Representation Robustness

MSN-pretrained ViT-B/16 models are evaluated on four alternative ImageNet validation sets against MAE-pretrained ViT-B/16 and a supervised ResNet50 baseline. MSN performance is comparable to or better than MAE across all four datasets.

  • ImageNet-A, ImageNet-R, and ImageNet-Sketch use top-1 accuracy, while ImageNet-C uses mean Corruption Error.The evaluation also reports a supervised ResNet50 baseline.
  • MSN-pretrained ViT-B/16 performance is comparable to or better than MAE-pretrained ViT-B/16 on all four robustness datasets.The comparison covers ImageNet-A, ImageNet-R, ImageNet-Sketch, and ImageNet-C.

E MSN Invariance to Masking

MSN is designed to denoise images at the representation level by matching masked and unmasked views. Experiments show greater robustness to missing patches, especially after more aggressive masking during pre-training.

  • MSN matches representations of masked and unmasked image views to denoise inputs at the representation level.The method applies a random mask to the anchor view while leaving the target view unchanged.
  • The masking robustness evaluation trains a linear classifier on masked images and tests it on unmasked ImageNet-1K validation images.Models are evaluated using 1% of ImageNet-1K labels.
  • Masked pre-training produces representations that are more robust to patch removal, with stronger robustness after more aggressive masking.This pattern is interpreted as representation-level image denoising.
  • 75.1% top-1 accuracy without dropping patches versus almost unchanged low-shot accuracy after removing 70% of patches for aggressively masked ViT-L/7.
  • MSN pre-training yields higher cosine similarity between masked and unmasked representations of the same image.This supports the observed robustness to patch removal.

F Qualitative Analysis

The qualitative analysis uses conditional diffusion sampling to inspect what MSN and DINO representations retain or discard, including under substantial patch masking. MSN representations preserve semantic information and image structure despite missing patches, while discarding instance-specific details such as background, pose, and lighting.

  • RCDM visualization: RCDM maps SSL representation vectors and random noise back to pixel space, enabling visual inspection of information retained by pre-trained representations.Information that varies across generated samples is absent from the representation, while semantically common qualities are retained.
  • Comparison with DINO: At 50% masking, RCDM samples conditioned on MSN representations preserve the masked image’s semantic category and remain visually close to the original.DINO-conditioned samples are blurrier and preserve the semantic category less effectively.
  • Comparison with DINO: At 80% masking, MSN-conditioned samples preserve some structure from the original images, whereas DINO conditioning produces almost uniform backgrounds.The comparison evaluates representations computed from heavily masked validation images.
  • Information retained by MSN: MSN representations discard instance-specific information such as background, pose, and lighting while retaining semantic image information under large masking ratios.The visualizations include unmasked images and images with 70% or 90% of patches masked.
Loading 2204.07141v1…