Diffusion Transformers with Representation Autoencoders

Overview

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).

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 / 1.13 at 256x256 and 512x512 (with guidance). RAE offers clear advantages and should be the new default for diffusion transformer training.

As with training VAEs, there are two questions to answer:

• How well can RAE reconstruct the input image?
• How well can diffusion transformers operate within the RAE latent space?

High Fidelity Reconstruction from RAE

We challenge the common assumption that pretrained representation encoders, such as DINOv2 and SigLIP2, are unsuitable for the reconstruction task because they “emphasize high-level semantics while downplaying low-level details” . We show that, with a properly trained decoder, frozen SEM can in fact serve as strong encoders for the diffusion latent space. Our Representation Autoencoders (RAE) pair frozen, pretrained SEM with a ViT decoder, yielding reconstructions on par with—or even better than—SD-VAE. More importantly, RAEs alleviate the fundamental limitations of VAEs , whose heavily compressed latent space (e.g., SD-VAE maps 256² images to 32²×4 latents) restricts reconstruction fidelity and, more importantly, representation quality.

Reconstruction, scaling, and representation.

As shown below, RAEs achieve consistently better reconstruction quality (rFID) than SD-VAE. For instance, RAE with MAE-B/16 reaches an rFID of 0.16, clearly outperforming SD-VAE and challenging the assumption that representation encoders cannot recover pixel-level detail.

We next study the scaling behavior of both encoders and decoders. As shown in Table 1c, reconstruction quality remains stable across DINOv2-S, B, and L, indicating that even small representation encoder models preserve sufficient low-level detail for decoding. On the decoder side (Table 1b), increasing capacity consistently improves rFID: from 0.58 with ViT-B to 0.49 with ViT-XL. Importantly, ViT-B already outperforms SD-VAE while being 14× more efficient in GFLOPs, and ViT-XL further improves quality at only one-third of SD-VAE’s cost. We also evaluate representation quality via linear probing on ImageNet-1K in Table 1d. Because RAEs use frozen pretrained encoders, they directly inherit the representations of the underlying encoders. Since RAEs use frozen pretrained encoders, they retain the strong representations of the underlying representation encoders. In contrast, SD-VAE achieves only approximately 8% accuracy.

Taming Diffusion Transformers for RAE

With RAE demonstrating good reconstruction quality, we now proceed to investigate the diffusability of its latent space; that is, how easily its latent distribution can be modeled by a diffusion model, and how good the generation performance can be. Empirically, despite the superior reconstruction quality of MAE encoders, we found that DINOv2 produces the strongest generation results, so we adopt it as the default, unless otherwise specified.

Following standard practice, we adopt the flow matching objective to train the diffusion model, for which we use LightningDiT, a variant of DiT, as our model backbone.

We adopt a patch size of 1, which results in a token length of 256 for all RAEs on 256x256 images, matching the sequence length used by VAE-based DiTs. We note that, Since the computational cost of DiT depends primarily on the sequence length, using RAE latents on DiTs will not incur additional computational cost compared to using VAE latents.

Scaling DiT Width to Match Token Dimensionality

To better understand the training dynamics of Diffusion Transformers in RAE latent space, we construct a simplified experiment: we randomly select a single image, encode it by RAE, and test whether the diffusion model can reconstruct it. Starting from a DiT-S, we first vary the model width while fixing depth. We fix the RAE encoder to DINOv2-B with token dimension of 768.

As shown in the following figure, sample quality is poor when the model width $d < $ token dimension $n=768$, but improves sharply and reproduces the input almost perfectly once $d >= n$. Training losses exhibit the same trend, converging only when $d >= n$. One might suspect that this improvement still arises from the larger model capacity, but again as shown in the following figure, even when doubling the depth from 12 to 24, the generated images remain artifact-heavy, and the training losses fail to converge to similar level of $d = 768$.

Together, the results indicate that for generation in RAE's latent space to succeed, the diffusion model's width must match or exceed the RAE's token dimension. This appears to contradict the common belief that data usually has low intrinsic dimension and thus allowing generative models like GAN to operate effectively without scaling width to the full data dimension . We note that this is due to the nature of diffusion models, where the injected Gaussian noise during training extends the data distribution’s support to the entire space, thereby "diffusing" the data manifold into a full-rank one, requiring model capacity that scales with the full data dimensionality.

