PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

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

arXiv:2601.16208v1cs.CV

TL;DR

The paper asks whether Representation Autoencoders can scale from ImageNet to large-scale, freeform text-to-image generation. It scales decoders and evaluates simplified RAE designs against FLUX VAE under matched diffusion-transformer training, finding faster convergence, better generation, and greater finetuning stability for RAE models. Shared semantic latents also support unified understanding and generation with direct processing of generated latents.

  • Problem

    Whether high-dimensional RAE latent diffusion can handle the diversity, open-ended compositions, and scale of freeform text-to-image generation remains unproven beyond ImageNet.

  • Method

    The paper scales SigLIP-2-based RAE decoders with web, synthetic, and text-rendering data, simplifies RAE design choices, and compares RAE with FLUX VAE under matched T2I training.

  • Results

    RAE-based diffusion models consistently outperform VAE baselines in convergence speed and generation quality while remaining less prone to overfitting during finetuning.

  • Takeaways & Limitations

    RAE provides a simpler and stronger foundation for large-scale text-to-image diffusion, while shared representation spaces enable direct reasoning over generated latents.

  • Takeaways & Limitations

    Performance gains from scaling DiTs diminish beyond 6B parameters without proportionally better data quality and diversity.

Abstract

from arXiv · show

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.

1. Introduction

The paper asks whether representation-based diffusion can scale beyond curated ImageNet to freeform text-to-image generation. It scales RAE decoders, simplifies the framework, and finds faster convergence, stronger generation, and shared-latent multimodal capabilities than VAE-based alternatives.

  • Motivation: ImageNet results do not establish whether high-dimensional RAE latents support diverse, open-ended, large-scale text-to-image generation.The paper identifies broader visual diversity, open-ended compositions, and larger models and compute as unresolved challenges.
  • Decoder scaling: RAE decoders generalize beyond ImageNet, but text reconstruction requires targeted text-rendering data rather than scale alone.Web-scale and synthetic data improve diverse natural-image reconstruction, while text-specific data is needed for fine glyph details.
  • Framework design: Dimension-aware noise scheduling remains essential, whereas wide diffusion heads and noise-augmented decoding provide diminishing benefits at large T2I scale.The wide head benefits smaller backbones, but its advantage fades as DiTs reach billion-parameter scale; noise-augmented decoding gains saturate quickly.
  • RAE versus VAE: 4.0× on GenEval and 4.6× on DPG-Bench are the reported RAE speedups over VAE for a 1.5B LLM plus 2.4B DiT.The comparison trains both models from scratch in their respective latent spaces for up to 60k iterations.
  • RAE versus VAE: RAE models converge faster, perform better, and overfit less than VAE counterparts across language-backbone and diffusion-model scales.The reported advantage spans Qwen-2.5 1.5B–7B backbones and DiT 0.5B–9.8B scales, with finetuning performance also favoring RAE.
  • Unified models: Shared semantic latents let unified models process generated representations directly while preserving understanding performance during generative training.The paper presents latent-space test-time scaling as a feasible and effective exploratory use of this shared space.
  • Conclusion: RAEs are presented as a simpler and stronger foundation than VAEs for scaling text-to-image diffusion in semantic representation spaces.The authors state that code, data, and model checkpoints will be released for open and reproducible multimodal-generation research.

2. Scaling Decoder Training Beyond ImageNet

