# ViT-AE++: Improving Vision Transformer Autoencoder for Self-supervised Medical Image Representations

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 32 references that are in the map.
- First posted on arXiv 18 Jan 2023.

## Builds on (most cited first, 23 of 23)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [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)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [Pyramid Scene Parsing Network](https://mlsanity.com/paper/1612.01105/) (2016)
- [An Empirical Study of Training Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.02057/) (2021)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Identifying the Best Machine Learning Algorithms for Brain Tumor Segmentation, Progression Assessment, and Overall Survival Prediction in the BRATS Challenge](https://mlsanity.com/paper/1811.02629/) (2018)
- [Deeply-Supervised Nets](https://mlsanity.com/paper/1409.5185/) (2014)
- [Contrastive learning of global and local features for medical image segmentation with limited annotations](https://mlsanity.com/paper/2006.10511/) (2020)
- [Models Genesis: Generic Autodidactic Models for 3D Medical Image Analysis](https://mlsanity.com/paper/1908.06912/) (2019)
- [Positional Contrastive Learning for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/2106.09157/) (2021)
- [Contrastive Learning with Continuous Proxy Meta-Data for 3D MRI Classification](https://mlsanity.com/paper/2106.08808/) (2021)
- [Imbalance-Aware Self-Supervised Learning for 3D Radiomic Representations](https://mlsanity.com/paper/2103.04167/) (2021)
- [Generic Perceptual Loss for Modeling Structured Output Dependencies](https://mlsanity.com/paper/2103.10571/) (2021)

## Built on by (most cited first, 12 of 12)

- [General surgery vision transformer: A video pre-trained foundation model for general surgery](https://mlsanity.com/paper/2403.05949/) (2024)
- [MiM: Mask in Mask Self-Supervised Pre-Training for 3D Medical Image Analysis](https://mlsanity.com/paper/2404.15580/) (2024)
- [DiCoM -- Diverse Concept Modeling towards Enhancing Generalizability in Chest X-Ray Studies](https://mlsanity.com/paper/2402.15534/) (2024)
- [Self-pruning Graph Neural Network for Predicting Inflammatory Disease Activity in Multiple Sclerosis from Brain MR Images](https://mlsanity.com/paper/2308.16863/) (2023)
- [Integration of deep generative Anomaly Detection algorithm in high-speed industrial line](https://mlsanity.com/paper/2603.07577/) (2026)
- [VariViT: A Vision Transformer for Variable Image Sizes](https://mlsanity.com/paper/2602.14615/) (2026)
- [H-SemiS: Hierarchical Fusion of Semi and Self-Supervised Learning for Knee Osteoarthritis Severity Grading](https://mlsanity.com/paper/2604.23335/) (2026)
- [Efficient Brain Imaging Analysis for Alzheimer's and Dementia Detection Using Convolution-Derivative Operations](https://mlsanity.com/paper/2411.13490/) (2024)
- [AI-Driven Radiology Report Generation for Traumatic Brain Injuries](https://mlsanity.com/paper/2510.08498/) (2025)
- [Google-MedGemma Based Abnormality Detection in Musculoskeletal radiographs](https://mlsanity.com/paper/2511.05600/) (2025)
- [Cryo-SWAN: the Multi-Scale Wavelet-decomposition-inspired Autoencoder Network for molecular density representation of molecular volumes](https://mlsanity.com/paper/2603.03342/) (2026)
- [A Multimodal Deep Learning Framework for Edema Classification Using HCT and Clinical Data](https://mlsanity.com/paper/2603.26726/) (2026)

## Abstract (arXiv)

Self-supervised learning has attracted increasing attention as it learns data-driven representation from data without annotations. Vision transformer-based autoencoder (ViT-AE) by He et al. (2021) is a recent self-supervised learning technique that employs a patch-masking strategy to learn a meaningful latent space. In this paper, we focus on improving ViT-AE (nicknamed ViT-AE++) for a more effective representation of 2D and 3D medical images. We propose two new loss functions to enhance the representation during training. The first loss term aims to improve self-reconstruction by considering the structured dependencies and indirectly improving the representation. The second loss term leverages contrastive loss to optimize the representation from two randomly masked views directly. We extended ViT-AE++ to a 3D fashion for volumetric medical images as an independent contribution. We extensively evaluate ViT-AE++ on both natural images and medical images, demonstrating consistent improvement over vanilla ViT-AE and its superiority over other contrastive learning approaches. Codes are here: https://github.com/chinmay5/vit_ae_plus_plus.git.