Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Robust Representation Learning in Masked Autoencoders

Anika Shrivastava, Renu Rameshan, Samar Agnihotri

TL;DR

The paper investigates why Masked Autoencoders (MAEs) yield strong downstream classification by analyzing how their latent representations organize across depth. It reveals that, even without supervision, MAE pretrained encoders develop class-separable subspaces that become increasingly distinct in deeper layers, coupled with persistent global attention from early layers. To assess robustness, the authors fine-tune MAEs and test under Gaussian blur and attention-guided occlusion, showing stable performance across many perturbations and strong results on ImageNet-C. They introduce two complementary robustness indicators—directional invariance and head-wise feature retention—to quantify representational stability, linking robust latent structure to robust classification performance and offering insights for robustness-aware representation learning.

Abstract

Masked Autoencoders (MAEs) achieve impressive performance in image classification tasks, yet the internal representations they learn remain less understood. This work started as an attempt to understand the strong downstream classification performance of MAE. In this process we discover that representations learned with the pretraining and fine-tuning, are quite robust - demonstrating a good classification performance in the presence of degradations, such as blur and occlusions. Through layer-wise analysis of token embeddings, we show that pretrained MAE progressively constructs its latent space in a class-aware manner across network depth: embeddings from different classes lie in subspaces that become increasingly separable. We further observe that MAE exhibits early and persistent global attention across encoder layers, in contrast to standard Vision Transformers (ViTs). To quantify feature robustness, we introduce two sensitivity indicators: directional alignment between clean and perturbed embeddings, and head-wise retention of active features under degradations. These studies help establish the robust classification performance of MAEs.

Robust Representation Learning in Masked Autoencoders

TL;DR

The paper investigates why Masked Autoencoders (MAEs) yield strong downstream classification by analyzing how their latent representations organize across depth. It reveals that, even without supervision, MAE pretrained encoders develop class-separable subspaces that become increasingly distinct in deeper layers, coupled with persistent global attention from early layers. To assess robustness, the authors fine-tune MAEs and test under Gaussian blur and attention-guided occlusion, showing stable performance across many perturbations and strong results on ImageNet-C. They introduce two complementary robustness indicators—directional invariance and head-wise feature retention—to quantify representational stability, linking robust latent structure to robust classification performance and offering insights for robustness-aware representation learning.

Abstract

Masked Autoencoders (MAEs) achieve impressive performance in image classification tasks, yet the internal representations they learn remain less understood. This work started as an attempt to understand the strong downstream classification performance of MAE. In this process we discover that representations learned with the pretraining and fine-tuning, are quite robust - demonstrating a good classification performance in the presence of degradations, such as blur and occlusions. Through layer-wise analysis of token embeddings, we show that pretrained MAE progressively constructs its latent space in a class-aware manner across network depth: embeddings from different classes lie in subspaces that become increasingly separable. We further observe that MAE exhibits early and persistent global attention across encoder layers, in contrast to standard Vision Transformers (ViTs). To quantify feature robustness, we introduce two sensitivity indicators: directional alignment between clean and perturbed embeddings, and head-wise retention of active features under degradations. These studies help establish the robust classification performance of MAEs.
Paper Structure (21 sections, 4 equations, 8 figures, 3 tables)

This paper contains 21 sections, 4 equations, 8 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Masked image modelling
  4. Understanding Masked Autoencoders
  5. Proposed analysis pipeline
  6. Layer-wise structural analysis of the pretrained encoder
  7. Overview of MAE:
  8. Layer-wise embedding extraction:
  9. Subspace geometry:
  10. Fine-tuning robustness
  11. Perturbations.
  12. Representational robustness indicators:
  13. Directional Robustness:
  14. Feature-level robustness of attention-heads:
  15. Experiments and Results
  16. ...and 6 more sections

Figures (8)

  • Figure 1: t-SNE visualizations of token embeddings across encoder layers: (a) CLS tokens and (b) mean patch tokens.
  • Figure 2: t-SNE visualization of patch tokens across encoder layers.
  • Figure 3: Each dot shows the mean attention distance across images for one head in left: ViT (adapted from dosovitskiy2020image) and right: MAE.
  • Figure 4: Layer-wise distribution of principal angles $\theta_1$ (in degrees) between classes across layers.
  • Figure 5: Occlusion level vs Accuracy plot on ImageNet-1K dataset.
  • ...and 3 more figures