Scaling dei Trasformatori di Diffusione Text-to-Image con Autoencoder di Rappresentazione

Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

January 22, 2026

Autori: Shengbang Tong, Boyang Zheng, Ziteng Wang, Bingda Tang, Nanye Ma, Ellis Brown, Jihan Yang, Rob Fergus, Yann LeCun, Saining Xie

cs.AI

Abstract

Gli Autoencoder di Rappresentazione (RAE) hanno dimostrato vantaggi distinti nella modellazione di diffusione su ImageNet addestrandosi in spazi latenti semantici ad alta dimensionalità. In questo lavoro, investigiamo se questo framework può scalare alla generazione testo-immagine (T2I) su larga scala e libera da vincoli. Inizialmente scaliamo i decoder RAE sul codificatore di rappresentazione congelato (SigLIP-2) oltre ImageNet addestrandoli su dati web, sintetici e di rendering testuale, riscontrando che sebbene la scala migliori la fedeltà generale, una composizione dei dati mirata è essenziale per domini specifici come il testo. Successivamente, sottoponiamo a rigorosi test di stress le scelte progettuali dei RAE originariamente proposte per ImageNet. La nostra analisi rivela che il ridimensionamento semplifica il framework: sebbene la pianificazione del rumore dipendente dalla dimensione rimanga critica, complessità architetturali come teste di diffusione ampie e decodifica aumentata con rumore offrono benefici trascurabili su larga scala. Basandoci su questo framework semplificato, conduciamo un confronto controllato tra RAE e il VAE FLUX allo stato dell'arte attraverso scale di transformer di diffusione da 0,5 a 9,8 miliardi di parametri. I RAE superano costantemente i VAE durante il pre-addestramento su tutte le scale del modello. Inoltre, durante il fine-tuning su dataset di alta qualità, i modelli basati su VAE vanno incontro a un overfitting catastrofico dopo 64 epoche, mentre i modelli RAE rimangono stabili fino a 256 epoche e raggiungono prestazioni costantemente migliori. In tutti gli esperimenti, i modelli di diffusione basati su RAE dimostrano una convergenza più rapida e una migliore qualità di generazione, stabilendo i RAE come una base più semplice e solida rispetto ai VAE per la generazione T2I su larga scala. Inoltre, poiché sia la comprensione che la generazione visiva possono operare in uno spazio di rappresentazione condiviso, il modello multimodale può ragionare direttamente sui latenti generati, aprendo nuove possibilità per modelli unificati.

English

Representation Autoencoders (RAEs) have shown distinct advantages in diffusion modeling on ImageNet by training in high-dimensional semantic latent spaces. In this work, we investigate whether this framework can scale to large-scale, freeform text-to-image (T2I) generation. We first scale RAE decoders on the frozen representation encoder (SigLIP-2) beyond ImageNet by training on web, synthetic, and text-rendering data, finding that while scale improves general fidelity, targeted data composition is essential for specific domains like text. We then rigorously stress-test the RAE design choices originally proposed for ImageNet. Our analysis reveals that scaling simplifies the framework: while dimension-dependent noise scheduling remains critical, architectural complexities such as wide diffusion heads and noise-augmented decoding offer negligible benefits at scale Building on this simplified framework, we conduct a controlled comparison of RAE against the state-of-the-art FLUX VAE across diffusion transformer scales from 0.5B to 9.8B parameters. RAEs consistently outperform VAEs during pretraining across all model scales. Further, during finetuning on high-quality datasets, VAE-based models catastrophically overfit after 64 epochs, while RAE models remain stable through 256 epochs and achieve consistently better performance. Across all experiments, RAE-based diffusion models demonstrate faster convergence and better generation quality, establishing RAEs as a simpler and stronger foundation than VAEs for large-scale T2I generation. Additionally, because both visual understanding and generation can operate in a shared representation space, the multimodal model can directly reason over generated latents, opening new possibilities for unified models.