Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders

TL;DR

This paper takes Representation Autoencoders (RAE) — a frozen pretrained representation encoder (SigLIP-2) paired with a learned decoder, used as the latent space for diffusion — and asks whether they scale from class-conditional ImageNet to large-scale, freeform text-to-image generation. Under matched conditions, RAE-based DiTs (0.5B–9.8B) converge 4–4.6× faster than FLUX-VAE counterparts on GenEval / DPG-Bench, are less prone to overfitting during high-quality finetuning (VAE collapses after 64 epochs; RAE stable through 256), and the framework simplifies at scale: dimension-aware noise scheduling stays essential, but the wide DDT head and noise-augmented decoding from the original ImageNet RAE recipe become unnecessary. Because understanding and generation now live in the same SigLIP-2 latent space, the LLM can verify its own generations directly in latent space (no pixel decode/re-encode), enabling latent test-time scaling.

Key claims

Scaling decoder training beyond ImageNet improves general fidelity, but text rendering specifically requires text-targeted supervision — data composition matters, not just scale [§2, Tab. 1].
Dimension-dependent noise schedule shift is critical at T2I scale: removing it drops GenEval 49.6 → 23.6 and DPG-Bench 76.8 → 54.8 [§3.2].
The wide DDT head (DDT^w) gives +11.2 GenEval at 0.5B but its advantage saturates at 2.4B+ where the backbone width already exceeds the latent dimension [§3.3, Fig. 4(b)].
Noise-augmented decoding helps before ~15k steps but becomes negligible at convergence — it functions as early-training regularization, not a fundamental requirement [§3.3, Fig. 4(a)].
At 1.5B LLM + 2.4B DiT, RAE (SigLIP-2) achieves 4.0× speedup on GenEval and 4.6× on DPG-Bench versus FLUX-VAE [§4.1, Fig. 1].
RAE consistently beats VAE across DiT scales 0.5B → 9.8B and LLM scales 1.5B → 7B, with diminishing returns past ~6B DiT params [§4.1, Fig. 5].
During finetuning, VAE-based models catastrophically overfit after 64 epochs (training loss collapses to near-zero) while RAE models remain stable through 256 epochs [§4.2, Fig. 6, App. Fig. 9].
The RAE advantage generalizes across encoder choice: WebSSL ViT-L (no text alignment) also beats FLUX-VAE, though slightly trails SigLIP-2 [§4.1, Tab. 4].
Adding generative training does not degrade understanding-task scores (MME, TextVQA, AI2D, SEED, MMMU); RAE vs VAE in the generative path has little effect on understanding since both share the same frozen encoder [§5, Tab. 6].
Latent test-time scaling — using the LLM as a verifier on generated latents without decoding to pixels — yields consistent GenEval gains; “Answer Logits” best-of-32 lifts a 7B LLM + 5.5B DiT from 55.5 → 67.8 [§5, Tab. 5].

Method

The system follows the MetaQuery architecture: a pretrained Qwen-2.5 LLM (1.5B or 7B) prepends 256 learnable query tokens to the text prompt; the LLM jointly processes both, and a 2-layer MLP projects the query-token outputs into a DiT (LightningDiT design, flow-matching objective, 0.5B–9.8B). The DiT operates not in a compressed VAE space but on the 256-token × 1152-dim latents produced by a frozen SigLIP-2 So400M encoder; an RAE decoder (ViT) reconstructs pixels from those latents. The same SigLIP-2 features are reused for visual understanding via a separate MLP projector, so understanding and generation share one latent space.

The paper does three things on top of this: (1) scales the RAE decoder beyond ImageNet using web + synthetic + text-rendering data; (2) stress-tests the original RAE design choices — dimension-aware noise schedule shift (kept), wide DDT head and noise-augmented decoding (dropped at scale); (3) runs apples-to-apples comparisons against FLUX-VAE under identical training configurations, varying only the latent space. Finetuning uses the BLIP-3o 60k high-quality dataset starting from the 30k-step pretrained checkpoint.

Results

Headline number: at 1.5B LLM + 2.4B DiT, RAE reaches GenEval ≈ 49.5 / DPG-Bench ≈ 76.9 at 4–4.6× fewer iterations than the FLUX-VAE baseline (39.6 / 70.5 at full schedule) [Fig. 1, Tab. 4]. Reconstruction rFID-50k on ImageNet/YFCC/Text: RAE SigLIP-2 0.435 / 0.702 / 1.621 vs FLUX-VAE 0.288 / 0.410 / 0.638 — RAE is behind on reconstruction yet ahead on generation [Tab. 2]. Scaling DiT 0.5B → 9.8B keeps RAE strictly ahead at every size; LLM 1.5B → 7B helps RAE more than VAE [Fig. 5]. Finetuning: at 64 epochs RAE and VAE are comparable; at 256 epochs RAE has held or improved while VAE has collapsed [Fig. 6, Fig. 9]. Latent TTS with best-of-32 Answer-Logits scoring lifts GenEval +12.3 (55.5 → 67.8) on the 7B + 5.5B configuration [Tab. 5]. Understanding benchmarks are essentially flat between RAE-based, VAE-based, and understanding-only variants on MME / SEED / AI2D / MMMU [Tab. 6].

Why it’s interesting

Three things matter for Luma. First, this is the strongest evidence to date that high-dimensional semantic latents (SigLIP-2: 256 × 1152 = 295k dims) are a better diffusion target than compressed VAE latents at T2I scale, not just on ImageNet — and the gap widens with model size, which is the regime that matters for production. Second, the simpler-at-scale finding (wide DDT head and noise-augmented decoding drop out; only dim-aware noise schedule stays) is the kind of architectural-debt-reduction result that should propagate into Luma’s DiT recipes. Third, sharing the latent space between understanding and generation enables latent-space verification — the LLM scoring its own samples without pixel round-trips — which is a much cheaper TTS knob than re-rendering candidates, and which connects to ongoing work on inference-time scaling.