Escalado de Transformadores de Difusión de Texto a Imagen con Autoencoders de Representación

Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

January 22, 2026

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

cs.AI

Resumen

Los Autoencoders de Representación (RAE) han demostrado ventajas distintivas en el modelado de difusión en ImageNet al entrenarse en espacios latentes semánticos de alta dimensión. En este trabajo, investigamos si este marco puede escalar a la generación de texto a imagen (T2I) libre y a gran escala. Primero escalamos los decodificadores RAE en el codificador de representación congelado (SigLIP-2) más allá de ImageNet entrenando con datos web, sintéticos y de renderizado de texto, encontrando que, si bien la escala mejora la fidelidad general, la composición de datos dirigida es esencial para dominios específicos como el texto. Luego, sometemos a rigurosas pruebas de estrés las decisiones de diseño de RAE propuestas originalmente para ImageNet. Nuestro análisis revela que el escalado simplifica el marco: mientras que la programación de ruido dependiente de la dimensión sigue siendo crítica, las complejidades arquitectónicas como las cabezas de difusión anchas y la decodificación con aumento de ruido ofrecen beneficios insignificantes a gran escala. Basándonos en este marco simplificado, realizamos una comparación controlada de RAE frente al VAE FLUX de última generación a través de escalas de transformadores de difusión de 0.5B a 9.8B de parámetros. Los RAE superan consistentemente a los VAE durante el preentrenamiento en todas las escalas del modelo. Además, durante el ajuste fino en conjuntos de datos de alta calidad, los modelos basados en VAE sufren un sobreajuste catastrófico después de 64 épocas, mientras que los modelos RAE se mantienen estables hasta 256 épocas y logran un rendimiento consistentemente mejor. En todos los experimentos, los modelos de difusión basados en RAE demuestran una convergencia más rápida y una mejor calidad de generación, estableciendo a los RAE como una base más simple y sólida que los VAE para la generación T2I a gran escala. Adicionalmente, dado que tanto la comprensión como la generación visual pueden operar en un espacio de representación compartido, el modelo multimodal puede razonar directamente sobre los latentes generados, abriendo nuevas posibilidades para modelos unificados.

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.