Masked VAE: Distributionally-Informed Self-Supervised Vision Learning

Gabe Guo, Richard Zemel

Keywords: Self-Supervised Learning, Transformers, Vision, Masked Autoencoders

TL;DR: We present a new self-supervised learning method for vision that respects the inherent multi-mode nature of masked autoencoding tasks.

Abstract:

Masked pre-training with transformers is a popular self-supervised representation learning paradigm, initially showing success in NLP before moving to CV.
However, an aspect of masked pre-training that is missing in CV is the ability to capture the distribution of possible outputs. In NLP pre-training methods, the distribution is expressed as a softmax output layer. In CV, the SoTA masked autoencoder (MAE) simply ignores the possibility of a distribution, only giving a point estimate of the masked pixels' RGB values. This formulation is fundamentally limited, as it models an under-constrained problem as well-posed.
This poses limitations when deployed for completion tasks: it can give only one possible completion, when in reality, a scene could be completed in many different ways, e.g., a partial kitchen could have spoons, cups, pizzas, etc. under the mask. This inability to complete multiple modes indicates the weakness of the underlying representation in capturing contextual relationships.
Towards creating a distributionally-aware formulation with contextually-aware representations, we propose the Masked VAE, a transformer-based self-supervised learning method that combines ideas from the MAE and the variational autoencoder (VAE). Like the VAE, we model the "masked" latent space tokens as samples from a multivariate Gaussian distribution, while keeping the MAE's deterministic latent codes for the visible tokens.
Evaluations show that our method can create contextually plausible masked completions in a distributionally-aware manner, while matching the state-of-the-art in representation performance in downstream classification tasks.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 2721