The decoder study trains reconstruction models on frozen SigLIP-2 representations using diverse web, synthetic, and text-rendering data. Broader data improves generalization, domain-matched composition is crucial for text, and WebSSL-DINO outperforms SigLIP-2 across evaluated domains.

  • Decoder setup: The decoder reconstructs 224×224 images from 16×16 SigLIP-2 tokens with channel dimension d = 1152.SigLIP-2 So400M is frozen while a ViT-based decoder is trained.
  • Decoder setup: The reconstruction objective combines ℓ1, LPIPS, Gram, and adversarial losses.The paper adds Gram Loss to the standard RAE reconstruction losses.
  • Training data: The training data combines roughly 73M web images, synthetic images, and RenderedText examples.RenderedText specifically targets text-rendering scenes.
  • Data scaling: Web-scale and synthetic data produce marginal ImageNet gains but moderate improvements on diverse YFCC images.The evaluation uses rFID-50k on ImageNet-1k, YFCC, and a held-out RenderedText set.
  • Data scaling: Text-specific data substantially improves typography reconstruction, showing that data composition matters more than size alone for domain-specific fidelity.Web plus synthetic data provides little improvement over ImageNet-only training for text reconstruction.
  • Encoder comparison: WebSSL-DINO reconstructs better than SigLIP-2 across natural-image and text domains, while both RAE variants outperform SDXL-VAE but trail FLUX-VAE.The comparison covers ImageNet, YFCC, and text-rendering reconstruction.
  • Pipeline: The pipeline trains a decoder from frozen encoder tokens, then jointly trains language and image generation with CE and flow-matching objectives.The unified stage includes an autoregressive model, diffusion transformer, and learnable query tokens.

3. RAE is Simpler in T2I

At large-scale T2I, dimension-aware noise scheduling remains essential, while architectural and decoding refinements from ImageNet become less important as backbone capacity grows.

  • Noise scheduling: Dimension-aware noise scheduling is critical for convergence in high-dimensional RAE latents.Removing the shift substantially worsens GenEval and DPG-Bench.
  • Noise-augmented decoding: Noise-augmented decoding provides early training gains but becomes negligible at later stages.Its gains are noticeable before approximately 15k steps and saturate quickly.
  • Wide DDT head: At 0.5B, DiTDH provides a +11.2 GenEval boost, but its advantage largely disappears at 2.4B and beyond.The wide head addresses narrow-backbone capacity constraints rather than a fundamental RAE requirement.
  • Resulting design: Standard DiT architectures with proper noise scheduling and no noise-augmented decoding are adopted for scalable T2I training.Backbone capacity increasingly dominates performance as models grow.

4. Training Diffusion Model with RAE vs. VAE

Under matched T2I training conditions, RAE models converge faster and outperform VAE models across diffusion scales, language backbones, and finetuning settings. RAE models also remain more stable during extended finetuning, whereas VAE models overfit after 64 epochs.

  • Data composition: 49.5 GenEval from combined data exceeds 48.0 from doubled synthetic data, showing that complementary data sources matter more than volume alone.The mixed dataset combines web and synthetic data, while doubling individual sources produces smaller improvements.
  • Scaling: RAE-based models outperform VAE-based models across DiT sizes from 0.5B to 9.8B parameters.The advantage holds with both 1.5B and 7B LLM backbones.
  • Scaling limitation: DiT scaling beyond 6B parameters shows diminishing returns without proportionally better data quality and diversity.The performance trend appears to plateau, suggesting underutilized capacity.
  • Finetuning: VAE-based models degrade significantly after 64 finetuning epochs, while RAE models remain stable with only a mild decline through 256 epochs.The authors hypothesize that RAE’s high-dimensional semantic latent space provides implicit regularization.
  • Generalization: RAE maintains its advantage when finetuning only the DiT or the full LLM+DiT system, with larger gains at larger model scales.The weakest RAE approach reaches 79.4, compared with 78.2 for the top-performing VAE configuration.

5. Implications for Unified Models

RAE places visual understanding and generation in a shared semantic latent space, allowing the LLM to evaluate generated latents directly. This supports latent-space test-time scaling while leaving visual understanding performance intact.

  • Shared latent space: A shared representation lets the LLM process diffusion-generated latents directly rather than decoding and re-encoding images.This decouples representation and pixel spaces in the unified model.
  • Latent test-time scaling: Latent-space test-time scaling uses the LLM as a verifier of generated representations without rendering pixels.Prompt Confidence and Answer Logits both produce consistent GenEval improvements with best-of-N selection.
  • Visual understanding: Generative training does not degrade visual understanding, and RAE and VAE achieve similar understanding performance.Both variants share the same frozen understanding encoder across the evaluated vision-language benchmarks.

6. Related Work

