Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Visualizing the loss landscape of Self-supervised Vision Transformer

Youngwan Lee, Jeffrey Ryan Willette, Jonghee Kim, Sung Ju Hwang

TL;DR

This paper addresses why self-supervised MAE-based Vision Transformers generalize better than fully supervised ViTs by visualizing the pretraining loss landscapes. Using filter-normalized 2D projections of the pretraining loss, it compares MAE, RC-MAE with an EMA teacher, and supervised ViT, revealing that MAE-based ViTs have smoother, wider convex regions and that the EMA teacher in RC-MAE further widens these regions to enable faster convergence. The findings offer qualitative insights into optimization dynamics, suggesting that gradient corrections from the EMA teacher foster more favorable loss geometry during both pretraining and linear probing. The work motivates quantitative follow-ups and broader comparisons across self-supervised methods, with potential implications for designing training dynamics that promote generalization.

Abstract

The Masked autoencoder (MAE) has drawn attention as a representative self-supervised approach for masked image modeling with vision transformers. However, even though MAE shows better generalization capability than fully supervised training from scratch, the reason why has not been explored. In another line of work, the Reconstruction Consistent Masked Auto Encoder (RC-MAE), has been proposed which adopts a self-distillation scheme in the form of an exponential moving average (EMA) teacher into MAE, and it has been shown that the EMA-teacher performs a conditional gradient correction during optimization. To further investigate the reason for better generalization of the self-supervised ViT when trained by MAE (MAE-ViT) and the effect of the gradient correction of RC-MAE from the perspective of optimization, we visualize the loss landscapes of the self-supervised vision transformer by both MAE and RC-MAE and compare them with the supervised ViT (Sup-ViT). Unlike previous loss landscape visualizations of neural networks based on classification task loss, we visualize the loss landscape of ViT by computing pre-training task loss. Through the lens of loss landscapes, we find two interesting observations: (1) MAE-ViT has a smoother and wider overall loss curvature than Sup-ViT. (2) The EMA-teacher allows MAE to widen the region of convexity in both pretraining and linear probing, leading to quicker convergence. To the best of our knowledge, this work is the first to investigate the self-supervised ViT through the lens of the loss landscape.

Visualizing the loss landscape of Self-supervised Vision Transformer

TL;DR

This paper addresses why self-supervised MAE-based Vision Transformers generalize better than fully supervised ViTs by visualizing the pretraining loss landscapes. Using filter-normalized 2D projections of the pretraining loss, it compares MAE, RC-MAE with an EMA teacher, and supervised ViT, revealing that MAE-based ViTs have smoother, wider convex regions and that the EMA teacher in RC-MAE further widens these regions to enable faster convergence. The findings offer qualitative insights into optimization dynamics, suggesting that gradient corrections from the EMA teacher foster more favorable loss geometry during both pretraining and linear probing. The work motivates quantitative follow-ups and broader comparisons across self-supervised methods, with potential implications for designing training dynamics that promote generalization.

Abstract

The Masked autoencoder (MAE) has drawn attention as a representative self-supervised approach for masked image modeling with vision transformers. However, even though MAE shows better generalization capability than fully supervised training from scratch, the reason why has not been explored. In another line of work, the Reconstruction Consistent Masked Auto Encoder (RC-MAE), has been proposed which adopts a self-distillation scheme in the form of an exponential moving average (EMA) teacher into MAE, and it has been shown that the EMA-teacher performs a conditional gradient correction during optimization. To further investigate the reason for better generalization of the self-supervised ViT when trained by MAE (MAE-ViT) and the effect of the gradient correction of RC-MAE from the perspective of optimization, we visualize the loss landscapes of the self-supervised vision transformer by both MAE and RC-MAE and compare them with the supervised ViT (Sup-ViT). Unlike previous loss landscape visualizations of neural networks based on classification task loss, we visualize the loss landscape of ViT by computing pre-training task loss. Through the lens of loss landscapes, we find two interesting observations: (1) MAE-ViT has a smoother and wider overall loss curvature than Sup-ViT. (2) The EMA-teacher allows MAE to widen the region of convexity in both pretraining and linear probing, leading to quicker convergence. To the best of our knowledge, this work is the first to investigate the self-supervised ViT through the lens of the loss landscape.
Paper Structure (8 sections, 3 equations, 2 figures, 1 table)

This paper contains 8 sections, 3 equations, 2 figures, 1 table.

Table of Contents

  1. Introduction
  2. Preminimary: MAE and RC-MAE
  3. Loss landscape
  4. Analysis
  5. Conclusion and Future works
  6. Acknowledgement
  7. Appendix
  8. Implementation Details

Figures (2)

  • Figure 1: Comparison between self-supervised vision transformer (ViT) dosovitskiy2021vit by MAE he2022mae and supervised ViT.MAE-ViT is drawn by the pre-training loss (masked patch prediction) and Supervised ViT is obtained by supervised image classification from scratch. We use the ViT-Base model. MAE-ViT converges at a smoother and wider convex region than the superivsed ViT.
  • Figure 2: Comparison between MAE he2022mae and RC-MAE lee2023rcmae using pre-training (left) and linear probing (right) weights. We use the ViT-B which is pre-trained for 1600 epochs with MAE or RC-MAE. While The pre-training computes the masked reconstruction loss through mean-square error loss, the linear probing task computes image classification loss with cross-entropy loss.