As shown above, convergence occurs only when the model width is at least as large as the token dimension (e.g., DiT-B with DINOv2-B), while the loss fails to converge otherwise (e.g., DiT-S with DINOv2-B).

Dimension-Dependent Noise Schedule Shift

Many prior works , , , have observed that, for inputs $\mathbf{z} \in \mathbb{R}^{C\times H\times W}$, increasing the spatial resolution ($H \times W$) reduces information corruption at the same noise level, impairing diffusion training. These findings, however, are based mainly on pixel- or VAE-encoded inputs with few channels (e.g., $C \leq 16$). In practice, the Gaussian noise is applied to both spatial and channel dimensions; as the number of channels increases, the effective “resolution” per token also grows, reducing information corruption further. We therefore argue that proposed resolution-dependent strategies in these prior works should be generalized to the \textit{effective data dimension}, defined as the number of tokens times their dimensionality.

This this yields significant performance gains, showing its importance for training diffusion models in the high-dimensional RAE latent space.

Noise-Augmented Decoding

Unlike VAEs, whose latents follow a continuous distribution $\mathcal{N}(\mu, \sigma^2\mathbf{I})$, the RAE decoder $D$ is trained to reconstruct images from a discrete latent distribution $p(\mathbf{z}) = \sum_i \delta(\mathbf{z} - \mathbf{z}_i)$, where ${\mathbf{z}_i}$ are encoder outputs from the training set. At inference, diffusion-generated latents may deviate slightly from this distribution due to imperfect training or sampling, leading to out-of-distribution artifacts and degraded sample quality. To address this, following prior work in Normalizing Flows, , , we add Gaussian noise $\mathbf{n} \sim \mathcal{N}(0, \sigma^2\mathbf{I})$ during decoder training. Instead of decoding from the clean distribution $p(\mathbf{z})$, we train on its smoothed variant $p_{\mathbf{n}}(\mathbf{z}) = \int p(\mathbf{z} - \mathbf{n})\mathcal{N}(0, \sigma^2\mathbf{I})(\mathbf{n})\mathrm{d}\mathbf{n}$, improving generalization to the denser latent space of diffusion models. We further sample $\sigma$ from $|\mathcal{N}(0, \tau^2)|$ to regularize training and enhance robustness.

Combining all the above techniques, we train a DiT-XL model on RAE latents, which achieves a gFID of 4.28 (Right) after only 80 epochs and 2.39 after 720 epochs. With the same model size, this not only surpasses prior diffusion baselines (e.g., SiT-XL) trained on VAE latents—achieving a 47× training speedup—but also outpaces recent representation-alignment methods (e.g., REPA-XL) with a 16× faster convergence.

In the following sections, we investigate ways to make RAE generation more efficient and effective, pushing it toward state-of-the-art performance.

Improving the Model Scalability with Wide Diffusion Head

As discussed in Section 2, within the standard DiT framework, handling higher-dimensional RAE latents requires scaling up the width of the entire backbone, which quickly becomes computationally expensive. To overcome this limitation, we draw inspiration from DDT and introduce the DDT head—a shallow yet wide transformer module dedicated to denoising. By attaching this head to a standard DiT, we effectively increase model width without incurring quadratic growth in FLOPs. We refer to this augmented architecture as DiT^DH.

Wide DDT Head. Formally, a DiT^DH model consists of a base DiT \(M\) and an additional wide, shallow transformer head \(H\). Given a noisy input \(x_t\), timestep \(t\), and an optional class label \(y\), the combined model predicts the velocity \(v_t\) as

DiT^DH converges faster than DiT. We train a series of DiT^DH models with varying backbone sizes (DiT^DH-S, B, L, and XL) on RAE latents. We use a 2-layer, 2048-dim DiT^DH head for all models. Performance is compared against the standard DiT-XL baseline. DiT^DH is substantially more FLOP-efficient than DiT. For example, DiT^DH-B requires only \( \sim 40\% \) of the training FLOPs yet outperforms DiT-XL by a large margin; when scaled to DiT^DH-XL under a comparable training budget, it achieves an FID of 2.16—nearly half that of DiT-XL.

Convergence Comparison. We compare the convergence behavior of DiT^DH-XL with previous state-of-the-art diffusion models , , , , and in terms of FID without guidance. We show the convergence curve of DiT^DH-XL with training epochs and GFLOPs, while baseline models are plotted at their reported final performance. DiT^DH-XL already surpasses REPA-XL, MDTv2-XL, and SiT-XL around \(5 \times 10^{10}\) GFLOPs, and by \(5 \times 10^{11}\) GFLOPs it achieves the best FID overall, requiring over 40× less compute.

