Diffusie Transformers met Representatie Autoencoders

Diffusion Transformers with Representation Autoencoders

October 13, 2025

Auteurs: Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie

cs.AI

Samenvatting

Latente generatieve modellering, waarbij een vooraf getrainde autoencoder pixels afbeeldt naar een latente ruimte voor het diffusieproces, is de standaardstrategie geworden voor Diffusion Transformers (DiT); de autoencoder-component is echter nauwelijks geëvolueerd. De meeste DiT's blijven vertrouwen op de originele VAE-encoder, wat verschillende beperkingen met zich meebrengt: verouderde backbones die architectonische eenvoud in gevaar brengen, laagdimensionale latente ruimtes die de informatiecapaciteit beperken, en zwakke representaties die voortkomen uit puur op reconstructie gebaseerde training en uiteindelijk de generatieve kwaliteit beperken. In dit werk onderzoeken we het vervangen van de VAE door vooraf getrainde representatie-encoders (bijv. DINO, SigLIP, MAE) in combinatie met getrainde decoders, wat we Representation Autoencoders (RAE's) noemen. Deze modellen bieden zowel hoogwaardige reconstructies als semantisch rijke latente ruimtes, terwijl ze een schaalbare transformer-gebaseerde architectuur mogelijk maken. Omdat deze latente ruimtes doorgaans hoogdimensionaal zijn, is een belangrijke uitdaging om diffusie-transformers effectief te laten werken binnen deze ruimtes. We analyseren de bronnen van deze moeilijkheid, stellen theoretisch onderbouwde oplossingen voor en valideren deze empirisch. Onze aanpak bereikt snellere convergentie zonder aanvullende representatie-uitlijningsverliezen. Met een DiT-variant uitgerust met een lichtgewicht, brede DDT-head behalen we sterke beeldgeneratieresultaten op ImageNet: 1.51 FID bij 256x256 (zonder begeleiding) en 1.13 bij zowel 256x256 als 512x512 (met begeleiding). RAE biedt duidelijke voordelen en zou de nieuwe standaard moeten zijn voor het trainen van diffusie-transformers.

English

Latent generative modeling, where a pretrained autoencoder maps pixels into a latent space for the diffusion process, has become the standard strategy for Diffusion Transformers (DiT); however, the autoencoder component has barely evolved. Most DiTs continue to rely on the original VAE encoder, which introduces several limitations: outdated backbones that compromise architectural simplicity, low-dimensional latent spaces that restrict information capacity, and weak representations that result from purely reconstruction-based training and ultimately limit generative quality. In this work, we explore replacing the VAE with pretrained representation encoders (e.g., DINO, SigLIP, MAE) paired with trained decoders, forming what we term Representation Autoencoders (RAEs). These models provide both high-quality reconstructions and semantically rich latent spaces, while allowing for a scalable transformer-based architecture. Since these latent spaces are typically high-dimensional, a key challenge is enabling diffusion transformers to operate effectively within them. We analyze the sources of this difficulty, propose theoretically motivated solutions, and validate them empirically. Our approach achieves faster convergence without auxiliary representation alignment losses. Using a DiT variant equipped with a lightweight, wide DDT head, we achieve strong image generation results on ImageNet: 1.51 FID at 256x256 (no guidance) and 1.13 at both 256x256 and 512x512 (with guidance). RAE offers clear advantages and should be the new default for diffusion transformer training.