Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

TL;DR

This work investigates scaling Representation Autoencoders (RAEs) for large-scale text-to-image diffusion, demonstrating faster pretraining convergence and stronger generation than VAE baselines across diffusion backbones and language models. By decoupling the diffusion process from pixels and operating in a high-dimensional semantic latent space, RAEs achieve better data efficiency, robustness to overfitting during finetuning, and enable new ideas like latent-space test-time scaling. The study shows that data composition—particularly text-focused data—is crucial for text fidelity, while simple scaling with proper dimension-aware noise scheduling suffices to maintain performance, diminishing the need for architectural cheats at scale. The results position RAEs as a simpler and more effective foundation for scalable T2I generation and hint at opportunities for unified multimodal models that reason directly in latent representations.

Abstract

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.

Figures (10)

• Figure 1: RAE converges faster than VAE in text-to-image pretraining. We train Qwen-2.5 1.5B + DiT 2.4B models from scratch on both RAE (SigLIP-2) and VAE (FLUX) latent spaces for up to 60k iterations. RAE converges significantly faster than VAE on both GenEval (4.0×) and DPG-Bench (4.6×).
• Figure 2: RAE decoders trained on more data (web, synthetic & text) generalize across domains. Decoders trained only on ImageNet reconstruct natural images well but struggle with text-rendering scenes (see second column). Adding web and text data greatly improves text reconstruction while maintaining natural-image quality. We also observe that both the language-supervised model and the SSL model learn representations suitable for reconstructing diverse images, including natural languages. Compared to proprietary VAEs, our RAE models achieve competitive overall fidelity.
• Figure 3: Overview of training pipeline. Left: RAE decoder training stage. We train a decoder on the representations (yellow tokens) produced by the frozen RAE encoder. Right: End-to-end unified training of the autoregressive model, diffusion transformer, and learnable query tokens (gray tokens) using cross-entropy (CE) loss for text prediction and a flow-matching objective for image prediction.
• Figure 4: Design choices that saturate at T2I scale.Left: Noise-augmented decoding provides substantial gains early in training but becomes negligible by 120k steps. Right: $\text{DiT}^{\text{DH}}$ yields large gains at 0.5B (+11.2 GenEval), but the advantage diminishes at $>$2.4B, where backbone capacity dominates.
• Figure 5: RAE outperforms VAE across LLM and DiT scales.Top: With a 1.5B LLM, RAE-based models outperform VAE-based ones at all DiT sizes (0.5B, 2.4B, 5.5B, 9.8B). Bottom: Using a larger 7B LLM, RAE continues to maintain its advantage.