Towards Latent Masked Image Modeling for Self-Supervised Visual Representation Learning

TL;DR

This work tackles the difficulty of extracting high-level semantics with pixel-based masked image modeling (MIM) by introducing Latent MIM, which learns localized latent representations and reconstructs masked regions in latent space. It identifies four core challenges—latent target optimization, reconstruction objectives, semantic correlation between nearby patches, and decoder design—and resolves them through asymmetric targets, a patch-discrimination objective, masking strategies, and a cross-attention decoder with visual cues and a latent projector. Scaling to ImageNet-1k demonstrates substantial gains in nearest-neighbor and linear probing, as well as strong performance in unsupervised segmentation, video object segmentation, and few-shot transfer, indicating robust, locally semantic representations without supervision. The results suggest Latent MIM can outperform traditional pixel-level MIM and prior latent MIM approaches, offering a scalable, semantically rich alternative for self-supervised visual representation learning.

Abstract

Masked Image Modeling (MIM) has emerged as a promising method for deriving visual representations from unlabeled image data by predicting missing pixels from masked portions of images. It excels in region-aware learning and provides strong initializations for various tasks, but struggles to capture high-level semantics without further supervised fine-tuning, likely due to the low-level nature of its pixel reconstruction objective. A promising yet unrealized framework is learning representations through masked reconstruction in latent space, combining the locality of MIM with the high-level targets. However, this approach poses significant training challenges as the reconstruction targets are learned in conjunction with the model, potentially leading to trivial or suboptimal solutions.Our study is among the first to thoroughly analyze and address the challenges of such framework, which we refer to as Latent MIM. Through a series of carefully designed experiments and extensive analysis, we identify the source of these challenges, including representation collapsing for joint online/target optimization, learning objectives, the high region correlation in latent space and decoding conditioning. By sequentially addressing these issues, we demonstrate that Latent MIM can indeed learn high-level representations while retaining the benefits of MIM models.

Figures (13)

• Figure 1: Challenges of Latent MIM. The representations learned by MIM approaches fail to capture high-level semantics, as shown by the poor performance in nearest neighbor and linear probe evaluation.
• Figure 2: Latent Masked Image Modeling Overview. Models are trained to reconstruct the latent representations generated by a target encoder at withheld locations. Four major challenges for effectively deploying Latent MIM are identified in this work, as well as potential solutions. These challenges relate to joint encoder optimization, direct reconstruction loss, the semantic correlation between visible and target patches, and the decoder design.
• Figure 3: Training Collapse of the Naive Latent MIM. This solution achieves a zero reconstruction loss but fails to capture any meaningful information about the input images. As a result, the nearest neighbor (NN) evaluation yields random performance. Top: NN accuracy; Bottom: training loss.
• Figure 4: Patch Generation Strategies. Left: Masking contiguous grids. Right: Non-contiguous stochastic masking.
• Figure 5: Comparison of Three Different Decoder Designs. Self-attention decoder is commonly used for low-level MIM models. Cross-attention decoder provides direct conditioning at each layer on the visible latents.