PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Transformers with Representation Autoencoders

Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie

arXiv:2510.11690v1cs.CVcs.LG

TL;DR

The paper addresses the limitations of reconstruction-trained, low-dimensional VAE latents and their separation from semantically rich representation learning. It introduces Representation Autoencoders using frozen pretrained encoders and trained decoders, together with adaptations for high-dimensional diffusion. On ImageNet, the approach achieves 1.51 FID at 256×256 without guidance and 1.13 FID at both 256×256 and 512×512 with guidance.

  • Problem

    Conventional DiTs largely use SD-VAE latents that are low-capacity, reconstruction-trained, and disconnected from semantically structured representations.

  • Method

    RAEs pair frozen pretrained representation encoders with trained decoders, while DiT width, noise scheduling, and decoder training are adapted to high-dimensional latents.

  • Results

    1.51 FID at 256×256 without guidance and 1.13 FID at both 256×256 and 512×512 with AutoGuidance are achieved on ImageNet.

  • Takeaways & Limitations

    RAEs provide a semantically rich latent foundation for diffusion transformers and enable efficient, strong image generation without auxiliary representation alignment losses.

  • Takeaways & Limitations

    FID comparisons depend on class-balanced sampling, which the authors observe produces around 0.1 lower FID than uniform random sampling.

Abstract

from arXiv · show

Latent generative modeling, where a pretrained autoencoder maps pixels into a latent space for the diffusion process, has become the standard strategy for Diffusion Transformers (DiT); however, the autoencoder component has barely evolved. Most DiTs continue to rely on the original VAE encoder, which introduces several limitations: outdated backbones that compromise architectural simplicity, low-dimensional latent spaces that restrict information capacity, and weak representations that result from purely reconstruction-based training and ultimately limit generative quality. In this work, we explore replacing the VAE with pretrained representation encoders (e.g., DINO, SigLIP, MAE) paired with trained decoders, forming what we term Representation Autoencoders (RAEs). These models provide both high-quality reconstructions and semantically rich latent spaces, while allowing for a scalable transformer-based architecture. Since these latent spaces are typically high-dimensional, a key challenge is enabling diffusion transformers to operate effectively within them. We analyze the sources of this difficulty, propose theoretically motivated solutions, and validate them empirically. Our approach achieves faster convergence without auxiliary representation alignment losses. Using a DiT variant equipped with a lightweight, wide DDT head, we achieve strong image generation results on ImageNet: 1.51 FID at 256x256 (no guidance) and 1.13 at both 256x256 and 512x512 (with guidance). RAE offers clear advantages and should be the new default for diffusion transformer training.

1 INTRODUCTION

The paper replaces conventional reconstruction-trained VAEs with representation autoencoders that use pretrained semantic encoders and trained decoders. It addresses the resulting high-dimensional diffusion challenge with architectural and scheduling changes, achieving strong ImageNet generation results.

  • Motivation: SD-VAE uses heavy compression, low-capacity latents, and a reconstruction-only objective that limits semantic structure and generative performance.Its legacy convolutional design is also computationally inefficient.
  • Motivation: Pretrained encoders such as DINO, MAE, and SigLIP provide semantically structured features, but latent diffusion has largely remained in reconstruction-trained VAE spaces.The paper challenges the assumed incompatibility between semantic representations and faithful reconstruction or generation.
  • Approach: Representation Autoencoders pair a pretrained representation encoder with a trained decoder to produce semantically rich, structurally coherent latent spaces for diffusion.This design links semantic and generative modeling through a shared latent representation.
  • Approach: High-dimensional RAE latents require matching DiT width to token dimensionality, dimension-dependent noise shifts, and decoder noise augmentation.These changes address transformer capacity, noise scheduling, and decoder robustness.
  • Results: 1.51 FID at 256×256 without guidance and 1.13 FID at both 256×256 and 512×512 with AutoGuidance demonstrate strong ImageNet generation.The results position RAEs as an alternative to conventional VAEs for diffusion transformer training.

2 RELATED WORKS

Related work improves reconstruction or generation by incorporating semantic representations into VAE or DiT pipelines. These approaches generally align or augment existing systems rather than replacing the VAE latent space with a pretrained representation encoder and trained decoder.

  • Representation for Reconstruction: Recent VAE methods align latents with pretrained representations or add representation-inspired objectives to improve reconstruction and generation.Their heavily compressed, low-dimensional latents still limit reconstruction fidelity and representation quality.
  • Positioning: RAE differs by directly pairing frozen pretrained representation encoders with trained decoders rather than aligning a conventional VAE latent space.This distinction motivates treating the representation encoder itself as the diffusion latent-space foundation.
  • Representation for Generation: Recent generative methods use representation alignment, architectural decoupling, learnable tokens, or joint latent prediction to improve diffusion modeling.These methods retain representation learning as an auxiliary or integrated component of the generative pipeline.

