# ViTCAE: ViT-based Class-conditioned Autoencoder

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 24 references that are in the map.
- First posted on arXiv 20 Sep 2025.

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

- [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)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [Ladder Variational Autoencoders](https://mlsanity.com/paper/1602.02282/) (2016)
- [TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up](https://mlsanity.com/paper/2102.07074/) (2021)
- [Koopcon: A new approach towards smarter and less complex learning](https://mlsanity.com/paper/2405.13866/) (2024)
- [Generative Expansion of Small Datasets: An Expansive Graph Approach](https://mlsanity.com/paper/2406.17238/) (2024)

## Abstract (arXiv)

Vision Transformer (ViT) based autoencoders often underutilize the global Class token and employ static attention mechanisms, limiting both generative control and optimization efficiency. This paper introduces ViTCAE, a framework that addresses these issues by re-purposing the Class token into a generative linchpin. In our architecture, the encoder maps the Class token to a global latent variable that dictates the prior distribution for local, patch-level latent variables, establishing a robust dependency where global semantics directly inform the synthesis of local details. Drawing inspiration from opinion dynamics, we treat each attention head as a dynamical system of interacting tokens seeking consensus. This perspective motivates a convergence-aware temperature scheduler that adaptively anneals each head's influence function based on its distributional stability. This process enables a principled head-freezing mechanism, guided by theoretically-grounded diagnostics like an attention evolution distance and a consensus/cluster functional. This technique prunes converged heads during training to significantly improve computational efficiency without sacrificing fidelity. By unifying a generative Class token with an adaptive attention mechanism rooted in multi-agent consensus theory, ViTCAE offers a more efficient and controllable approach to transformer-based generation.