Scaling Behavior. We compare DiT^DH with recent methods across different model scales. Increasing the size of DiT^DH consistently improves FID performance. The smallest model, DiT^DH-S, achieves a competitive FID of 6.07—already outperforming the much larger REPA-XL. Scaling up to DiT^DH-B yields a substantial improvement from 6.07 to 3.38, surpassing all prior works of similar or even larger scale. The performance continues to improve with DiT^DH-XL, reaching a new state-of-the-art FID of 2.16 at 80 training epochs.

Performance. we provide a quantitative comparison between DiT^DH-XL, our most performant model, with recent state-of-the-art diffusion models on ImageNet 256×256 and 512×512 in Table 1 and Table 2. Our method outperforms all prior diffusion models by a large margin, setting new state-of-the-art FID scores of 1.51 without guidance and 1.13 with guidance at 256×256. On 512×512, with 400-epoch training, DiT^DH-XL achieves an FID of 1.13 with guidance, surpassing the previous best performance achieved by EDM-2 (1.25).

Discussions

How can RAE extend to high-resolution synthesis efficiently?

A central challenge in generating high-resolution images is that resolution scales with the number of tokens: doubling image size in each dimension requires roughly four times as many tokens. To address this, we let the decoder handle resolution scaling by allowing its patch size \(p_d\) to differ from the encoder patch size \(p_e\). When \(p_d = p_e\), the output matches the input resolution; setting \(p_d = 2p_e\) produces a 2× upsampled image, reconstructing a 512×512 image from the same tokens used at 256×256.

Since the decoder is decoupled from both the encoder and the diffusion process, we can reuse diffusion models trained at 256×256 resolution, simply swapping in an upsampling decoder to produce 512×512 outputs without retraining. This approach slightly increases both rFID and gFID, but being 4× more efficient than quadrupling the number of tokens.

Does DiT^DH work without RAE?

In this work, we propose and study RAE and DiT^DH. In Section 2, we showed that RAE combined with DiT already brings substantial benefits, even without the additional head. Here, we turn the question around: can DiT^DH still provide improvements without the latent space of RAE?

To investigate, we train both DiT-XL and DiT^DH-XL on SD-VAE latents with a patch size of 2, alongside DINOv2-B for comparison, for 80 epochs, and report unguided FID. As shown on the left, DiT^DH performs even worse than DiT on SD-VAE, despite the additional computation introduced by the diffusion head. This indicates that the head provides little benefit in low-dimensional latent spaces, and its primary strength arises in high-dimensional diffusion tasks introduced by RAE.

The role of structured representation in high-dimensional diffusion

DiT^DH achieves strong performance when paired with the high-dimensional latent space of RAE. This raises a key question: is the structured representation of RAE essential, or would DiT^DH work equally well on unstructured high-dimensional inputs such as raw pixels?

To evaluate this, we train DiT-XL and DiT^DH-XL directly on raw pixels. For 256×256 images with a patch size of 16, the resulting DiT input token dimensionality is 16×16×3 = 768, matching that of the DINOv2-B latents. We report unguided FID after 80 epochs. As shown on the right, DiT^DH outperforms DiT on pixels, but both models perform far worse than their counterparts trained on RAE latents. These results demonstrate that high dimensionality alone is not sufficient—the structured representation provided by RAE is crucial for achieving strong performance gains.

Conclusion

In this work, we challenge the belief that pretrained representation encoders are too high-dimensional and too semantic for reconstruction or generation. We show that a frozen representation encoder, paired with a lightweight trained decoder, forms an effective Representation Autoencoder (RAE). On this latent space, we train Diffusion Transformers in a stable and efficient way with three added components: (1) match DiT width to the encoder token dimensionality, (2) apply a dimension-dependent shift to the noise schedule, and (3) add decoder noise augmentation so the decoder handles diffusion outputs. We also introduce DiT^DH, a shallow-but-wide diffusion transformer head that increases width without quadratic compute. Empirically, RAEs enable strong visual generation: on ImageNet, our RAE-based DiT^DH-XL achieves an FID of 1.51 at 256×256 (no guidance) and 1.13 / 1.13 at 256×256 and 512×512 (with guidance). We believe RAE latents serve as a strong candidate for training diffusion transformers efficiently and robustly in future generative modeling research.