3 HIGH FIDELITY RECONSTRUCTION FROM FROZEN ENCODERS

The paper shows that frozen semantic encoders can support high-fidelity reconstruction when paired with suitable ViT decoders. RAEs improve reconstruction and representation quality over SD-VAE while offering favorable decoder efficiency and scalable encoder choices.

  • Architecture: RAEs pair frozen DINOv2 or SigLIP2 encoders with ViT decoders and achieve reconstructions on par with or better than SD-VAE.This challenges the assumption that semantic encoders cannot preserve pixel-level detail.
  • Training: The decoder reconstructs pixels from encoder tokens and is trained with L1, LPIPS, and adversarial losses.The default ViT decoder uses the encoder patch size so reconstruction matches the input resolution.
  • Reconstruction Quality: RAEs with frozen encoders consistently outperform SD-VAE in reconstruction quality, with MAE-B/16 reaching an rFID of 0.16.The result directly challenges the assumption that representation encoders cannot recover pixel-level detail.
  • Scaling: Decoder scaling improves rFID from 0.58 with ViT-B to 0.49 with ViT-XL.ViT-B already outperforms SD-VAE while using 14× fewer GFLOPs, and ViT-XL reaches higher quality at one-third of SD-VAE’s cost.
  • Representation Quality: RAEs inherit the strong representations of their frozen encoders, whereas SD-VAE achieves only ∼8% ImageNet-1K linear-probing accuracy.Thus, the RAE latent space supports both reconstruction and semantic evaluation.

4 TAMING DIFFUSION TRANSFORMERS FOR RAE

RAE latents require diffusion transformers to be redesigned for high dimensionality. Matching model width to token dimension, shifting noise schedules, and augmenting decoder training enable faster and stronger generation.

  • Motivation: Standard DiT training struggles on RAE latents: DiT-S fails completely, while DiT-XL underperforms its SD-VAE counterpart.This motivates changes to transformer design, noise scheduling, and decoder training.
  • Scaling DiT Width to Match Token Dimensionality: Width drives overfitting more than depth: with token dimension n = 768, quality and loss improve sharply once width d ≥n, while doubling depth leaves artifacts and nonconvergent losses.Increasing width from 384 to 784 enables near-perfect reproduction; doubling depth from 12 to 24 at width 384 does not reach the same loss level.
  • Scaling DiT Width to Match Token Dimensionality: Gaussian noise expands the effective support of the latent distribution to full rank, explaining why diffusion capacity must scale with full data dimensionality.The paper contrasts this with models that can operate on low-intrinsic-dimensional data manifolds.
  • Scaling DiT Width to Match Token Dimensionality: Convergence occurs only when model width is at least as large as the RAE token dimension.For example, DiT-B with DINOv2-B converges, whereas DiT-S with DINOv2-B does not.
  • Dimension-Dependent Noise Schedule Shift: Dimension-dependent schedule shifting improves diffusion training in high-dimensional RAE latent spaces.The shift accounts for effective data dimension rather than relying only on spatial resolution.
  • Noise-Augmented Decoding: Noise-augmented decoding improves gFID but slightly worsens rFID by smoothing latent distributions and reducing fine details.The method addresses decoder out-of-distribution issues caused by noisy generated latents.
  • Results: The combined recipe reaches gFID 4.28 after 80 epochs and 2.39 after 720 epochs, with 47× and 16× training speedups over VAE and representation-alignment baselines, respectively.These results use a DiT-XL trained on RAE latents.

5 IMPROVING THE MODEL SCALABILITY WITH WIDE DIFFUSION HEAD

DiTDH adds a shallow, wide denoising head to standard DiT, increasing effective width for high-dimensional RAE latents without quadratic compute growth. Across efficiency, encoder scale, convergence, and ImageNet evaluation, DiTDH consistently delivers strong results.

  • 5 IMPROVING THE MODEL SCALABILITY WITH WIDE DIFFUSION HEAD: A shallow, wide DDT head increases DiT width for high-dimensional latents without incurring quadratic FLOP growth.The augmented architecture combines a base DiT with a dedicated denoising head.
  • 5 IMPROVING THE MODEL SCALABILITY WITH WIDE DIFFUSION HEAD: DiTDH-B uses approximately 40% of DiT-XL’s training FLOPs while outperforming it by a large margin.At comparable training budgets, DiTDH-XL reaches an FID of 2.16, nearly half that of DiT-XL.
  • 5 IMPROVING THE MODEL SCALABILITY WITH WIDE DIFFUSION HEAD: DiTDH consistently outperforms DiT across DINOv2-S, DINOv2-B, and DINOv2-L, with the advantage increasing for larger encoders.With DINOv2-L, DiTDH improves FID from 6.09 to 2.73.
  • 5.1 STATE-OF-THE-ART DIFFUSION TRANSFORMERS: DiTDH-XL surpasses several prior diffusion models around 5 × 10^10 GFLOPs and achieves the best overall FID by 5 × 10^11 GFLOPs using over 40× less compute.The comparison includes REPA-XL, MDTv2-XL, and SiT-XL.
  • 5.1 STATE-OF-THE-ART DIFFUSION TRANSFORMERS: DiTDH-XL achieves ImageNet FID scores of 1.51 at 256×256 without guidance and 1.13 at 256×256 and 512×512 with guidance.At 512×512, the 1.13 guided FID uses 400-epoch training and surpasses the previous best reported result of 1.25.
  • 5.1 STATE-OF-THE-ART DIFFUSION TRANSFORMERS: Qualitative samples show high semantic diversity and fine-grained detail consistent with the reported state-of-the-art FID.The samples are compared with ground-truth ImageNet samples.

