Diffusion Transformers with Representation Autoencoders

TL;DR

Representation Autoencoders (RAE) replace SD-VAE’s compressed convolutional latent with a frozen pretrained representation encoder (DINOv2, SigLIP-2, MAE) plus a trained ViT decoder, then run a DiT directly on those higher-dimensional semantic tokens. Three architectural adjustments make this work — DiT width must match or exceed the latent token dimension, a dimension-dependent noise-schedule shift, and noise-augmented decoder training — and a shallow-but-wide DDT head (DiTDH) scales width without quadratic FLOP cost. On ImageNet, RAE-based DiTDH-XL reaches gFID 1.51 at 256² without guidance and 1.13 with AutoGuidance at both 256² and 512², while converging ~16× faster than REPA on SD-VAE.

Key claims

Frozen representation encoders (DINOv2, SigLIP-2, MAE) paired with a trained ViT decoder beat SD-VAE on reconstruction rFID — MAE-B 0.16, DINOv2-B 0.49, SigLIP2-B 0.53 vs SD-VAE 0.62 [Table 1a].
The RAE decoder is also more compute-efficient than SD-VAE: ViT-B decoder reaches rFID 0.58 at 22.2 GFLOPs vs SD-VAE 0.62 at 310.4 GFLOPs; ViT-XL hits 0.49 at 106.7 GFLOPs [Table 1b].
Diffusion does not work out of the box in the RAE latent space: standard DiT-S on RAE blows up to gFID 215.76 vs 51.74 on SD-VAE, and DiT-XL underperforms SD-VAE 23.08 vs 7.13 [Table 2].
DiT width must match or exceed the encoder’s token dimension for training to converge — proven in a single-image overfit experiment and supported by a theoretical loss bound; below the width threshold the model cannot fit even one image [§4.1, Fig. 3, Thm. 3].
The dimension-dependent noise-schedule shift (generalizing Esser et al. 2024 from spatial resolution to the full effective data dimension N×d) is critical: removing it collapses gFID from 4.81 to substantially worse [§4.2, Table 4].
Noise-augmented decoder training (add 𝒩(0, σ²) to latents during decoder training, σ sampled from 𝒰(0, σ_max)) trades a small rFID hit (0.49 → 0.57) for a generation gain (gFID 4.81 → 4.28) by closing the train-test gap between RAE’s discrete training latents and the diffusion model’s continuous samples [§4.3, Table 5].
With these three fixes, plain DiT-XL on RAE reaches gFID 4.28 after 80 epochs and 2.39 after 720 epochs — a 47× speedup over SiT-XL on SD-VAE and 16× over REPA-XL [§4 closing].
DiTDH — augmenting DiT with a shallow, wide diffusion head — gives further FLOP-efficient scaling: DiTDH-B at 40% of DiT-XL’s training FLOPs already beats DiT-XL, and DiTDH-XL hits FID 2.16 (nearly half of DiT-XL’s) [§5, Fig. 6a].
DiTDH-XL with DINOv2-B sets a new ImageNet-256 SOTA: gFID 1.51 without guidance, 1.13 with AutoGuidance; at 512² it also reaches 1.13, surpassing EDM-2’s 1.25 [Table 7, Table 8].
Generation quality and reconstruction quality decouple: MAE has the best rFID, but DINOv2 produces the strongest generation, so DINOv2 is adopted as the default encoder for all generation experiments [§4 preamble].

Method

RAE has two stages. Stage 1 — RAE training: given a frozen pretrained representation encoder E (DINOv2-B at default; also SigLIP-2-B and MAE-B), train a ViT-XL decoder D on L1 + LPIPS + adversarial loss to reconstruct images from E’s 256 tokens × 768-dim representation. Encoder remains frozen; only the decoder is trained. Stage 2 — diffusion training: train a LightningDiT variant with the flow-matching objective directly on E’s latents (no patchification — patch size 1, so the DiT operates on 256 tokens to match SD-VAE-DiT sequence length). Three non-standard ingredients: (i) DiT hidden dimension ≥ latent token dimension (e.g. DiT-B for DINOv2-B’s d=768); (ii) noise-schedule shifted by N×d / N_base×d_base; (iii) decoder retrained with additive noise augmentation 𝒩(0, σ²) on latents.

DiTDH adds a shallow-but-wide transformer head h(x_t, h_DiT(z, t, y), t) after the DiT backbone — a 2-layer, 2048-dim head that takes the noisy latent and the DiT’s conditioning output and predicts velocity. This separates “conditioning processing” (deep, narrow DiT backbone) from “denoising” (shallow, wide head). FLOPs scale linearly in head width rather than quadratically in backbone width, so DiTDH attains the wide-DiT requirement of §4.1 cheaply.

Results

ImageNet-256 reconstruction: RAE+DINOv2-B 0.49 rFID, +MAE-B 0.16, +SigLIP2-B 0.53 vs SD-VAE 0.62 [Table 1a]. DINOv2-B linear probe: 84.5% top-1 vs SD-VAE 8.0% [Table 1d].
DiT-XL on RAE (DINOv2-B): gFID 4.28 at 80 epochs; 2.39 at 720 epochs — 47× / 16× speedup over SiT-XL on SD-VAE / REPA-XL respectively [§4 closing].
DiTDH-XL on RAE (DINOv2-B): gFID 1.51 (no guidance) / 1.13 (AutoGuidance) at 256² in 800 epochs [Table 7]; gFID 1.13 at 512² in 400 epochs, beating EDM-2’s 1.25 [Table 8]. Early-epoch numbers: 3.71 at 20 epochs, 2.16 at 80 epochs [Table 7].
DiTDH-XL beats DiT-XL across all three encoder sizes (DINOv2-S/B/L), and the advantage grows with encoder size — with DINOv2-L the DiTDH improvement is 6.09 → 2.73 [Table 6].
DiTDH-S already beats REPA-XL despite being smaller [§5.1 scaling discussion, Fig. 6c].

Why it’s interesting

This is the foundational paper of the Representation Autoencoders strand and the load-bearing reference for almost everything in Diffusion training efficiency that involves DINO features in the latent. The wiki already holds the two direct follow-ups — Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders scales the same recipe from class-conditional ImageNet to T2I (4–4.6× speedup over FLUX-VAE, dropping the wide DDT head and noise-augmented decoder at scale), and Improved Baselines with Representation Autoencoders (RAEv2) improves the recipe itself with multi-layer encoder aggregation and shows RAE and REPA are complementary. Filing the original paper closes the citation loop: the conflict on Diffusion training efficiency‘s key claims (RAE / REPA / Self-Flow / SRA) now has its anchor. The three architectural rules — width ≥ token dim, dim-aware noise shift, noise-augmented decoder — are the practical contribution most likely to propagate into Luma’s DiT recipes, especially because (1) and (2) are zero-cost design constraints rather than training-time additions.