Diffusion Transformers With Representation Autoencoders

September 19, 2026

frechet inception distance (FID) - measure quality of images created by a generative model by using the distribution of an Inception NN

-reconstruction objectives for VAE encoders produce low-capacity latents lacking global semantic structure
-higher dimensional encoder representations can be an advantage
-no extra compute/memory cost because they can be projected into DiT first-layer expected dims
R256×256×3  →encoder  R16×16×denc  →flatten  R256×denc  →linear projection  R256×dDiT\mathbb{R}^{256 \times 256 \times 3}\;\xrightarrow{\text{encoder}}\;\mathbb{R}^{16 \times 16 \times d_{\text{enc}}}\;\xrightarrow{\text{flatten}}\;\mathbb{R}^{256 \times d_{\text{enc}}}\;\xrightarrow{\text{linear projection}}\;\mathbb{R}^{256 \times d_{\text{DiT}}}

representation autoencoder (RAE) - replace VAE with a pretrained representation encoder (i.e. DINO) paired with a trained decoder

-basically, do diffusion in DINO feature space

Main Problems:

1.DiT input width from DINO outputs will have to be R256×768\mathbb{R}^{256 \times 768}, currently support a lot less (i.e. n×384)n \times 384), meaning that we lose a lot of information
1.Solution: introduce DiT^DH, which is a shallow but wide head to process RAE encoder outputs
2.Noise scheduling for 768 dimensions is much harder than for traditional lower dimensional denoising (i.e. R32×32×4\mathbb{R}^{32 \times 32 \times 4}), because we need to simultaneously denoise 768 dimensions instead of 4
1.going from 32x32 → 16x16 is a function of the transformer depth, can add layers + more attention blocks
2.going from learning a 4-parameter velocity field to a 768-parameter field is much harder
3.We must be able to decode images that are slightly off distribution from DINO feature representations

learned perceptual image patch similarity (LPIPS) - comparing deep feature representations extracted from various intermediate layers of a pretrained convolutional neural network

adversarial loss - GAN discriminator

-decoder is trained with L1, LPIPS, and adversarial loss
z=E(x),x^=D(z),Lrec=ωLLPIPS(x^,x)+L1(x^,x)+ωGλGAN(x^,x)z = E(x), \hat{x}=D(z),\\\mathcal{L}_{rec}= \omega_L \text{LPIPS}(\hat{x}, x) + \text{L1}(\hat{x}, x) + \omega_G \lambda \text{GAN}(\hat{x}, x)
-noise-augmented decoding - since the DINO outputs are discrete, but VAE encoder representations are continuous, add noise to decoding to train reconstruction that is robust
-The main DiT processes the high-dimensional RAE latent using a smaller hidden dimension (e.g. 384) to learn useful context/features.
-The wide DDT head receives both the DiT features and the original noisy latent.
-The DDT head predicts the full-dimensional velocity (e.g. all 1024 dimensions).
-This is cheaper than making every layer of the main DiT very wide.
1024-D latent+384-D DiT output→1024-D DDT denoising velo1024 \text{-D latent}+384\text{-D DiT output} \rightarrow 1024\text{-D DDT denoising velo}