6 DISCUSSIONS

Decoder upsampling separates output resolution from diffusion token count, allowing 512×512 generation from models trained at 256×256. Additional experiments show that DiTDH’s benefit depends on structured RAE latents rather than high dimensionality alone.

  • 6 DISCUSSIONS: Doubling image size in each dimension requires roughly four times as many tokens, creating a high-resolution scaling challenge.The approach shifts resolution scaling to the decoder by allowing decoder and encoder patch sizes to differ.
  • 6 DISCUSSIONS: Decoder upsampling reuses 256×256 diffusion models to produce 512×512 outputs without retraining.The method slightly increases rFID, achieves competitive gFID, and is 4× more efficient than quadrupling the token count.
  • 6 DISCUSSIONS: DiTDH-XL performs worse than DiT-XL on SD-VAE latents despite the extra diffusion-head computation.This indicates that the DDT head provides little benefit in low-dimensional latent spaces.
  • 6 DISCUSSIONS: On raw pixels with matched token dimensionality, DiTDH outperforms DiT, but both perform far worse than models trained on RAE latents.The comparison shows that high dimensionality alone is insufficient; structured representations also matter.

7 CONCLUSION

The paper presents RAEs as a practical alternative to conventional VAE-based diffusion-transformer training. Its reported conclusion combines structured representation latents, dimension-aware diffusion design, and strong ImageNet generation results.

  • 7 CONCLUSION: Pixel diffusion has much worse FID than diffusion on DINOv2-B latents.
  • 7 CONCLUSION: RAEs pair frozen pretrained representation encoders with lightweight trained decoders for reconstruction and generation.The method adds width matching, a dimension-dependent noise-schedule shift, decoder noise augmentation, and the DiTDH head.
  • 7 CONCLUSION: RAE-based DiTDH-XL achieves FID 1.51 at 256×256 without guidance and 1.13 at both 256×256 and 512×512 with guidance.

8 ACKNOWLEGMENTS

The supplied passages cover related work, theoretical assumptions, model definitions, and encoder normalization details surrounding the paper’s RAE framework. They describe how semantic encoders, diffusion models, and dimensionality constraints motivate the proposed design.

  • Related Work: Related work distinguishes semantic representation encoders, compressed image tokenizers, and diffusion or autoregressive generative models.Representation encoders learn semantically rich features, while quantized tokenizers retain an information-capacity bottleneck.
  • Theoretical Analysis: The lower-bound analysis assumes Gaussian noise, continuous interpolation xt = (1 − t)x + tε, and centered data in parts of the derivation.The proof uses projection properties and covariance eigenvalues to characterize approximation error.
  • Theoretical Analysis: The theoretical setup models standard DiT blocks whose width d is smaller than the representation token dimension n, with linear input and output projections.The function family is expressed as g(xt, t) = Bf(Axt, t).
  • Theoretical Analysis: When d ≥ n, the projection lower bound becomes trivial because the projected energy can equal the full energy.
  • Encoder Normalization: The encoder keeps patch tokens while discarding [CLS] or [REG] tokens, then independently layer-normalizes each token across channels.The affine parameters of the encoder’s final layer normalization are canceled without changing representation quality because this is a linear transformation.
  • Encoder Normalization: DINOv2 with Registers is interpolated to 224×224 and decoded with patch size 16 to produce 256 tokens while reconstructing 256×256 images.

C.2 DECODER TRAINING DETAILS