Prior work largely uses compressed VAE latents for generation or separate representations for understanding and generation. RAE instead trains a decoder over frozen, high-dimensional semantic representations and extends this approach to T2I.

  • VAE-based generation: VAEs compress images into low-dimensional latents, causing information loss through aggressive compression or quantization.RAE follows a different route by retaining high-dimensional semantic features.
  • Representation Autoencoders: RAE trains only a decoder on a frozen pretrained representation encoder to reconstruct pixels from high-dimensional semantic features.Diffusion then operates directly in the semantic latent space.
  • T2I extension: This work extends RAE from ImageNet to text-to-image generation and reports faster convergence and better performance than VAEs.The representation route is presented as an alternative to modifying VAE architectures.
  • Unified multimodal models: Unified multimodal approaches seek to reconcile distinct representations for visual understanding and generation, including shared continuous latent spaces.Existing approaches may introduce substantial downsampling or use separate encoders for the two capabilities.

7. Conclusion

The study finds that Representation Autoencoders scale effectively to large-scale text-to-image generation, with simpler designs outperforming VAEs and remaining more robust during finetuning. Shared representation spaces also support unified multimodal modeling.

  • RAE-based diffusion models consistently outperform state-of-the-art VAE baselines in convergence speed and generation quality.
  • Scaling simplifies RAE design: dimension-dependent noise scheduling remains essential, while architectural modifications such as DiTDH provide diminishing returns.
  • Larger decoder training data improves general fidelity, but targeted data composition remains necessary for domains such as text.
  • RAEs remain less prone to overfitting during finetuning than VAE-based models.
  • A shared representation space lets understanding and generation operate together, enabling possibilities such as latent-space test-time scaling in unified models.

A. Implementation

The implementation uses large-scale accelerator training, tuned decoder and discriminator recipes, controlled RAE–VAE finetuning, and synthetic data generation for decoder and text-to-image corpora.

  • Decoder training uses a ViT-XL architecture with 28 blocks, hidden size 1152, and 16 heads.The decoder architecture largely follows RAE, while the GAN recipe is tuned for web-scale images.
  • Web-scale decoder training uses DINO-S/16 as a strong discriminator after the original recipe proved unstable.
  • Pretraining runs primarily use TPU-v5p-128 and TPU-v6e-64, with separate optimizers for the pretrained LLM and newly trained DiT.SPMD sharding and TorchXLA are used for the LLM, adapters, and DiT.
  • Finetuning compares RAE and VAE models under identical configurations across 4, 16, 64, and 256 epochs.The experiments use BLIP3o-60k, a global batch size of 1024, and the same codebase and infrastructure.
  • Synthetic-data generation produces 24.7M images with FLUX.1-schnell for decoder and text-to-image training.

B. Models

The unified models connect visual representations to language and diffusion spaces, while the diffusion transformer scaling strategy prioritizes width and keeps its hidden dimension above the target latent dimension.

  • Pretrained Qwen2.5 language models at 1.5B and 7B scales are used in the unified-model experiments.
  • A 2-layer MLP projects encoder visual features into the LLM embedding space, and a linear layer maps LLM query outputs into the diffusion input space.
  • The DiT architecture follows LightningDiT and scales primarily by increasing model width rather than depth.
  • The DiT hidden dimension remains strictly larger than the target latent dimension, including in the DiT-0.5B model.The example target latent dimension is 1152 for the SigLIP2 ViT-So.

C. Additional Results

Additional finetuning analyses show that VAE models overfit rapidly, whereas RAE models retain robust generation performance over substantially longer training horizons.

  • VAE loss decreases rapidly to a very low value during finetuning, correlating with performance degradation and overfitting.RAE loss decreases more gradually and stabilizes at a higher value while generation remains robust.
  • Table 7 provides separate decoder and discriminator training configurations for SigLIP2-So and WebSSL ViT-L.The caption notes that different encoders require slightly different recipes.
  • Table 8 specifies optimization hyperparameters for the unified model’s LLM backbone and DiT diffusion head.
  • RAE remains stable through 512 finetuning epochs, showing only a small performance decline.VAE models already suffer substantial performance drops by 256 epochs and are not trained further.
  • Table 10 lists architectural specifications for the DiT variants.
Loading 2601.16208v1…