Pseudo Labelling for Enhanced Masked Autoencoders

TL;DR

This paper introduces a novel masked autoencoder enhancement using pseudo labelling for class and data tokens, improving performance across multiple vision tasks by decoupling label generation and reconstruction with separate teacher models.

Contribution

It proposes a new approach that integrates pseudo labelling and token-level reconstruction in masked autoencoders, with a dual-teacher setup for improved performance.

Findings

Improved ImageNet-1K classification accuracy

Enhanced downstream task performance in segmentation and detection

Negligible impact on throughput and memory

Abstract

Masked Image Modeling (MIM)-based models, such as SdAE, CAE, GreenMIM, and MixAE, have explored different strategies to enhance the performance of Masked Autoencoders (MAE) by modifying prediction, loss functions, or incorporating additional architectural components. In this paper, we propose an enhanced approach that boosts MAE performance by integrating pseudo labelling for both class and data tokens, alongside replacing the traditional pixel-level reconstruction with token-level reconstruction. This strategy uses cluster assignments as pseudo labels to promote instance-level discrimination within the network, while token reconstruction requires generation of discrete tokens encapturing local context. The targets for pseudo labelling and reconstruction needs to be generated by a teacher network. To disentangle the generation of target pseudo labels and the reconstruction of the token features, we decouple the teacher into two distinct models, where one serves as a labelling teacher and the other as a reconstruction teacher. This separation proves empirically superior to a single teacher, while having negligible impact on throughput and memory consumption. Incorporating pseudo-labelling as an auxiliary task has demonstrated notable improvements in ImageNet-1K and other downstream tasks, including classification, semantic segmentation, and detection.