The decoder reconstructs encoder token embeddings into images at the encoder’s spatial resolution, using reconstruction and adversarial objectives with several stabilization choices. Visual examples show satisfactory reconstruction fidelity across RAEs and SD-VAE.

  • Decoder Architecture: RAE decoders map frozen encoder token embeddings back to pixel space using the encoder’s patch size and spatial resolution.A learnable [CLS] token is prepended and discarded after decoding.
  • Discriminator Architecture: Dino-S/8 is used as the frozen discriminator backbone because it stabilizes training and avoids adversarial patches.Virtual batch normalization is replaced with standard batch normalization.
  • Losses: Decoder training combines L1, LPIPS, and adversarial losses, with adaptive weighting to balance reconstruction and GAN objectives.The reconstruction loss uses ωL = 1 and ωG = 0.75.
  • Augmentations: Training resizes images to 384 × 384, randomly crops them to 256 × 256, and applies differentiable augmentations before discrimination.
  • Visualizations: RAE visualizations show satisfactory reconstruction fidelity across the evaluated representation encoders.Figure 8 compares DINOv2-B, SigLIP2-B, MAE-B, and SD-VAE reconstructions against the input image.

D DIFFUSION MODEL IMPLEMENTATION

Diffusion experiments use ImageNet at 256 × 256 and 512 × 512 with LightningDiT-based models, matched token sequence lengths, and standardized sampling and evaluation protocols. Ablations identify noise-augmented decoding and a wide, shallow DDT head as important design choices.

  • Datasets: Diffusion training primarily uses ImageNet-1K at 256 × 256, while 512-resolution models are trained directly on 512 × 512 images.
  • Compute: RAE, VAE, and pixel experiments all process 256 tokens, keeping DiT backbone computational cost identical apart from patchification.RAE uses patch size 1, VAE uses 2, and pixel inputs use 16; patchification contributes less than 1% of total GFLOPs.
  • Optimization: DiTDH training uses a modified optimization schedule because the standard LightningDiT recipe causes later loss spikes and slow early EMA convergence.
  • FID Evaluation: Class-balanced FID evaluation samples 50 images per class for 50,000 total images, and all methods improve consistently under this sampling strategy.The authors also evaluate their method with random sampling.
  • FID Evaluation: The ImageNet training set contains approximately 732–1,300 samples per class, and FID absolute values become less meaningful as generation fidelity improves.The authors note that sampling details matter more as FID approaches lower ranges.
  • Encoder Ablations: DINOv2-B achieves the best overall performance, while MAE’s lower reconstruction FID does not translate into better generation.MAE generation FID is substantially worse, motivating DINOv2-B as the default encoder.
  • Noise-Augmented Decoding: Noise-augmented decoder training improves generative FID across encoders while mildly worsening reconstruction FID.For MAE-B, gFID improves from 16.14 to 8.38 with noisy training; the benefit also holds across encoder sizes.
  • DDT Head: A 2-layer, 2048-dimensional DDT head outperforms deeper or narrower alternatives, indicating that wide and shallow heads are more effective for denoising.A 4-layer, 2048-dimensional head does not improve over the 2-layer version despite double the GFLOPs.

H ADDITIONAL SCALING RESULTS

Additional results examine scaling, guidance, flow-based training, and unconditional generation. Larger DiTDH models converge faster and lower final loss, while AutoGuidance is preferred over unguided CFG for the reported RAE results.

  • Scaling: Increasing DiTDH computational capacity produces faster convergence and a lower final training loss.Figure 9 reports the smoothed training loss on DINOv2-B.
  • Guidance: AutoGuidance uses a weaker diffusion model to guide a stronger model and is easier to tune than CFG with interval.The smallest DiTDH-S variant and typically an early checkpoint serve as the guiding model.
  • Guidance: CFG without an interval does not improve FID, whereas interval-tuned CFG is competitive but remains inferior to AutoGuidance on the final model.
  • Flow-Based Models: Flow-based models learn to reverse a noising process, generating samples by integrating an ODE or SDE from noise toward clean data.The described linear interpolation is x_t = (1 − t)x + tε, with t ranging from 0 to 1.
  • Flow-Based Models: The velocity and score functions are equivalent up to a constant factor, so training need estimate only one of them.
  • Unconditional Generation: Unconditional DiTDH-XL trained on RAE latents substantially outperforms DiT-XL trained on VAE latents and is competitive with RCG.The approach is described as simpler than RCG because it does not require two-stage generation.

M VISUAL RESULTS

Visual results include an unconditional ImageNet comparison and uncurated 512 × 512 DiTDH-XL samples generated with DINOv2-B and AutoGuidance scale 1.5.

  • Unconditional Generation: Table 18 compares unconditional generation on ImageNet at 256 × 256.
  • 512 × 512 Samples: The showcased 512 × 512 samples come from DiTDH-XL using DINOv2-B with AutoGuidance scale 1.5.
  • Class-Conditional Samples: Figures 10–15 present uncurated 512 × 512 samples for golden retriever, husky, cliff, macaw, arctic fox, and balloon labels.
Loading 2510.11690v1…