Semi-Supervised Masked Autoencoders: Unlocking Vision Transformer Potential with Limited Data

TL;DR

This paper introduces SSMAE, a semi-supervised framework for training Vision Transformers effectively with limited labeled data by combining masked autoencoding and pseudo-labeling, reducing bias and improving performance.

Contribution

The paper presents a novel semi-supervised training method for ViTs that uses a validation-driven gating mechanism for pseudo-labeling, enhancing data efficiency and accuracy.

Findings

SSMAE outperforms supervised ViT and fine-tuned MAE on CIFAR datasets.

Significant gains in low-label regimes, e.g., +9.24% on CIFAR-10 with 10% labels.

Effective reduction of confirmation bias through dynamic pseudo-labeling activation.

Abstract

We address the challenge of training Vision Transformers (ViTs) when labeled data is scarce but unlabeled data is abundant. We propose Semi-Supervised Masked Autoencoder (SSMAE), a framework that jointly optimizes masked image reconstruction and classification using both unlabeled and labeled samples with dynamically selected pseudo-labels. SSMAE introduces a validation-driven gating mechanism that activates pseudo-labeling only after the model achieves reliable, high-confidence predictions that are consistent across both weakly and strongly augmented views of the same image, reducing confirmation bias. On CIFAR-10 and CIFAR-100, SSMAE consistently outperforms supervised ViT and fine-tuned MAE, with the largest gains in low-label regimes (+9.24% over ViT on CIFAR-10 with 10% labels). Our results demonstrate that when pseudo-labels are introduced is as important as how they are generated for data-efficient transformer training. Codes are available at https://github.com/atik666/ssmae.