Diffusion Transformer con Autoencoder di Rappresentazione

Diffusion Transformers with Representation Autoencoders

October 13, 2025

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

cs.AI

Abstract

La modellazione generativa latente, in cui un autoencoder pre-addestrato mappa i pixel in uno spazio latente per il processo di diffusione, è diventata la strategia standard per i Diffusion Transformers (DiT); tuttavia, il componente autoencoder è rimasto pressoché invariato. La maggior parte dei DiT continua a fare affidamento sull'originario encoder VAE, che introduce diverse limitazioni: backbone obsoleti che compromettono la semplicità architetturale, spazi latenti a bassa dimensionalità che riducono la capacità informativa e rappresentazioni deboli derivanti da un addestramento puramente basato sulla ricostruzione, limitando infine la qualità generativa. In questo lavoro, esploriamo la sostituzione del VAE con encoder di rappresentazione pre-addestrati (ad esempio, DINO, SigLIP, MAE) abbinati a decoder addestrati, formando ciò che definiamo Representation Autoencoders (RAE). Questi modelli offrono sia ricostruzioni di alta qualità che spazi latenti semanticamente ricchi, consentendo al contempo un'architettura scalabile basata su transformer. Poiché questi spazi latenti sono tipicamente ad alta dimensionalità, una sfida chiave è consentire ai diffusion transformer di operare efficacemente al loro interno. Analizziamo le fonti di questa difficoltà, proponiamo soluzioni teoricamente motivate e le validiamo empiricamente. Il nostro approccio raggiunge una convergenza più rapida senza l'uso di perdite ausiliarie di allineamento delle rappresentazioni. Utilizzando una variante DiT dotata di una testa DDT leggera e ampia, otteniamo risultati forti nella generazione di immagini su ImageNet: 1.51 FID a 256x256 (senza guida) e 1.13 sia a 256x256 che a 512x512 (con guida). RAE offre chiari vantaggi e dovrebbe diventare il nuovo standard per l'addestramento dei diffusion transformer.

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.