Understanding Masked Autoencoders via Hierarchical Latent Variable Models

TL;DR

This paper provides a theoretical understanding of masked autoencoders (MAE) by modeling data with hierarchical latent variables, explaining how hyperparameters influence the semantic level of learned representations.

Contribution

It introduces a hierarchical latent variable model to theoretically justify MAE's ability to extract high-level features and analyzes the impact of hyperparameters on the learned representations.

Findings

MAE can identify latent variables under certain assumptions

Hyperparameters like masking ratio and patch size affect the semantic level of representations

Theoretical explanations align with empirical observations

Abstract

Masked autoencoder (MAE), a simple and effective self-supervised learning framework based on the reconstruction of masked image regions, has recently achieved prominent success in a variety of vision tasks. Despite the emergence of intriguing empirical observations on MAE, a theoretically principled understanding is still lacking. In this work, we formally characterize and justify existing empirical insights and provide theoretical guarantees of MAE. We formulate the underlying data-generating process as a hierarchical latent variable model and show that under reasonable assumptions, MAE provably identifies a set of latent variables in the hierarchical model, explaining why MAE can extract high-level information from pixels. Further, we show how key hyperparameters in MAE (the masking ratio and the patch size) determine which true latent variables to be recovered, therefore influencing the level of semantic information in the representation. Specifically, extremely large or small masking ratios inevitably lead to low-level representations. Our theory offers coherent explanations of existing empirical observations and provides insights for potential empirical improvements and fundamental limitations of the masking-reconstruction paradigm. We conduct extensive experiments to validate our theoretical insights.