Масштабирование диффузионных трансформеров для преобразования текста в изображение с помощью репрезентационных автоэнкодеров

Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

January 22, 2026

Авторы: Shengbang Tong, Boyang Zheng, Ziteng Wang, Bingda Tang, Nanye Ma, Ellis Brown, Jihan Yang, Rob Fergus, Yann LeCun, Saining Xie

cs.AI

Аннотация

Автокодировщики представлений (RAE) продемонстрировали явные преимущества в диффузионном моделировании на ImageNet за счет обучения в высокоразмерных семантических латентных пространствах. В данной работе мы исследуем, может ли эта методология масштабироваться до задач крупномасштабной генерации изображений по свободному текстовому описанию (T2I). Сначала мы масштабируем декодеры RAE на основе замороженного кодировщика представлений (SigLIP-2) за пределы ImageNet, обучая их на веб-данных, синтетических данных и данных с рендерингом текста, и обнаруживаем, что хотя масштабирование улучшает общую точность, целевой состав данных критически важен для специфических доменов, таких как текст. Затем мы тщательно тестируем в стресс-условиях архитектурные решения RAE, изначально предложенные для ImageNet. Наш анализ показывает, что масштабирование упрощает фреймворк: хотя зависимое от размерности планирование шума остается критически важным, архитектурные усложнения, такие как широкие диффузионные головы и декодирование с добавлением шума, дают незначительный выигрыш при масштабировании. На основе этого упрощенного фреймворка мы проводим контролируемое сравнение RAE с передовым FLUX VAE в диапазоне масштабов диффузионных трансформеров от 0.5 до 9.8 миллиардов параметров. RAE стабильно превосходят VAE на этапе предварительного обучения на всех масштабах моделей. Более того, при дообучении на высококачественных наборах данных модели на основе VAE катастрофически переобучаются после 64 эпох, в то время как модели RAE остаются стабильными в течение 256 эпох и демонстрируют стабильно лучшие результаты. Во всех экспериментах диффузионные модели на основе RAE показывают более быструю сходимость и лучшее качество генерации, что утверждает RAE как более простую и мощную основу по сравнению с VAE для крупномасштабной T2I-генерации. Дополнительно, поскольку как визуальное понимание, так и генерация могут работать в общем пространстве представлений, мультимодальная модель может напрямую выполнять логический вывод на сгенерированных латентных переменных, открывая новые возможности для унифицированных моделей.

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.