# Visualizing the loss landscape of Self-supervised Vision Transformer

arXiv: https://arxiv.org/abs/2405.18042 · HTML: https://mlsanity.com/paper/2405.18042/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 38 references that are in the map.
- First posted on arXiv 28 May 2024.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [Improved Baselines with Momentum Contrastive Learning](https://mlsanity.com/paper/2003.04297/) (2020)
- [Is Space-Time Attention All You Need for Video Understanding?](https://mlsanity.com/paper/2102.05095/) (2021)
- [An Empirical Study of Training Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.02057/) (2021)
- [Do ImageNet Classifiers Generalize to ImageNet?](https://mlsanity.com/paper/1902.10811/) (2019)
- [Big Self-Supervised Models are Strong Semi-Supervised Learners](https://mlsanity.com/paper/2006.10029/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [VideoMAE: Masked Autoencoders are Data-Efficient Learners for Self-Supervised Video Pre-Training](https://mlsanity.com/paper/2203.12602/) (2022)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [iBOT: Image BERT Pre-Training with Online Tokenizer](https://mlsanity.com/paper/2111.07832/) (2021)
- [Multiscale Vision Transformers](https://mlsanity.com/paper/2104.11227/) (2021)
- [Masked Autoencoders As Spatiotemporal Learners](https://mlsanity.com/paper/2205.09113/) (2022)
- [Entropy-SGD: Biasing Gradient Descent Into Wide Valleys](https://mlsanity.com/paper/1611.01838/) (2016)
- [XCiT: Cross-Covariance Image Transformers](https://mlsanity.com/paper/2106.09681/) (2021)
- [Masked Siamese Networks for Label-Efficient Learning](https://mlsanity.com/paper/2204.07141/) (2022)
- [When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations](https://mlsanity.com/paper/2106.01548/) (2021)
- [Understanding and Robustifying Differentiable Architecture Search](https://mlsanity.com/paper/1909.09656/) (2019)
- [A Bayesian Perspective on Generalization and Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06451/) (2017)
- [Benchmarking Detection Transfer Learning with Vision Transformers](https://mlsanity.com/paper/2111.11429/) (2021)
- [MPViT: Multi-Path Vision Transformer for Dense Prediction](https://mlsanity.com/paper/2112.11010/) (2021)

## Abstract (arXiv)

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.