ArXiv: 2510.11690
🎯 Pitch
Frozen semantic encoders like DINOv2 paired with a trained ViT decoder can reconstruct images faithfully from high-dimensional latents—overturning the belief that semantic features sacrifice reconstruction quality. These Representation Autoencoders (RAEs) make DiTs converge up to 47× faster than prior work, achieving a record 1.13 FID on ImageNet 512×512 by simply widening the diffusion head to match the token dimension.
1. Executive Summary
This paper introduces Representation Autoencoders (RAEs), a new class of autoencoders that replace the standard VAE encoder with a frozen, pretrained representation encoder—such as DINOv2, SigLIP2, or MAE—paired with a trained ViT decoder, establishing that such semantically rich, high-dimensional latents are not only viable for reconstruction but yield substantially faster convergence and better generation quality in diffusion transformers than the conventional SD-VAE. The authors identify three technical obstacles that prevent standard DiTs from training effectively in these high-dimensional spaces—insufficient model width relative to token dimensionality, suboptimal noise scheduling that ignores channel dimension, and decoder fragility to OOD diffusion outputs—and propose targeted solutions for each: matching DiT width to the token dimension (e.g., using DiT-XL with DINOv2-B's 768-dim latents), applying a dimension-dependent schedule shift, and noise-augmented decoder training. They further introduce DiT, a DiT variant with a lightweight, wide DDT head that increases model width without quadratic FLOPs growth, achieving state-of-the-art ImageNet generation results: 1.51 FID at 256×256 without guidance and 1.13 at both 256×256 and 512×512 with AutoGuidance, outperforming all prior VAE-based methods while converging up to 47× faster than SiT-XL and 16× faster than REPA-XL—establishing that the benefits are specific to the combination of high-dimensional semantic latents and the wide-head architecture, as DiT degrades performance on low-dimensional VAE latents, and pixel-space diffusion with matched dimensionality performs far worse than RAE-based training.
2. Context and Motivation
The Core Problem: The Autoencoder in Latent Diffusion Has Barely Evolved
The fundamental question this paper tackles is: why does the autoencoder component in latent diffusion models remain frozen in time while everything else advances? Since the emergence of Latent Diffusion Models (Rombach et al., 2022), the standard strategy for high-quality image generation has been a two-stage pipeline: first, an autoencoder compresses pixels into a low-dimensional latent space; second, a diffusion model learns to generate within that latent space. The diffusion backbone has seen rapid architectural innovation—from U-Nets to Diffusion Transformers (DiT) (Peebles & Xie, 2023), from standard DiT to LightningDiT (Yao et al., 2025), and from purely generative training to representation-aligned training (REPA; Yu et al., 2025). Yet the autoencoder defining the latent space itself has remained the SD-VAE (Rombach et al., 2022), a convolutional architecture trained with a reconstruction-only objective that maps images to latents.
This is a significant gap for several reasons the authors surface throughout Section 1 and Section 3:
Architectural stagnation. The SD-VAE relies on a legacy convolutional design with aggressive downsampling and upsampling—a U-Net style encoder-decoder that is computationally expensive and architecturally misaligned with the transformer-based diffusion backbones that now dominate. As shown in Figure 2, the SD-VAE encoder requires ~135 GFLOPs and the decoder ~310 GFLOPs for a single image, making the autoencoder a substantial fraction of the total compute budget despite being an auxiliary component.
Information capacity bottleneck. The SD-VAE's aggressive compression—from pixels to latents—represents a ~48× reduction in dimensionality. While this compression was originally motivated by computational efficiency (diffusing in lower dimensions is cheaper), it fundamentally restricts how much information the latent space can carry. The authors argue that this low capacity "captures local appearance but lacks global semantic structure crucial for generalization and generative performance" (Section 1). Song et al. (2025) is cited in support of this claim, showing that VAE latents are impoverished representations that limit what downstream diffusion models can achieve.
Purely reconstruction-based training. SD-VAE is trained solely to reconstruct pixels, with no mechanism for capturing semantic structure. The consequence is a latent space optimized for faithful pixel-level recovery but semantically unstructured. The paper demonstrates this starkly in Table 1d: a linear probe on SD-VAE latents achieves only ~8% ImageNet-1K accuracy, while DINOv2-B achieves 84.5%. This means the diffusion model must learn semantic organization from scratch during training—a significant burden that slows convergence and limits final quality.
The Missed Opportunity: Representation Learning Has Transformed, But Latent Diffusion Hasn't Benefited
In parallel with generative modeling, the field of visual representation learning has undergone a revolution. Self-supervised and multimodal encoders—DINOv2 (Oquab et al., 2023), SigLIP2 (Tschannen et al., 2025), MAE (He et al., 2021)—learn semantically structured features through diverse pretraining objectives that go far beyond reconstruction. These encoders produce representations that generalize across tasks and scales, capture high-level semantic properties, and have been shown to be powerful foundation models for visual understanding.
Yet latent diffusion remains isolated from this progress. The authors identify this as a missed opportunity: if representation encoders produce richer, more structured latent spaces than the SD-VAE, why not simply diffuse in those spaces? The answer, they argue, lies in two long-standing assumptions that have prevented the field from exploring this direction.
Two Assumptions That Have Blocked Progress—And Why They May Be Wrong
The paper is motivated by what it identifies as two widely held but potentially incorrect beliefs:
Assumption 1: Semantic encoders cannot reconstruct faithfully. It is commonly stated that encoders trained for semantics (e.g., DINO, SigLIP) "focus on high-level information and can only reconstruct an image with high-level semantic similarities" (Yu et al., 2024b, quoted in Section 1). The intuition is straightforward: objectives like contrastive learning or self-distillation prioritize invariance to low-level details (color, texture, exact pixel positions) in favor of high-level semantic invariance. As a result, the latent space should discard the precise information needed for pixel-level reconstruction. This assumption has driven the community to treat reconstruction and representation as competing objectives that require careful trade-offs—hence the continued reliance on VAEs explicitly trained for reconstruction.
The paper challenges this head-on in Section 3. With a properly trained decoder (a ViT-based architecture trained with L1, LPIPS, and adversarial losses), frozen representation encoders can achieve reconstruction quality comparable to or better than SD-VAE. Table 1a shows that DINOv2-B achieves an rFID of 0.49, SigLIP2-B achieves 0.53, and MAE-B achieves 0.16—all lower (better) than SD-VAE's 0.62. This directly contradicts the assumption that semantic training destroys reconstruction capability.
Assumption 2: Diffusion models perform poorly in high-dimensional latent spaces. The SD-VAE produces -dimensional latents (when flattened). Representation encoders like DINOv2-B produce -dimensional latents—a ~48× increase in dimensionality. It is widely believed that diffusion models struggle in such high-dimensional spaces (Skorokhodov et al., 2025; Yao et al., 2025; Esser et al., 2024; Liu et al., 2024, all cited in Section 1). The intuition is that adding Gaussian noise to a high-dimensional vector corrupts less information per dimension, making the denoising task harder because the model must learn finer distinctions. This has led practitioners to strongly favor low-dimensional VAE latents.
The paper argues this assumption is also flawed—or at least incomplete. With the right architectural adjustments (matching model width to token dimension, Section 4.1) and training recipe modifications (dimension-dependent noise scheduling, Section 4.2), diffusion training in high-dimensional spaces is not only stable but advantageous. The higher dimensionality provides richer structure that actually accelerates convergence and improves final quality, provided the model has sufficient capacity to exploit it.
Where Prior Approaches Fall Short
The paper identifies specific limitations in existing work along several axes:
VAE-based latent diffusion is hitting diminishing returns. State-of-the-art DiT variants—SiT (Ma et al., 2024), REPA (Yu et al., 2025), DDT (Wang et al., 2025c), MDTv2 (Gao et al., 2023), VA-VAE (Yao et al., 2025)—all operate on SD-VAE latents. While these methods improve the diffusion backbone (through better interpolants, representation alignment, or architectural modifications), they are fundamentally bottlenecked by the latent space they inherit. The paper's Figure 1 illustrates this: VAE-based methods cluster in a particular FID regime, and progress has been incremental, with methods like REPA-E achieving 1.70 FID after 800 epochs at 256×256. The authors argue that continuing to optimize the diffusion model while keeping the VAE fixed is a form of diminishing returns—the autoencoder is the limiting factor, not the diffusion backbone.
Representation-aligned methods add complexity without solving the root cause. Recent work attempts to bridge the gap between representation learning and generation through auxiliary alignment losses. REPA (Yu et al., 2025) aligns the intermediate features of the diffusion model with a pretrained representation encoder during training. DDT (Wang et al., 2025c) decouples the DiT into an encoder-decoder and applies REPA loss to the encoder output. REG (Wu et al., 2025) introduces a learnable token explicitly aligned with a representation encoder's output. ReDi (Kouzelis et al., 2025b) generates both VAE latents and PCA components of DINOv2 features jointly. These methods demonstrate that injecting representation knowledge improves convergence, but they do so indirectly—the diffusion model still operates in the VAE's latent space, and the representation signal enters through auxiliary losses that require careful balancing, additional training stages, and tuning complexity.
The fundamental limitation of these approaches is that they treat the VAE latent space as fixed and try to enrich it from the outside. They do not address the core issue: the latent space itself is impoverished. The authors frame RAEs as a more direct solution—rather than aligning a VAE-based diffusion model with a representation encoder through auxiliary losses, simply use the representation encoder as the latent space directly.
Recent attempts at representation-based autoencoding have been incomplete. The paper acknowledges related work that explores using representation encoders as autoencoders:
- VA-VAE (Yao et al., 2025) aligns VAE latents with a pretrained encoder but retains the VAE's compressive bottleneck, limiting both reconstruction fidelity and representation quality.
- MAETok (Chen et al., 2025a) and DC-AE 1.5 (Chen et al., 2025d) incorporate MAE-inspired objectives into VAE training, but still operate within the compressed VAE paradigm.
- l-DEtok (Yang et al., 2025) combines masking and denoising losses for VAE training but requires joint encoder-decoder training.
- UniLIP (Tang et al., 2025) uses a one-step convolutional decoder on top of InternViT features, achieving reconstruction quality surpassing SD-VAE. However, the authors note that UniLIP argues frozen representation encoders lack sufficient visual detail and requires additional large-scale fine-tuning of the pretrained ViT.
The key distinction the paper makes (expanded in Appendix A): prior work either modifies the VAE training objective or fine-tunes the representation encoder, operating under the assumption that frozen encoders cannot provide sufficient detail for reconstruction. The paper's counterclaim is that this assumption is wrong—frozen encoders, with a sufficiently capable decoder, can reconstruct faithfully and serve as a superior latent space for diffusion without any encoder modification.
How This Paper Positions Itself
The paper's positioning can be understood as a direct challenge to the latent diffusion status quo along two dimensions:
From compression to representation. The standard view treats the autoencoder as a compression mechanism—pixels in, compressed latents out, with reconstruction quality as the sole metric of success. The paper reframes this: the autoencoder should be a representation foundation, where the latent space inherits the semantic structure and generalization properties of pretrained encoders. The empirical demonstration that frozen DINOv2 + trained decoder reconstructs almost as well as SD-VAE (and better in some cases, e.g., MAE-B's 0.16 rFID vs. SD-VAE's 0.62) while providing dramatically richer representations (84.5% vs. 8% linear probe accuracy) is the evidence that compression and representation need not be in tension.
From auxiliary alignment to direct diffusion. Rather than training diffusion models on VAE latents and aligning them with representation encoders post-hoc (the REPA/DDT/REG paradigm), the paper argues for diffusing directly in the representation encoder's space. This eliminates auxiliary losses, reduces training complexity, and—as Section 4 demonstrates—accelerates convergence substantially (47× faster than SiT-XL, 16× faster than REPA-XL). The technical challenge is making diffusion work in high-dimensional spaces, which Sections 4.1–4.3 address through the width-matching principle, dimension-dependent noise scheduling, and noise-augmented decoding.
From universal to architecture-dependent. A subtle but important aspect of the paper's positioning is that the benefits of RAE are not universal—they depend on proper architectural alignment between the latent space and the diffusion model. The DiT architecture is specifically designed to address the width bottleneck that arises when the token dimension exceeds standard DiT widths. The fact that DiT degrades performance on low-dimensional VAE latents (Table 10: 11.70 FID vs. 7.13 for DiT-XL) while dramatically improving performance on RAE latents (2.16 vs. 4.28) demonstrates that the contribution is not simply "RAE is better" or "DiT is better"—it is the combination of high-dimensional semantic latents with an architecture designed to exploit them that yields the gains. This is a more nuanced claim than a simple method proposal.
The paper also explicitly positions itself against a third alternative: pixel-space diffusion. Section 6.3 asks whether the structured representation of RAE is essential, or whether DiT would work equally well on unstructured high-dimensional inputs. Training on raw pixels with matched dimensionality (patch size 16, producing -dim tokens) yields dramatically worse results (DiT-XL: 51.09 FID; DiT-XL: 30.56) compared to RAE latents (4.28 and 2.16 respectively). This demonstrates that high dimensionality alone is insufficient—the structured nature of representation encoder features is crucial.
3. Technical Approach
3.1 Reader Orientation
This paper builds a representation-based latent diffusion pipeline that replaces the conventional SD-VAE autoencoder with a frozen, pretrained semantic encoder (like DINOv2) paired with a trained ViT decoder—called a Representation Autoencoder (RAE)—and then trains diffusion transformers directly in this semantically rich, high-dimensional latent space. The core problem is that standard diffusion transformers fail catastrophically in high-dimensional latent spaces for reasons that were previously misunderstood, and the solution is a combination of three principled fixes to the diffusion recipe (width matching, dimension-dependent noise scheduling, noise-augmented decoding) plus a new architecture variant (DiT) that adds a shallow-but-wide denoising head to handle the increased dimensionality efficiently.
3.2 Big-Picture Architecture (Diagram in Words)
The system consists of five major components connected in a two-stage pipeline:
-
Representation Encoder (frozen, pretrained): Takes an input image
$x \in \mathbb{R}^{3 \times H \times W}$and produces token embeddings$z = E(x) \in \mathbb{R}^{N \times d}$, where$N = HW/p_e^2$is the number of patch tokens and$d$is the encoder's hidden dimension (e.g., 768 for DINOv2-B). This encoder is pretrained via self-supervised or multimodal objectives and never updated. The output tokens are rich semantic representations, not compressed pixels. -
ViT Decoder (trained from scratch): Takes the frozen encoder's token embeddings
$z$and maps them back to pixels$\hat{x} = D(z) \in \mathbb{R}^{3 \times H \times W}$. By default, the decoder uses the same patch size$p_d = p_e$, so the output resolution matches the input. The decoder is trained with a combination of L1 loss, LPIPS perceptual loss, and adversarial loss against a frozen DINO-S/8 discriminator. -
Diffusion Transformer Backbone (trained from scratch): A flow-matching diffusion model—by default LightningDiT—that learns to denoise samples in the RAE latent space. It takes noisy latent tokens
$z_t = (1-t)z + t\varepsilon$(where$\varepsilon \sim \mathcal{N}(0, I_n)$and$t \in [0, 1]$) and predicts the velocity$v(z_t, t) \approx \varepsilon - z$. The model uses a patch size of 1, meaning it processes all$N$encoder tokens directly as the input sequence. -
DDT Head (for DiT only): A shallow, wide transformer module attached to the output of the standard DiT backbone. It takes the DiT's intermediate features as conditioning and applies additional denoising capacity with a much wider hidden dimension (2048) but only 2 layers, increasing model width without the quadratic FLOPs cost of scaling the entire backbone.
-
Noise-Augmented Decoder Interface: At inference time, the diffusion model generates latents
$\tilde{z}$that may deviate from the clean training distribution. The noise-augmented decoder—trained on$z + n$where$n \sim \mathcal{N}(0, \sigma^2 I)$with$\sigma \sim |\mathcal{N}(0, \tau^2)|$—is robust to this distribution shift and can decode both clean and slightly noisy latents into high-quality images.
Information flows in two phases. Phase 1 (Autoencoder Training): Image $\rightarrow$ frozen encoder $\rightarrow$ latent tokens $\rightarrow$ decoder $\rightarrow$ reconstruction $\rightarrow$ compute L1 + LPIPS + adversarial losses $\rightarrow$ update decoder parameters only. Phase 2 (Diffusion Training): Image $\rightarrow$ frozen encoder $\rightarrow$ latent tokens $\rightarrow$ add noise $\rightarrow$ noisy latents $\rightarrow$ DiT/DiT backbone $\rightarrow$ predicted velocity $\rightarrow$ flow matching loss $\rightarrow$ update diffusion model parameters. At generation time: random noise $\rightarrow$ DiT/DiT $\rightarrow$ denoised latents $\rightarrow$ noise-augmented decoder $\rightarrow$ output image.
3.3 Roadmap for the Deep Dive
- First, the RAE architecture and decoder training recipe, because the decoder's reconstruction quality establishes the upper bound on what the diffusion model can generate. Understanding what the decoder needs (L1, LPIPS, adversarial training, encoder normalization) sets the foundation for everything downstream.
- Second, the core theoretical puzzle: why standard DiT training fails on RAE latents. This includes the single-image overfitting experiment and Theorem 1, which together show that model width must match or exceed token dimensionality. This is the most counterintuitive finding and the key to everything that follows.
- Third, the dimension-dependent noise schedule shift, which generalizes prior resolution-dependent scheduling to the full effective data dimension (tokens × channels). This is a relatively simple modification but yields large practical gains.
- Fourth, noise-augmented decoder training, which bridges the distribution gap between the decoder's training distribution (clean latents from the real dataset) and the diffusion model's output distribution (slightly noisy latents due to imperfect denoising).
- Fifth, the DiT architecture with its wide DDT head, including the design rationale (why shallow and wide), the scaling behavior across encoder sizes, and the critical interaction with latent space dimensionality.
- Sixth, the full training recipe and how all components combine, including optimization hyperparameters, the AutoGuidance mechanism, and the 512×512 upsampling strategy.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily a systems and architecture paper whose core idea is that pretrained representation encoders—despite being designed for semantic understanding rather than reconstruction—can serve as superior latent spaces for diffusion transformers, provided three obstacles are addressed: (1) model capacity must scale with token dimensionality, (2) noise scheduling must account for the full effective data dimension, and (3) the decoder must be made robust to the diffusion model's output distribution. The method involves no auxiliary representation alignment losses; instead, diffusion happens directly in the semantically structured space.
RAE Architecture and Decoder Training
Encoder.
Given an input image $x \in \mathbb{R}^{3 \times H \times W}$ (typically $256 \times 256$), a frozen pretrained representation encoder $E$ with patch size $p_e$ and hidden dimension $d$ produces $N = HW / p_e^2$ patch tokens, each of dimension $d$. For DINOv2-B, which the paper selects as the default encoder after empirical comparison (Section G.1), $p_e = 14$ and $d = 768$. The authors interpolate input images to $224 \times 224$ before encoding (because DINOv2 only provides $p_e = 14$ variants), which yields exactly $(224/14)^2 = 256$ tokens. Critically, they set the decoder patch size $p_d = 16$, so the decoder reconstructs $256 \times 256$ images from those 256 tokens—matching the token count of standard VAE-based DiTs (which process $32 \times 32 = 1024$ VAE latents with patch size 2, also yielding 256 tokens). This means that, despite the much higher per-token dimensionality (768 vs. 4 for SD-VAE), the sequence length—and therefore the computational cost of the DiT backbone—is identical. The paper explicitly states that "since the token count is fixed (determined by the patch size) and the channels are projected to the DiT hidden dimension in the first layer," the higher dimensionality "introduces effectively no extra compute or memory costs" (Section 1).
Encoder normalization. The authors discard any [CLS] or [REG] tokens produced by the encoder and keep only the patch tokens. They then apply layer normalization to each token independently, ensuring each token has zero mean and unit variance across channels. Since the representation encoders already apply layer normalization after their final transformer block (standard ViT architecture per Dosovitskiy et al., 2021), the authors only need to cancel the affine parameters of this existing layer norm—a linear transformation that does not affect representation quality. This normalization step is a practical detail that ensures the latent distribution fed to the diffusion model is well-conditioned.
Decoder architecture.
The decoder $D$ is a standard ViT (no special modifications beyond the architecture family) that takes the $N$ token embeddings from the frozen encoder, prepends a learnable [CLS] token (following the MAE design from He et al., 2021), and maps the patch tokens back to pixels. The decoder uses patch size $p_d$, which by default equals the encoder's patch size $p_e$, meaning one output patch per input token. For $p_d = p_e = 14$ (DINOv2-B), the decoder maps 256 tokens to $256 \times 256 \times 3$ images (since, as noted above, the input is interpolated to $224 \times 224$ but $p_d = 16$ is used to reconstruct $256 \times 256$). The [CLS] token is discarded after decoding. The authors study three decoder sizes—ViT-B, ViT-L, and ViT-XL—with ViT-XL as the default (Table 1b). Even the smallest decoder (ViT-B, 22.2 GFLOPs) already outperforms SD-VAE's decoder (310.4 GFLOPs) in rFID (0.58 vs. 0.62), representing a ~14× reduction in compute.
Decoder training losses.
The decoder is trained with three losses combined—this is a critical design choice explained in Appendix C. The encoder produces clean latents $z = E(x)$, and the decoder predicts $\hat{x} = D(z)$. The training objective is:
where $\text{LPIPS}(\hat{x}, x)$ is the Learned Perceptual Image Patch Similarity metric (Zhang et al., 2018), measuring perceptual distance in deep feature space; $\ell_1(\hat{x}, x)$ is the element-wise L1 distance between the reconstruction and the original; and $\mathcal{L}_{\text{GAN}}(\hat{x}, x)$ is the adversarial loss from a discriminator trained to distinguish real from reconstructed images. The coefficients are $\omega_L = 1.0$ and $\omega_G = 0.75$.
What this computes: The LPIPS term penalizes reconstructions that look perceptually different to a deep network, even if per-pixel L1 error is small; it captures high-level structure. The L1 term penalizes pixel-level discrepancies and anchors the reconstruction to exact input content, preventing the GAN loss from hallucinating plausible-but-incorrect textures. The GAN term, weighted by an adaptive factor $\lambda$, provides a realism signal that pushes the decoder to produce sharp, visually convincing images rather than blurry averages. The adaptive $\lambda$ is computed as:
where $\nabla_{\hat{x}} \mathcal{L}_{\text{rec}}$ is the gradient of the reconstruction losses (LPIPS + L1) with respect to the decoder output, and $\nabla_{\hat{x}} \mathcal{L}_{\text{GAN}}$ is the gradient of the GAN loss. The $\epsilon$ is a small constant preventing division by zero.
Why this form: The adaptive $\lambda$ ensures the reconstruction and adversarial gradients have comparable magnitudes, preventing either term from dominating training. Without this balancing, typical GAN + reconstruction training can oscillate—the GAN term can overpower the reconstruction signal, causing the decoder to ignore input content and hallucinate, or vice versa. The three-loss combination directly follows the VAE training literature (Esser et al., 2021; Rombach et al., 2022) but has not previously been applied to frozen representation encoders. The explicit LPIPS term is important because semantic encoders can discard pixel-level detail that L1 alone would heavily penalize; LPIPS provides a more forgiving perceptual distance that allows the decoder to recover missing detail through learned priors.
Discriminator architecture and training.
The discriminator largely follows the design of StyleGAN-T (Sauer et al., 2023) with some modifications. It uses a frozen DINO-S/8 (Caron et al., 2021) as the backbone—the authors found DINO-S/8 stabilizes training and avoids adversarial patches compared to DINO-S/16. All inputs to the discriminator are interpolated to $224 \times 224$ resolution. The discriminator uses standard batch normalization (removing virtual batch norm from StyleGAN-T). Training starts at epoch 6 (the decoder is trained with only L1 + LPIPS for the first 6 epochs), and the adversarial loss component starts at epoch 8. The discriminator is trained with a non-saturating GAN loss. Data augmentation includes differentiable augmentations (Zhao et al., 2020) applied before the discriminator, and input images are resized to $384 \times 384$ then randomly cropped to $256 \times 256$ for data augmentation during decoder training. Both decoder and discriminator use Adam optimizer with learning rate $2 \times 10^{-4}$ decaying to $2 \times 10^{-5}$ via cosine schedule, betas $(0.5, 0.9)$, weight decay 0.0, batch size 512, and 1 epoch of warmup. The decoder trains for 16 epochs; the discriminator for 10.
Design choice—frozen encoder:
The encoder is completely frozen throughout training. This is the central architectural claim of the paper: that pretrained semantic features, without any fine-tuning or adaptation, contain sufficient information for high-fidelity reconstruction. The paper explicitly contrasts this with UniLIP (Tang et al., 2025), which argues that frozen encoders lack visual detail and requires large-scale fine-tuning. The empirical evidence in Table 1 supports the frozen approach: DINOv2-B + ViT-XL decoder achieves 0.49 rFID vs. SD-VAE's 0.62, and even DINOv2-S (a small model with $d = 384$) achieves 0.52. The scaling analysis (Table 1c) shows rFID is stable across DINOv2-S, B, and L, indicating that even small representation encoders preserve sufficient low-level detail. This is a critical practical advantage: it means practitioners do not need to retrain or fine-tune large encoders for each new autoencoder configuration.
Design choice—ViT decoder rather than convolutional: The decoder uses a standard ViT architecture rather than the U-Net style convolutional decoder of SD-VAE. This architectural alignment between encoder and decoder (both are ViTs) enables the decoder to process tokenized representations directly—patch-to-patch mapping—without the need for reshaping, upsampling, or architectural bridging between transformer outputs and convolutional inputs. The GFLOPs comparison in Figure 2 makes the efficiency case concrete: the RAE encoder is ~6× cheaper than the SD-VAE encoder (22 GFLOPs vs. 135 GFLOPs), and the RAE decoder is ~3× cheaper (106 GFLOPs vs. 310 GFLOPs for ViT-XL) while achieving better reconstruction quality.
What the reconstruction results establish: Table 1 demonstrates three things simultaneously. First, multiple representation encoders (DINOv2, SigLIP2, MAE) from different pretraining paradigms all achieve better rFID than SD-VAE (Table 1a: 0.49, 0.53, 0.16 vs. 0.62). Second, reconstruction quality scales with decoder size—ViT-XL > ViT-L > ViT-B—and the scaling curve has not saturated (Table 1b). Third, representation quality (linear probe accuracy on ImageNet-1K) is dramatically higher: 84.5% for DINOv2-B vs. 8.0% for SD-VAE (Table 1d). This third point is crucial for what follows: the diffusion model will be trained in a latent space that is not just reconstructable but semantically structured, carrying category-level information in ways the VAE latent does not. The paper's central claim is that this semantic structure is the source of the accelerated convergence and improved generation quality.
Why Standard DiTs Fail: The Width Bottleneck
The empirical puzzle. When the authors first attempted to train a standard DiT directly on RAE latents (DINOv2-B, 768-dimensional tokens), the results were catastrophic: DiT-S achieved 215.76 gFID (essentially random), and even DiT-XL achieved only 23.08 gFID—far worse than the same architecture on SD-VAE latents (DiT-XL: 7.13 FID; Table 2). This is the central technical challenge the paper must solve. The authors decompose the failure into three hypotheses: (1) suboptimal transformer design for high-dimensional tokens, (2) suboptimal noise scheduling that ignores the channel dimension, and (3) decoder fragility to out-of-distribution diffusion outputs. They address each with a targeted solution.
The single-image overfitting experiment. To isolate the fundamental capacity issue from dataset complexity, the authors construct a minimal experiment: take a single ImageNet image, encode it with RAE, and attempt to overfit a diffusion model to reconstruct just that one image from noise. If a model with sufficient capacity should be able to memorize a single datapoint, failure to do so reveals a fundamental structural limitation rather than a training-difficulty issue.
Using a DiT-S backbone (12 layers, default width 384) on DINOv2-B latents (token dimension $n = 768$), the authors sweep the model width $d$ from 384 to 896 while keeping depth fixed at 12. The results are presented in Figure 3 (left): when $d = 384$ (model width less than token dimension), the loss plateaus around 0.5 and the generated reconstruction is unrecognizable noise. When $d$ increases to 512 and 640, the loss decreases but still fails to converge to near-zero. At $d = 768$ (matching the token dimension), the loss drops sharply to near-zero and the reconstruction perfectly reproduces the input. At $d = 896$, the behavior is similar.
To rule out the alternative explanation that it is simply total model capacity (parameters) that matters, the authors fix width at $d = 384$ and vary depth from 4 to 24 layers (Figure 3, right). Even with 24 layers (substantially more parameters than a 12-layer, 768-width model), the loss remains around 0.5 and reconstructions remain artifact-heavy. The loss does not converge. This confirms that the bottleneck is specifically width relative to token dimension, not total FLOPs or parameter count.
The general pattern.
Table 3 extends this to combinations of different model sizes and encoder sizes. With DiT-S (width 384), overfitting succeeds only on DINOv2-S (token dimension 384), achieving loss $3.6 \times 10^{-2}$. On DINOv2-B (768-dim) and DINOv2-L (1024-dim), the loss is $5.2 \times 10^{-1}$ and $6.5 \times 10^{-1}$ respectively—failure. DiT-B (width 768) succeeds on DINOv2-S and B but fails on DINOv2-L. DiT-L (width 1024) succeeds on all three, because its width meets or exceeds all token dimensions. The pattern is precisely: convergence requires $d \geq n$, where $d$ is the DiT's hidden dimension and $n$ is the token channel dimension.
Theoretical justification—Theorem 1.
The authors provide a formal statement that explains why width matters categorically, not just as a capacity measure. The theorem considers the function family $\mathcal{G}_d$ representing DiT models whose internal representations are constrained to dimension $d < n$, where $n$ is the token dimension (and the full data dimension is $N \cdot n$ for $N$ tokens—though the theorem operates per-dimension). The family is defined as:
where $A$ is the input projection (from $n$ dimensions to the DiT's $d$-dimensional hidden space), $B$ is the output projection (from $d$ back to $n$), and $f$ is a stack of DiT blocks operating in $d$ dimensions. The theorem states that for any $g \in \mathcal{G}_d$ with $d < n$:
where $\lambda_i$ are the eigenvalues of the covariance matrix of $W = \varepsilon - x$, and $x_t = (1-t)x + t\varepsilon$.
What this means in operational terms: The flow matching objective requires the model $g$ to predict the velocity $v = \varepsilon - x$, which is an $n$-dimensional vector for each token. Since $g$ compresses through a $d$-dimensional bottleneck ($A$ maps $n \to d$, $B$ maps $d \to n$), the output of $g$ can only span at most a $d$-dimensional subspace of the full $n$-dimensional target space. The irreducible error is the variance of the target $W$ in the $n-d$ dimensions that the model cannot represent. When $p(x)$ has full support (the data distribution is not degenerate to a lower-dimensional subspace), this irreducible error is non-zero, and the loss can never reach zero—even on a single training example.
Why this is different from common intuitions about manifold learning: It is widely believed that image data lies on a low-dimensional manifold (Pope et al., 2021), and generative models can operate effectively within that manifold without needing capacity equal to the full data dimension. The authors argue that this intuition breaks for diffusion models because of a specific mechanism: injecting Gaussian noise throughout training extends the data distribution's support to the entire $\mathbb{R}^n$ space. During training, the model sees targets $x_t = (1-t)x + t\varepsilon$ for all $t \in [0, 1]$. At $t = 1$, $x_t$ is pure noise. At intermediate $t$, the noisy latents fill the ambient space because Gaussian noise has full support. The model must therefore learn a function defined on the entire ambient space $\mathbb{R}^n$, not just the data manifold. The velocity field $v(x_t, t)$ is a well-defined function everywhere in $\mathbb{R}^n$ (it is the conditional expectation of $\varepsilon - x$ given the noisy observation), and it has variation in all $n$ dimensions. A model with $d < n$ internal dimension simply cannot represent this function faithfully.
Applying the bound to the single-image case: When $p(x) = \delta(x - x_0)$ (a single training image, a point mass), the random variable $W = \varepsilon - x$ follows $\mathcal{N}(-x_0, I_n)$. The covariance matrix of $W$ is simply the identity $I_n$, so all eigenvalues $\lambda_i = 1$. The lower bound becomes $\mathcal{L}(g, \theta) \geq \sum_{i=d+1}^n 1 = n - d$. Normalizing by $n$ (since the loss averages over all $n$ dimensions), the minimum achievable average loss is $(n-d)/n$. When $d = 384$ and $n = 768$, this predicts a minimum average loss of $(768-384)/768 = 0.5$—exactly matching the empirical plateau in Figure 3. This is a strong validation of the theoretical model.
What the full-image scenario implies: For the full ImageNet dataset, $p(x)$ is not a point mass; it is the distribution of RAE latents over all images. The covariance of $W = \varepsilon - x$ will have a more complex eigenvalue spectrum, but the same principle applies: dimensions where the data distribution has variance cannot be represented by a model whose internal dimension is smaller than the ambient dimension. The practical consequence is that the DiT's width must be at least as large as the RAE token dimension. For DINOv2-B with $d = 768$, this means using at least DiT-XL (width 1152) or larger. The paper defaults to DiT-XL for all RAE experiments using DINOv2-B.
Design choice—why not just always use a wider model? The obvious solution to the width bottleneck is simply to scale up the DiT width. The problem is computational: DiT's self-attention cost scales quadratically with width (since the hidden dimension determines both the sequence dimension for attention and the feed-forward dimensions). A DiT with width 2048 would be substantially more expensive than DiT-XL (width 1152). This motivates the DiT architecture (Section 5), which adds width only in a shallow head module, achieving the benefits of width without the quadratic cost throughout the entire backbone.
Dimension-Dependent Noise Schedule Shift
The problem with standard noise scheduling on high-dimensional tokens.
Prior work (Chen, 2023; Hoogeboom et al., 2023; Esser et al., 2024) established that for inputs $z \in \mathbb{R}^{C \times H \times W}$, increasing the spatial resolution $H \times W$ reduces information corruption at a given noise level. More pixels mean that Gaussian noise, which is independently sampled per dimension, corrupts a smaller fraction of the total information content—there are more "redundant" dimensions carrying overlapping information. This is why resolution-based noise schedule shifts were introduced: they effectively increase the noise level as resolution grows, compensating for the reduced per-dimension corruption.
However, the authors point out that prior work derived these shifts for pixel-space or VAE-latent inputs, where the channel dimension $C$ is small (typically 3 for pixels, 4–16 for VAE latents). In RAE, each token has $d = 768$ channels for DINOv2-B. Gaussian noise is applied to all $n = N \cdot d$ dimensions simultaneously (across both tokens and channels). The effective data dimension is therefore $n = N \cdot d$, not just $N$. If you have 256 tokens each of dimension 768, adding noise with the same per-dimension variance corrupts vastly less information than adding the same noise to 256 tokens each of dimension 4 (SD-VAE), because the high-dimensional tokens carry information content distributed across many more dimensions.
The solution: dimension-dependent shifting.
The authors adopt the shifting strategy from Esser et al. (2024). Given a base schedule $t_n \in [0, 1]$ designed for a reference dimension $n$, and an actual input with effective dimension $m$, the shifted timestep is:
where the scaling factor is $\alpha = \sqrt{m / n}$.
What this computes: For each timestep $t_n$ in the original schedule, the function maps it to a new timestep $t_m$ that accounts for the difference in dimensionality. When $m > n$ (RAE has higher effective dimension than the reference), $\alpha > 1$, and the shift compresses the schedule toward later timesteps (higher noise levels). This means that for a given $t_n$, the model sees $x_{t_m}$ with more noise than it would at $t_n$ under the original schedule, compensating for the reduced per-dimension corruption. The $\sqrt{m/n}$ form arises from the fact that the expected L2 norm of Gaussian noise in $m$ dimensions scales as $\sqrt{m}$; the shift adjusts the noise-to-signal ratio to match what a $n$-dimensional model would experience.
Practical implementation: The paper follows Esser et al. (2024) in using $n = 4096$ as the base dimension (matching the SD-VAE latent dimensionality: $32 \times 32 \times 4 = 4096$). For RAE with DINOv2-B at $256 \times 256$, the effective data dimension is $m = 256 \times 768 = 196,608$. The scaling factor is $\alpha = \sqrt{196608 / 4096} = \sqrt{48} \approx 6.93$. This is a substantial shift that concentrates training on higher noise levels—exactly where the high-dimensional structure makes the denoising task more informative.
Empirical impact: Table 4 shows the effect of this shift on gFID for DiT-XL on DINOv2-B latents: without the shift, gFID = 23.08; with the shift, gFID = 4.81—a dramatic improvement. This confirms that standard noise schedules (implicitly designed for low-dimensional inputs) are severely mismatched to RAE latent spaces. The shift is a "free" improvement in the sense that it requires no architectural changes and no additional computation—it only changes which timesteps the model sees during training.
Why this form rather than simply training with more noise:
The shifting function $t_m = \frac{\alpha t_n}{1 + (\alpha - 1) t_n}$ has the specific property that it is monotonic, maps $[0, 1]$ to $[0, 1]$, and preserves the relative ordering of timesteps. Simply increasing the noise scale without shifting the schedule (e.g., multiplying $\varepsilon$ by a constant) would change the flow dynamics entirely—the interpolation path $x_t = (1-t)x + t\varepsilon$ assumes $\varepsilon \sim \mathcal{N}(0, I)$, and scaling $\varepsilon$ would break the flow matching formulation. The shift instead changes which $t$ corresponds to which noise level, keeping the underlying stochastic interpolant framework intact. This form is theoretically grounded in the observation that the effective noise level should be a function of the ratio of dimensions, not an additive correction.
Noise-Augmented Decoder Training
The problem: distribution mismatch between training and inference.
The RAE decoder $D$ is trained to reconstruct images from the exact latents $z = E(x)$ produced by the encoder on real images. Formally, the training distribution for the decoder is the discrete empirical distribution $p(z) = \sum_i \delta(z - z_i)$, where $\{z_i\}$ are the latent encodings of the training set images. At inference time, the diffusion model generates latents $\tilde{z}$ by solving an ODE starting from random noise. Due to imperfect training and the inherent approximation error of the diffusion model, $\tilde{z}$ will typically deviate from any exact $z_i$ in the training set—it will be noisy or slightly off-distribution. This is particularly problematic for RAEs because: (1) the latent space has high dimensionality (196,608 dimensions for DINOv2-B), giving the diffusion model many degrees of freedom to deviate; (2) unlike VAEs, which are trained to encode images as continuous Gaussian distributions $\mathcal{N}(\mu, \sigma^2 I)$ (Kingma & Welling, 2014) and whose decoders are therefore naturally exposed to noisy latents during training, the RAE decoder has only seen exact, clean latents.
When $\tilde{z}$ falls outside the support of $p(z)$, the decoder's behavior becomes unpredictable—it may produce artifacts, blur, or semantically inconsistent reconstructions. This is a specific manifestation of the out-of-distribution generalization problem.
The solution: train the decoder on noise-smoothed latents.
Inspired by prior work on Normalizing Flows (Dinh et al., 2017; Ho et al., 2019), the authors augment the decoder training by adding Gaussian noise to the latent tokens before decoding. Instead of training $D$ to map $z \to x$, they train it to map $z' \to x$ where $z' = z + n$ and $n \sim \mathcal{N}(0, \sigma^2 I)$. The training distribution becomes a smoothed version:
which is the convolution of the discrete empirical distribution with a Gaussian kernel. This fills in the "gaps" between training latents in the high-dimensional space, so the decoder learns to handle small perturbations gracefully.
Stochastic noise magnitude:
The authors further introduce stochasticity into the noise standard deviation $\sigma$ by sampling it as $\sigma \sim |\mathcal{N}(0, \tau^2)|$ (the absolute value of a zero-mean Gaussian with standard deviation $\tau$). This means each training batch (or each sample) sees a different noise magnitude, forcing the decoder to be robust across a range of perturbation scales. The default value is $\tau = 0.8$, determined through ablation (Table 15c).
Trade-off between reconstruction and generation:
Table 5 quantifies the impact. Training with $z \sim p_n(z)$ (noise-augmented) vs. $z \sim p(z)$ (clean):
- gFID (generation quality): improves from 4.81 to 4.28—the decoder generalizes better to diffusion outputs.
- rFID (reconstruction quality): worsens slightly from 0.49 to 0.57—adding noise smooths the latent distribution and removes fine detail that the clean decoder could capture.
This trade-off is expected: the noise augmentation acts as a regularizer. It trades some reconstruction fidelity for robustness to input perturbations. The improvement in gFID demonstrates that decoder robustness is a bottleneck for generation quality in the RAE pipeline. The fact that rFID degrades modestly (0.49 to 0.57, still better than SD-VAE's 0.62) while gFID improves substantially (4.81 to 4.28) indicates that the bottleneck is real and that noise augmentation is a well-targeted solution.
Ablation across encoders and sizes:
Table 15a shows the effect of noise-augmented decoding across different encoders. For DINOv2-B, gFID improves from 4.81 to 4.28. For SigLIP2-B, from 6.69 to 4.93. For MAE-B, the effect is most dramatic: from 16.14 to 8.38—nearly halving the gFID. The fact that MAE-B, which has the best reconstruction (0.16 rFID clean) but worst generation (16.14 gFID without noise augmentation), benefits most strongly suggests that MAE's latent space is particularly susceptible to the distribution mismatch problem. Table 15b shows the effect across DINOv2 sizes: DINOv2-S improves from 3.83 to 3.50, DINOv2-B from 4.81 to 4.28, DINOv2-L from 6.77 to 6.09. The benefit is consistent across scales. Table 15c ablates $\tau$ for DINOv2-B: $\tau = 0.0$ (no noise) gives gFID 4.81, rFID 0.49; $\tau = 0.5$ gives gFID 4.39, rFID 0.54; $\tau = 0.8$ gives gFID 4.28, rFID 0.57; $\tau = 1.0$ gives gFID 4.20, rFID 0.60. The trend is monotonic: more noise improves generation at the cost of reconstruction, and the paper selects $\tau = 0.8$ as the default balance point.
Why this is different from VAE training: In a VAE, the encoder outputs parameters of a Gaussian distribution $\mu, \sigma^2$, and the decoder receives samples $z \sim \mathcal{N}(\mu, \sigma^2 I)$ through the reparameterization trick. This means the VAE decoder is trained from the start on noisy latents, with the noise level determined by the learned $\sigma^2$ at each latent position. The noise-augmented training for RAEs mimics this property but with an important difference: the noise level $\sigma$ is a fixed hyperparameter rather than being learned from data. This is simpler and avoids the well-known posterior collapse problem where VAEs learn $\sigma^2 \to 0$ and effectively become deterministic autoencoders. The stochastic $\sigma \sim |\mathcal{N}(0, \tau^2)|$ provides a range of noise scales, similar to how a VAE's learned posterior variance varies across inputs, but without the training instability.
Design choice—absolute value of Gaussian noise for $\sigma$: Sampling $\sigma \sim |\mathcal{N}(0, \tau^2)|$ (absolute value of zero-mean Gaussian) produces a distribution over $[0, \infty)$ that is concentrated near zero but has an exponential tail. The mean of $|z|$ for $z \sim \mathcal{N}(0, \tau^2)$ is $\tau \sqrt{2/\pi}$, so for $\tau = 0.8$, the average noise standard deviation is approximately $0.8 \times 0.798 = 0.638$. The zero-mean Gaussian without absolute value would have half the mass on negative $\sigma$, which is not a valid standard deviation; the absolute value is the standard trick for generating non-negative random variables with a mode near zero and unbounded positive support.
The DiT Architecture with Wide DDT Head
Motivation: the tension between width requirements and computational cost.
Section 4.1 establishes that the diffusion model's width must match or exceed the token dimension ($d \geq n$). For DINOv2-B with $n = 768$, this requires at least DiT-XL (width 1152). But what if even wider models would be better? The standard DiT architecture applies the same width throughout all transformer blocks—input projection, self-attention, feed-forward networks in every layer. Doubling the width of a DiT roughly quadruples the FLOPs per token because both the attention matrix computation and the feed-forward layers scale quadratically with width (attention: $O(d^2 \cdot L)$ for sequence length $L$, feed-forward: $O(d \cdot d_{\text{ff}} \cdot L)$ where typically $d_{\text{ff}} \approx 4d$). This makes scaling width extremely expensive.
The insight from Section 4.1's single-image overfitting experiment is that the width bottleneck is specifically about the model's ability to represent functions in the full ambient dimension. However, much of the semantic processing—attention between tokens, integration of positional information, conditioning on timestep and class—may not require this full dimensionality. It is primarily the final denoising prediction (mapping from internal representations to the $n$-dimensional velocity target) that needs the width. This motivates a decoupled architecture where a standard-width backbone handles token interactions and a wide but shallow head handles the high-dimensional output.
The DDT head design.
The DiT architecture augments a standard DiT backbone $\mathcal{M}$ with a DDT head $\mathcal{H}$. Given a noisy input $x_t$, timestep $t$, and optional class label $y$, the combined model predicts the velocity as:
where $z_t$ is the intermediate representation from the standard DiT backbone, and $\mathcal{H}$ is the DDT head that takes both the original noisy input $x_t$ and the backbone's output $z_t$ as conditioning (plus timestep $t$).
What this computes: The backbone $\mathcal{M}$ processes the noisy tokens through all its layers (self-attention, cross-attention with timestep and class embeddings, feed-forward networks) at the standard width (1152 for DiT-XL). Its output $z_t$ is a contextualized representation for each token that captures interactions with other tokens, temporal conditioning, and class information. The DDT head $\mathcal{H}$ takes the original noisy tokens $x_t$ as its main input (providing the high-dimensional signal) and conditions on $z_t$ and $t$ to produce the final velocity prediction in the full $n$-dimensional space. The head operates at a much wider dimension (2048 by default) but is very shallow (2 layers).
Why this form rather than simply making the backbone wider: The DDT head is inspired by DDT (Wang et al., 2025c), which also decouples the DiT into an encoder-decoder architecture. However, the motivation here is different. In DDT, the decoupling is primarily to apply representation alignment (REPA) loss to the encoder output. In DiT, the motivation is to satisfy the width requirement ($d \geq n$) for the output prediction without paying the quadratic cost throughout the backbone. The head operates on the full token dimension at width 2048, but since it has only 2 layers, the total FLOPs overhead is modest. The standard DiT backbone handles all the heavy computation of attending across tokens and integrating conditioning information; the head handles the final "expansion" to the wide representation needed for accurate velocity prediction.
Architecture details: The default DDT head configuration is 2 layers with hidden dimension 2048 (denoted as "G" width following the DiT naming convention in Table 13). The head does not reapply absolute positional embeddings; it receives the noisy input $x_t$ (which already has positional embeddings from the backbone's input processing) and the backbone output $z_t$. When the backbone hidden dimension differs from the head dimension (which it always does—1152 vs. 2048 for DiT-XL), a linear projection maps $z_t$ to the head dimension. The head applies standard transformer blocks (self-attention + feed-forward) operating on the full sequence of tokens, then projects back to dimension $n = 768$ (for DINOv2-B) with a final linear layer.
Empirical scaling results (Figure 6a): DiT variants across sizes (S, B, L, XL) are compared against the standard DiT-XL baseline. DiT-B achieves better FID (lower is better) than DiT-XL with only ~40% of the training FLOPs. DiT-L, DiT-XL, and even the small variant DiT-S all substantially outperform DiT-XL in the FLOPs-matched regime. When scaled to comparable training budgets, DiT-XL reaches an FID of 2.16 (at 80 epochs), nearly half that of DiT-XL (4.28). This demonstrates that the DDT head is not just a capacity trick but a fundamentally more efficient way to use compute for high-dimensional latent diffusion.
Dependency on encoder size (Table 6): Comparing DiT-XL vs. DiT-XL across DINOv2-S ($n=384$), DINOv2-B ($n=768$), and DINOv2-L ($n=1024$):
- DINOv2-S: DiT FID 2.42 vs. DiT FID 3.50 (gap: 1.08)
- DINOv2-B: DiT FID 2.16 vs. DiT FID 4.28 (gap: 2.12)
- DINOv2-L: DiT FID 2.73 vs. DiT FID 6.09 (gap: 3.36)
The advantage of DiT grows with encoder size (and thus token dimension). The authors attribute this to two mechanisms. First, larger encoders produce higher-dimensional tokens, which amplify the width bottleneck—the DiT must compress through its hidden dimension, losing more information as the gap between $d$ and $n$ grows. DiT's wide head directly addresses this by providing a $d_{\text{head}} = 2048 \gg n$ for the final prediction. Second, the wide head can "filter out noisy information that becomes more prevalent in high-dimensional RAE latents." The shallow depth means the head applies only local transformations (limited receptive field across tokens through 2 self-attention layers), acting as a per-token or near-per-token denoiser that doesn't attempt global reasoning—that's the backbone's job.
Design choice—shallow and wide rather than deep and narrow: Table 16 ablates the DDT head architecture: a 2-layer, 2048-dim head (G width) achieves FID 2.16; a 4-layer, 2048-dim head (double the FLOPs in the head) achieves FID 2.31—slightly worse; a 6-layer, 1152-dim head (similar total capacity but distributed across more layers at narrower width) achieves FID 2.36—also worse. The conclusion is that width matters more than depth for the head. This is consistent with the theoretical picture: the head's primary job is to map from the backbone's compressed representation to the full $n$-dimensional velocity prediction. This requires width to represent the full target space, not depth (which would help with compositional reasoning, already handled by the backbone). The 2-layer depth is a pragmatic minimum—a single projection would be a linear map, which cannot exploit the conditioning signal $z_t$ nonlinearly. Two layers with a nonlinearity between them provides minimal but sufficient nonlinear processing while maintaining the width advantage.
Table 17—the optimal head width grows with encoder size: For DINOv2-S ($n=384$), the optimal head width is 2048 (G), with diminishing returns beyond that (2688 gives FID 2.43 vs. 2.42). For DINOv2-B ($n=768$), 2048 is also optimal (2.16, slightly better than 2688 at 2.22). For DINOv2-L ($n=1024$), 2688 gives 2.64, improving over 2048 at 2.73. The pattern suggests that the head width should be roughly 2–3× the token dimension for optimal performance, with larger ratios needed for larger encoders. The default of 2048 for DINOv2-B ($2048/768 \approx 2.67$) is in this regime.
Critical interaction—DiT does NOT help on low-dimensional VAE latents (Table 10): When tested on SD-VAE latents (token dimension 4, sequence length 1024 with patch size 2), the results reverse: DiT-XL achieves 11.70 FID, significantly worse than DiT-XL's 7.13 FID. This is a crucial finding that demonstrates the specificity of the contribution. The DDT head adds parameters and computation (the wide head) that are beneficial when the latent tokens are high-dimensional and the width bottleneck is active, but harmful when the latent tokens are low-dimensional—presumably because the extra capacity leads to overfitting or introduces optimization challenges without the corresponding need. This result also makes the broader point that DiT is not a universally better architecture; it is an architecture specifically designed to address the width bottleneck created by RAE's high-dimensional latents.
Pixel-space comparison (Table 11): To isolate whether it is high dimensionality or structured representation that matters, the authors train on raw pixels with patch size 16, producing tokens of dimension $16 \times 16 \times 3 = 768$—matching DINOv2-B's token dimension. DiT-XL on pixels: 51.09 FID. DiT-XL on pixels: 30.56 FID. The DDT head does help on pixels (30.56 vs. 51.09), consistent with the width bottleneck argument. But both pixel results are dramatically worse than RAE-based results (4.28 and 2.16). This demonstrates that high dimensionality alone is not sufficient—the structured semantic representation provided by RAE's frozen encoder is essential. The pixel tokens are 768-dimensional but unstructured (they represent raw RGB patches with no semantic organization), making the diffusion task fundamentally harder than on RAE's structured semantic tokens.
Full Training Recipe and Integration
Diffusion model backbone.
The default backbone is LightningDiT (Yao et al., 2025), a variant of DiT optimized for training efficiency. It uses continuous-time flow matching with the linear interpolation path $x_t = (1-t)x + t\varepsilon$ and trains the model to predict the velocity $v(x_t, t) = \varepsilon - x$. The timestep $t \in [0, 1]$ is embedded using Gaussian Fourier features, following Song et al. (2021). Absolute Positional Embeddings (APE) are added to the input tokens in addition to Rotary Position Embeddings (RoPE), though the authors note they "do not observe significant performance difference with or without APE."
Flow matching objective. The training loss is the standard velocity prediction objective from flow matching (Lipman et al., 2023; Liu et al., 2023):
where $v_\theta$ is the model (DiT or DiT) parameterized by $\theta$, $x$ is a clean RAE latent from the training set, $\varepsilon$ is standard Gaussian noise, $t$ is sampled during training (typically with a weighting that accounts for the schedule shift), and the integral is approximated by Monte Carlo sampling of $t$ per batch. The dimension-dependent schedule shift (Section 4.2) modifies which $t$ values the model sees, but the loss formulation itself remains unchanged.
Sequence length and computational equivalence to VAE-based DiTs.
For RAE on $256 \times 256$ images, the encoder produces $N = 256$ tokens (for DINOv2-B with $p_e = 14$ and input interpolated to $224 \times 224$, decoded at $p_d = 16$). The diffusion model uses a patch size of 1, meaning it operates directly on these 256 tokens with no further patching. This exactly matches the sequence length of standard VAE-based DiTs, which use SD-VAE latents of $32 \times 32 \times 4$ and a patch size of 2, yielding $(32/2)^2 = 256$ tokens. Therefore, the self-attention cost in the DiT backbone is identical between RAE-based and VAE-based training. The only difference is the input projection: the VAE-based DiT projects from 4 channels to the hidden dimension, while the RAE-based DiT projects from 768 channels to the hidden dimension—a difference of ~192× in the input projection FLOPs, which is negligible compared to the backbone cost (less than 1% of total GFLOPs, as the authors note).
Optimization for DiT:
Following LightningDiT, the DiT backbone uses AdamW with a constant learning rate of $2.0 \times 10^{-4}$, batch size 1024, and an exponential moving average (EMA) weight of 0.9999 for evaluation. The authors note that they "do not observe instability or abnormal training dynamics with this recipe on DiT." Training runs for 80 epochs by default (the standard for comparisons in the paper), with extended 400-epoch and 800-epoch runs for state-of-the-art results. The DiT uses model configurations as specified in Table 13: DiT-XL has hidden dimension 1152, 16 attention heads, and 28 layers.
Optimization for DiT: The optimization recipe for DiT requires several modifications from the DiT recipe. Using the LightningDiT recipe (constant learning rate, EMA 0.9999) leads to "loss spikes at later epochs and slow EMA model convergence at early epochs." The authors switch to:
- Learning rate schedule: Linear decay from
$2.0 \times 10^{-4}$to$2.0 \times 10^{-5}$, with a constant warmup of 40 epochs. The warmup keeps the learning rate at the maximum value for the first 40 epochs, then linearly decays over the remaining epochs (typically 40 more for an 80-epoch run). - EMA weight: Reduced from 0.9999 to 0.9995. A smaller EMA weight means the EMA model places more weight on recent parameter values, encouraging faster convergence of the evaluation model. The standard 0.9999 EMA is very conservative (effectively averaging over thousands of steps), which works for stable DiT training but is too slow when the DiT architecture changes more rapidly.
- Gradient clipping: Set to 1.0, which is not used in the standard DiT recipe. This prevents the loss spikes that occurred without clipping, suggesting that the DDT head introduces some training instability—perhaps due to the wide-but-shallow architecture creating larger gradient norms in the head layers.
Other hyperparameters (AdamW optimizer, betas, weight decay, batch size) remain the same as DiT. Only the EMA model is evaluated for all reported results.
Sampling:
All generation uses standard ODE sampling with the Euler sampler and 50 steps. The authors find that "performance generally converges above 50 steps," so this is chosen as the default for efficiency. No stochasticity (SDE sampling) is used. The sampling process starts from pure noise $x_1 \sim \mathcal{N}(0, I_n)$ and iteratively applies:
with $\Delta t = 1/50$ for 50 steps. The final output $x_0$ is the denoised latent, which is then passed through the noise-augmented decoder to produce the image.
Guidance methods. The paper primarily uses AutoGuidance (Karras et al., 2025) for improved sample quality. The idea: use a weaker diffusion model (or an earlier checkpoint of the same architecture) to guide a stronger one, analogous to Classifier-Free Guidance (CFG) but without requiring a conditional/unconditional model pair. Specifically, a DiT-S model trained for only 20 epochs (or 14 epochs for the best results) serves as the guiding model. At sampling time, the guided velocity is:
where $w$ is the guidance scale (typically 1.5, swept to 1.42 for best results). The weak model provides an "unbiased" estimate that is deliberately less accurate; the difference $v_{\text{strong}} - v_{\text{weak}}$ points toward regions where the strong model is more confident, and scaling this difference sharpens the distribution. Training the guiding DiT-S for 20 epochs takes "only about 0.05% of the compute used to train the guided model (DiT-XL for 800 epochs)," making AutoGuidance extremely compute-efficient.
The paper also experiments with Classifier-Free Guidance but finds it inferior: "CFG without interval does not improve FID; in fact, applying it from the first diffusion step increases FID" (Appendix I). CFG with Guidance Interval (Kynkääanniemi et al., 2024) can achieve competitive FID after careful grid search, but the best CFG result remains inferior to AutoGuidance on the final model. Only the DiT-XL + DINOv2-S result in Table 8 uses CFG with interval rather than AutoGuidance.
512×512 generation via decoder upsampling (Section 6.1).
A key practical advantage of RAEs is that the decoder is decoupled from both the encoder and the diffusion model. This enables a simple resolution-scaling trick: the diffusion model is trained at $256 \times 256$ resolution (on 256 tokens), but the decoder can be configured with a different patch size. By setting $p_d = 2p_e$, the decoder upsamples each token to a 2× larger output patch, reconstructing a $512 \times 512$ image from the same 256 tokens. The diffusion model does not need to be retrained—the same $256 \times 256$ latent-space model is reused. Table 9 compares this approach to directly training on 512×512 images (which requires 1024 tokens, 4× more). Direct training achieves gFID 1.13, rFID 0.53; upsampling achieves gFID 1.61, rFID 0.97. The upsampling approach is somewhat worse in quality (as expected, since the decoder must hallucinate high-frequency detail not present in the 256-token latent representation) but is 4× more efficient in the diffusion model. For practical deployment, this provides a flexible quality-efficiency trade-off: train the diffusion model once at 256×256, and use different decoder configurations for different output resolutions.
FLOPs accounting for the full pipeline. The paper is careful to note that the computational equivalence to VAE-based DiTs holds for the diffusion model component. The autoencoder cost, however, is substantially different: the RAE encoder (DINOv2-B) requires 22 GFLOPs vs. SD-VAE encoder at 135 GFLOPs (6× less), and the RAE decoder (ViT-XL) requires 106 GFLOPs vs. SD-VAE decoder at 310 GFLOPs (3× less). At training time, the encoder cost is amortized over many diffusion model updates if latents are pre-computed, but at inference time, both encoder (if needed for conditioning) and decoder contribute to the total cost. The decoder efficiency advantage means that, for a single generation, the RAE pipeline is faster end-to-end than the VAE pipeline, all else being equal.
Summary of combined recipe:
To train a state-of-the-art RAE-based diffusion model, the full recipe is: (1) Select a frozen pretrained representation encoder (default: DINOv2-B, patch size 14, dimension 768). (2) Train a ViT decoder (default: ViT-XL) with L1 + LPIPS + adversarial losses on noise-augmented latents ($\tau = 0.8$). (3) Train a DiT diffusion model (default: XL backbone, 2-layer 2048-dim DDT head) with flow matching on RAE latents, using dimension-dependent noise schedule shift ($\alpha = \sqrt{m/4096}$), linear-decay learning rate from $2 \times 10^{-4}$ to $2 \times 10^{-5}$ with 40-epoch warmup, EMA 0.9995, gradient clipping 1.0. (4) At inference, sample with 50-step Euler ODE, apply AutoGuidance with a weak DiT-S checkpoint (guidance scale ~1.5), and decode with the noise-augmented decoder. This pipeline produces 1.51 FID without guidance and 1.13 with guidance at 256×256, converging 47× faster than SiT-XL and 16× faster than REPA-XL.
4. Key Insights and Innovations
Innovation 1: The Width-Matching Principle as a Fundamental Bottleneck, Not a Capacity Concern
Prior to this work, the dominant assumption was that diffusion models struggle in high-dimensional latent spaces due to some combination of optimization difficulty, information-theoretic challenges in denoising, or the common intuition that data manifolds are low-dimensional and generative models only need capacity matching the intrinsic dimension (Pope et al., 2021). The paper reframes the problem entirely: the failure is architectural, not statistical. The inability of a standard DiT to overfit even a single image when its hidden dimension $d$ is smaller than the token dimension $n$ (Figure 3) demonstrates that this is not about dataset complexity, generalization, or manifold learning—it is about representational capacity in the strict linear algebra sense.
This is a conceptual shift from how the field typically thinks about model scaling. The standard narrative—more parameters, more layers, more compute—implies that capacity is fungible: a deeper narrow model should be able to compensate for a shallower wide one, given enough total FLOPs. The paper's critical finding is that for diffusion in high-dimensional spaces, width and depth are not interchangeable. Increasing depth from 12 to 24 layers at width 384 does nothing to close the loss gap (Figure 3, right), while increasing width from 384 to 768—matching the token dimension—collapses the loss to near-zero. This is a categorical rather than continuous bottleneck: you either satisfy the width requirement or you don't, and no amount of depth can substitute.
Theorem 1 formalizes this insight in a way that should change how practitioners think about diffusion model architecture. The function family $\mathcal{G}_d$ representing DiTs with bottleneck dimension $d < n$ can only produce outputs in a $d$-dimensional subspace of the $n$-dimensional target space, and the irreducible loss of $(n-d)/n$ (for the point-mass case where all eigenvalues equal 1) exactly matches the empirical plateau. This is not an approximation bound or a sample-complexity result—it is a representation theorem showing that the model class itself is structurally incapable of solving the problem, regardless of training data or optimization. The mechanism—that Gaussian noise injection during training extends the data distribution's support to the full ambient space $\mathbb{R}^n$, forcing the model to learn a velocity field defined everywhere, not just on the data manifold—explains why this bottleneck is specific to diffusion models and wouldn't apply to, say, GANs or autoregressive models that operate differently.
The practical implication is a new design rule for diffusion transformers: the hidden dimension must meet or exceed the latent token dimension, full stop. This is a crisp, falsifiable constraint that directly determines model sizing—use DiT-XL (width 1152) for DINOv2-B (token dimension 768), or the model will fail regardless of how long you train it. It also explains why the SD-VAE compatible DiT configurations (e.g., DiT-XL with width 1152 processing 4-dimensional VAE tokens) have enormous headroom that is never tested in low-dimensional spaces—the width requirement is trivially satisfied, so practitioners never had to discover this principle.
Innovation 2: High-Dimensional Semantic Latents as an Asset, Not a Liability
The second conceptual shift is the inversion of the field's relationship to latent dimensionality. Since the inception of latent diffusion (Rombach et al., 2022), the entire pipeline has been organized around compression: the VAE aggressively reduces dimensionality (48× for SD-VAE) because diffusing in lower dimensions is assumed to be easier, faster, and more stable. This assumption is so deeply embedded that when practitioners encounter high-dimensional representation encoder features (e.g., DINOv2-B producing 196,608-dimensional latents vs. SD-VAE's 4,096), the immediate reaction is that these spaces are unsuitable for diffusion—as evidenced by the multiple citations the paper provides of this exact claim (Skorokhodov et al., 2025; Yao et al., 2025; Esser et al., 2024; Liu et al., 2024).
The paper inverts this framing. The evidence is in the convergence curves (Figure 4): DiT-XL on RAE latents reaches gFID 4.28 in 80 epochs, while SiT-XL on VAE latents requires ~1400 epochs to reach 8.61—a 47× training speedup that is not about faster iterations but about fundamentally more efficient learning per gradient step. The rich semantic structure of DINOv2 latents (84.5% linear probe accuracy vs. 8.0% for SD-VAE, Table 1d) provides a learning signal that the diffusion model can exploit, converting what was viewed as a liability (high dimensionality) into what is actually an asset (structured representations that already encode much of what the diffusion model needs to learn).
This reframing has implications beyond the specific method. It suggests that the field's historical focus on compression was solving a problem (computational cost of high-dimensional diffusion) that was partly an artifact of using the wrong architectures—standard DiTs couldn't handle high dimensionality, so the solution was compression. But once the width-matching principle is understood and architectures like DiT$^\text{DH}$ are designed accordingly, the compression bottleneck becomes unnecessary. The autoencoder's job shifts from "compress pixels" to "provide a structured latent space," which representation encoders already do as a byproduct of their pretraining objectives. This is a more elegant division of labor: pretraining handles semantic structure, and the diffusion model handles the remaining generative modeling, rather than forcing the autoencoder to learn both compression and semantics from scratch through reconstruction alone.
The comparison between encoders (Table 15a) reinforces this point with a surprising empirical finding: MAE-B achieves the best reconstruction (rFID 0.16, far better than SD-VAE's 0.62) but the worst generation quality (gFID 16.14 without noise augmentation), while DINOv2-B has slightly worse reconstruction (rFID 0.49) but dramatically better generation (gFID 4.81). Reconstruction quality and "diffusability" are decoupled properties of the latent space—a finding that the paper's noise-augmented decoder training partially mitigates (MAE-B improves to 8.38 gFID with augmentation) but doesn't fully eliminate. This decoupling suggests that future work on autoencoder design for diffusion should optimize for diffusability directly, not just reconstruction fidelity, and that representation learning objectives (like DINOv2's self-distillation) may incidentally produce more diffusion-friendly latent spaces than pure reconstruction objectives (like MAE).
Innovation 3: DiT$^\text{DH}$ as a Width-Targeted Architecture, Not a General Model Improvement
The DiT$^\text{DH}$ architecture could easily be mischaracterized as "a better DiT"—but the paper's experiments explicitly show it is worse than standard DiT on low-dimensional VAE latents (11.70 FID vs. 7.13, Table 10). This negative result is one of the paper's most important contributions because it demonstrates that the DDT head specifically addresses the width bottleneck, not some general deficiency in the DiT architecture. The architectural innovation is not "add a wide head to everything" but rather "identify the specific component of the model that is bottlenecked, and add capacity there rather than scaling uniformly."
This is a demonstration of what might be called dimension-aware architecture design: the model's internal dimensionality must be matched to the task's representational requirements at the point where those requirements are most stringent. The standard DiT applies uniform width throughout all layers, which is computationally wasteful because different parts of the model likely have different dimensionality requirements—the early layers that compute token interactions may operate effectively at moderate width, while the final layers that project to the $n$-dimensional velocity prediction need the full width. The DDT head localizes the width increase to where it's needed (the output head), keeping the backbone efficient.
The ablation of head architecture (Table 16) provides a crisp design principle: shallower and wider beats deeper and narrower for this specific function. A 2-layer, 2048-dim head achieves FID 2.16, while a 6-layer, 1152-dim head of similar capacity achieves 2.36—the width matters more than the depth. This is the opposite of what typical transformer scaling would suggest (where deeper is usually better), and it makes sense only in the context of the width-matching principle: the head's primary job is to provide sufficient representational capacity to map into the full $n$-dimensional space, and two layers with a nonlinearity is sufficient for the transformation to be nonlinear, while additional depth doesn't help because the semantic work is already done by the backbone.
The practical impact is a new scaling strategy: rather than building ever-wider DiTs uniformly (which scales FLOPs quadratically), practitioners can attach wider but shallower heads to standard-width backbones, achieving the benefits of width without the full computational cost. This is not an incremental tweak—it's a structural change to how capacity is allocated within the model, motivated by a theoretical understanding of where the bottleneck actually is.
Innovation 4: The Two Assumptions That Held Back the Field, Identified and Falsified
This is more of a meta-contribution, but it is genuinely novel as a diagnostic contribution. The paper explicitly names and falsifies two widespread beliefs that had prevented the community from exploring representation-based latent spaces for diffusion:
Assumption 1: Semantic encoders cannot reconstruct faithfully. The paper shows that with a properly trained ViT decoder (L1 + LPIPS + GAN losses), frozen DINOv2-B achieves 0.49 rFID—better than SD-VAE's 0.62, with a decoder that is 3× more computationally efficient. This directly contradicts the claim quoted from Yu et al. (2024b) that representation encoders "focus on high-level information and can only reconstruct an image with high-level semantic similarities." The key is that the decoder, given sufficient capacity (ViT-XL) and the right training objectives (adversarial loss in particular), can recover the low-level detail that the frozen encoder discards—it learns to invert the semantic abstraction, not just decode it. The fact that even MAE-B, a model trained with a reconstruction objective (masked autoencoding), achieves 0.16 rFID—superior to SD-VAE—further demonstrates that the reconstruction bottleneck in prior work was likely the decoder architecture and training recipe, not the encoder's information content.
Assumption 2: Diffusion fails in high-dimensional latent spaces. The paper shows that with width matching and dimension-dependent noise scheduling, diffusion not only works but converges faster and achieves better final quality in 768-dimensional token spaces than in 4-dimensional VAE spaces. The key diagnostic is that prior observations of poor performance in high-dimensional spaces were likely due to misattributed causation: practitioners observed that diffusion failed in high-dimensional spaces, assumed it was an inherent property of diffusion, but it was actually an architectural mismatch (width $<$ token dimension) and a scheduling issue (noise schedules designed for low-dimensional inputs). Once these are corrected, high dimensionality is revealed as beneficial.
The significance of falsifying these assumptions extends beyond this paper's method. It opens the door for the field to explore a much broader space of autoencoder designs—pretrained encoders from different modalities, different pretraining objectives, different dimensionalities—without being constrained by the belief that only compression-focused VAEs with low-dimensional latents are viable. It also suggests that some of the difficulty attributed to "diffusion in high dimensions" in prior work may have been an architectural artifact, and that revisiting those negative results with width-matched architectures could be productive.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All experiments use the ImageNet-1K dataset (Russakovsky et al., 2015) at resolutions of 256 × 256 and 512 × 512. The training set contains approximately 1.28 million images across 1,000 classes; the validation set (50,000 images, 50 per class) is used for FID computation. The paper notes that while ImageNet classes are not perfectly balanced (ranging from ~732 to 1,300 samples), 895 classes contain exactly 1,300 samples, making the distribution approximately uniform.
-
Base model(s). Three frozen pretrained representation encoders serve as the RAE encoder: DINOv2-B (Oquab et al., 2023) with patch size 14 and hidden dimension 768; SigLIP2-B (Tschannen et al., 2025) with patch size 16 and hidden dimension 768; and MAE-B (He et al., 2021) with patch size 16 and hidden dimension 768. DINOv2-S (dim 384) and DINOv2-L (dim 1024) are also studied for scaling analyses. The diffusion backbone is LightningDiT (Yao et al., 2025) at scales S, B, L, XL, with DiT-XL serving as the default RAE backbone. For the DiT
$^\text{DH}$architecture, a 2-layer, 2048-dim DDT head is attached. The SD-VAE (Rombach et al., 2022) serves as the reference autoencoder for VAE-based baselines. -
Metrics. The primary metric is gFID (generative Fréchet Inception Distance; Heusel et al., 2017), computed on 50,000 generated samples against the ADM reference statistics (Dhariwal & Nichol, 2021) over the full ImageNet training set. rFID (reconstruction FID) measures autoencoder quality on the reconstructed ImageNet validation set. Additional metrics include IS (Inception Score), Precision, and Recall (Kynkäänniemi et al., 2019). For representation quality, linear probing accuracy on ImageNet-1K is reported (Table 1d). Diffusion training loss is also tracked but gFID is the primary quality measure. All FID evaluations use 50 sampling steps with the Euler ODE sampler.
-
Baselines. The paper compares against a comprehensive set of prior methods. VAE-based diffusion models: DiT-XL/2 (Peebles & Xie, 2023), SiT-XL (Ma et al., 2024), MDTv2 (Gao et al., 2023), MaskDiT (Zheng et al.), DiffiT (Hatamizadeh et al., 2024). Representation-aligned diffusion models: REPA-XL (Yu et al., 2025), REPA-E (Leng et al., 2025), DDT (Wang et al., 2025c), REG (Wu et al., 2025). Autoregressive models: VAR (Tian et al., 2024), MAR (Li et al., 2024b), xAR (Ren et al., 2025). Pixel-space diffusion: ADM (Dhariwal & Nichol, 2021), RIN (Jabri et al., 2023), PixelFlow (Chen et al., 2025e), PixNerd (Wang et al., 2025b), SiD2 (Hoogeboom et al., 2025). GANs: BigGAN-deep (Brock et al., 2019), StyleGAN-XL (Sauer et al., 2022). For autoencoder comparisons, SD-VAE (Rombach et al., 2022) is the primary autoencoder baseline. Ablation baselines specific to this paper: DiT-XL trained on RAE latents without any of the three fixes (Table 2), DiT-XL on SD-VAE latents (Table 10), and DiT-XL/DiTDH-XL on raw pixels (Table 11).
-
Generation budget / compute accounting. The primary unit of compute for diffusion models is training GFLOPs, measured as the total floating-point operations consumed during training (Figures 4, 6a, 6b). For convergence comparisons, "training speedup" is measured as the ratio of training epochs or FLOPs required to reach a given gFID (e.g., "47× faster than SiT-XL"). For the autoencoder, decoder GFLOPs are reported per-image at 256 × 256 resolution (Figure 2, Table 1b). For DiTDH variants, the total GFLOPs include both the standard DiT backbone and the DDT head; bubble area in Figure 6c represents model FLOPs. The paper emphasizes that the diffusion backbone sequence length is 256 tokens for both RAE and VAE settings (when using patch size 2 for VAE), so self-attention cost is identical—only the input/output projection dimensions differ, contributing less than 1% to total GFLOPs. At inference, generation cost is measured by the number of sampling steps (50 for Euler ODE).
-
Cross-validation / statistical protocol. No cross-validation is used for the main results—all models are trained once and evaluated on the standard ImageNet validation set. This is standard practice in the image generation literature. The authors acknowledge a subtle evaluation protocol inconsistency: some prior works (VAR, MAR, xAR, DDT) use class-balanced sampling (exactly 50 images per class) while others use uniform random sampling across classes. The paper re-evaluates several recent methods (SiT, REPA, REPA-E, DDT) using class-balanced sampling and reports updated scores in Table 14, finding that class-balanced sampling consistently yields ~0.1 lower (better) FID. All RAE results use class-balanced sampling by default, and Table 14 reports both random and balanced sampling numbers for transparency. For the 800-epoch DiTDH-XL result, FID is evaluated on CPU due to lack of TensorFlow GPU support, though this is noted as a one-off practical constraint.
Main Quantitative Results
Autoencoder Reconstruction Quality and Efficiency
The foundational claim is that frozen representation encoders can reconstruct at quality competitive with or exceeding SD-VAE. Table 1a reports reconstruction FID (rFID) for three encoder types, all paired with a ViT-XL decoder: DINOv2-B achieves 0.49, SigLIP2-B achieves 0.53, MAE-B achieves 0.16, compared to SD-VAE at 0.62. All RAEs outperform SD-VAE; MAE-B's 0.16 is striking given that MAE has no semantic pretraining beyond masked reconstruction. The decoder efficiency advantage is stark: ViT-B already achieves 0.58 rFID at 22.2 GFLOPs—better than SD-VAE's 0.62 at 310.4 GFLOPs (Table 1b). Scaling the decoder from ViT-B to ViT-L to ViT-XL improves rFID from 0.58 to 0.50 to 0.49, with ViT-XL using only ~1/3 of SD-VAE's decoder FLOPs (106.7 vs. 310.4).
Table 1c shows that encoder scaling has minimal impact on rFID: DINOv2-S (0.52), DINOv2-B (0.49), DINOv2-L (0.52)—essentially flat. This suggests that even small representation encoders retain sufficient low-level detail for the decoder to recover.
Table 1d reports linear probe accuracy: DINOv2-B achieves 84.5%, SigLIP2-B 79.1%, MAE-B 68.0%, versus SD-VAE at 8.0%. This exhibits the core claim that RAEs provide semantically rich latents while VAE latents are semantically impoverished—a key enabler for downstream diffusion performance.
Diffusion Training on RAE Latents: The Baseline Failure and Fixes
The first attempted training of a standard DiT (LightningDiT) directly on RAE latents reveals catastrophic failure. Table 2 reports: DiT-S on RAE achieves 215.76 gFID (essentially random), while DiT-S on SD-VAE achieves 51.74. Even DiT-XL on RAE scores 23.08 gFID—far worse than DiT-XL on SD-VAE at 7.13. This establishes the starting point before any fixes are applied.
Section 4.1—Width matching. The single-image overfitting experiment (Figure 3, left) demonstrates the pattern: with width d < n = 768 (token dimension), loss plateaus; at d = 768, loss collapses to near-zero and the image is perfectly reconstructed; at d = 896, behavior is similar. The normalized average loss plateau at (n-d)/n is: for d = 384, approximately (768-384)/768 = 0.5, matching the empirical loss level. Figure 3 (right) shows that increasing depth from 4 to 24 at fixed width 384 does not close the gap—loss remains around 0.5. Table 3 extends the pattern across model/encoder combinations: DiT-S (width 384) achieves loss 3.6e-2 on DINOv2-S (dim 384) but fails completely on DINOv2-B (5.2e-1) and DINOv2-L (6.5e-1). DiT-B (width 768) succeeds on DINOv2-S and B but fails on DINOv2-L. DiT-L (width 1024) succeeds on all. The pattern is binary: convergence requires d >= n.
The solution is straightforward: use DiT-XL (width 1152) with DINOv2-B (token dimension 768). Training on RAE latents with this pairing and none of the other fixes yields the 23.08 gFID reported above (this is the "w/o shift" result in Table 4, as discussed next).
Section 4.2—Dimension-dependent schedule shift. Table 4 reports the impact: without schedule shift, DiT-XL on DINOv2-B achieves gFID 23.08; with the shift (α = sqrt(196608/4096) ≈ 6.93), gFID improves to 4.81. This is a dramatic improvement from a zero-parameter, zero-compute change to the training procedure.
Section 4.3—Noise-augmented decoding. Table 5 reports: training the decoder on clean latents (z ~ p(z)) gives gFID 4.81, rFID 0.49; training on noise-augmented latents (z ~ p_n(z) with τ = 0.8) improves gFID to 4.28 while rFID degrades slightly to 0.57. The gFID improvement demonstrates that decoder robustness to diffusion outputs is a real bottleneck.
Combined effect. With all three fixes applied (width matching via DiT-XL, schedule shift, noise-augmented decoding), DiT-XL achieves gFID 4.28 in 80 epochs (Figure 4). This is compared against SiT-XL (Ma et al., 2024) and REPA-XL (Yu et al., 2025), both operating in VAE latent space. SiT-XL reaches gFID 8.26 at 1400 epochs—RAE-based DiT-XL achieves similar performance (gFID ~8.26) in approximately 30 epochs, a ~47× training speedup. REPA-XL reaches gFID 5.78 at 800 epochs—RAE-based DiT-XL achieves similar or better performance in approximately 50 epochs, a ~16× speedup. At 720 epochs, RAE-based DiT-XL achieves gFID 2.39.
DiT$^\text{DH}$ Architecture: Scaling and Performance
Figure 6a compares DiT$^\text{DH}$ variants (S, B, L, XL) against the DiT-XL baseline on DINOv2-B latents, plotting FID vs. training GFLOPs. DiT$^\text{DH}$-B requires only ~40% of the training FLOPs yet outperforms DiT-XL (lower FID). At comparable training budgets, DiT$^\text{DH}$-XL achieves gFID 2.16 at 80 epochs—nearly half of DiT-XL's 4.28. The scaling curves are steeper for DiT$^\text{DH}$, indicating better FLOP efficiency. Table 6 compares DiT$^\text{DH}$-XL vs. DiT-XL across encoder sizes: on DINOv2-S, gFID 2.42 vs. 3.50 (gap: 1.08); on DINOv2-B, 2.16 vs. 4.28 (gap: 2.12); on DINOv2-L, 2.73 vs. 6.09 (gap: 3.36). The advantage grows with encoder dimension, consistent with the width bottleneck argument—larger encoders produce higher-dimensional tokens, making the wide DDT head more beneficial.
Figure 6b places DiT$^\text{DH}$-XL on DINOv2-B in the context of prior VAE-based methods, plotting FID vs. training GFLOPs. DiT$^\text{DH}$-XL surpasses REPA-XL, MDTv2-XL, and SiT-XL around 5 × 10^10 GFLOPs, and by 5 × 10^11 GFLOPs it achieves the best FID overall, requiring over 40× less compute than the next-best method.
Figure 6c shows scaling across model sizes for DiT$^\text{DH}$ vs. prior methods. DiT$^\text{DH}$-S reaches gFID 6.07, already outperforming the much larger REPA-XL. DiT$^\text{DH}$-B improves to 3.38, surpassing all prior methods of similar or larger scale. DiT$^\text{DH}$-L and XL continue the trend, with XL achieving 2.16 at 80 epochs—the best reported at that training budget. Bubble area reflects model FLOPs; DiT$^\text{DH}$ variants show better FID at lower or comparable FLOPs.
State-of-the-Art Generation Results
Table 8 reports class-conditional performance on ImageNet 256 × 256. At 800 epochs, DiT$^\text{DH}$-XL on DINOv2-B achieves gFID 1.51 without guidance (IS 242.9, Precision 0.79, Recall 0.63)—the best reported among all methods compared. The nearest competitors are REPA-E at 1.70 (800 epochs) and VAR at 1.92 (350 epochs, 2.0B parameters). At 20 epochs, DiT$^\text{DH}$-XL already achieves 3.71; at 80 epochs, 2.16. With AutoGuidance (guidance scale 1.42, using a 14-epoch DiT$^\text{DH}$-S checkpoint as the guiding model), the 800-epoch model achieves gFID 1.13 (IS 262.6, Precision 0.78, Recall 0.67). The guiding model training cost is approximately 0.05% of the guided model's compute. An alternative configuration—DiT-XL (without DDT head) on DINOv2-S latents, trained for 800 epochs—achieves gFID 1.87 without guidance and 1.41 with CFG interval guidance (Table 8). This result uses a smaller encoder (DINOv2-S, dim 384) and the standard DiT architecture, showing that the RAE approach provides benefits even without the DiT$^\text{DH}$ architecture, though the best results require both.
Table 7 reports class-conditional performance on ImageNet 512 × 512. With 400 epochs of training and AutoGuidance, DiT$^\text{DH}$-XL on DINOv2-B achieves gFID 1.13 (IS 259.6, Precision 0.80, Recall 0.63). This surpasses the previous best result of EDM2 (1.25; Karras et al., 2024). Other strong baselines at 512 × 512 include DDT at 1.28, SiD2 at 1.50, and REPA at 2.08.
Table 9 reports the decoder upsampling approach for 512 × 512: training the diffusion model at 256 × 256 and swapping in an upsampling decoder (p_d = 2p_e) achieves gFID 1.61 with AutoGuidance, compared to 1.13 for direct 512 resolution training. The upsampling approach is 4× more efficient in the diffusion model (256 tokens vs. 1024 tokens) at a modest quality cost.
Figure 7 presents qualitative samples from the 512 × 512 model, showing diverse categories (golden retriever, husky, cliff, macaw, arctic fox, balloon) with fine-grained detail and high visual quality. Additional uncurated samples appear in Appendix M (Figures 10–15).
Convergence Speed
Figure 4 plots FID vs. training epochs for DiT-XL on RAE (DINOv2-B), alongside SiT-XL and REPA-XL on VAE latents. The RAE-based model reaches gFID 4.28 at 80 epochs, gFID 2.39 at 720 epochs. SiT-XL requires approximately 1400 epochs to reach 8.61. The ~47× speedup claim is referenced in the text and supported by this curve: RAE-based DiT-XL reaches ~8.26 FID at approximately 30 epochs, while SiT-XL requires ~1400 epochs to reach similar performance. However, note that the gFID values don't perfectly overlap at the claimed crossover point; the speedup is approximate and based on visual extrapolation of the SiT-XL curve.
Figure 6b extends this analysis to DiT$^\text{DH}$-XL, which achieves faster convergence than DiT-XL on RAE. The convergence curve shows DiT$^\text{DH}$-XL surpassing REPA-XL, MDTv2-XL, and SiT-XL (all VAE-based) at significantly lower training FLOPs.
Figure 9 (Appendix H) reports training loss curves for DiT$^\text{DH}$ at scales S, B, L, XL. All models show stable convergence; larger models converge to lower final loss. An EMA weight of 0.9 is used to smooth the loss for visualization.
Guidance Methods
The paper primarily adopts AutoGuidance (Karras et al., 2025). For the best 256 × 256 result (gFID 1.13), a DiT$^\text{DH}$-S checkpoint at 14 epochs serves as the guiding model with guidance scale 1.42. For the best 512 × 512 result (gFID 1.13), guidance scale 1.5 is used with a 20-epoch DiT$^\text{DH}$-S checkpoint. CFG with guidance interval is used only for the DiT-XL + DINOv2-S result in Table 8. The paper notes that CFG without interval actually increases FID on RAE models, and even with careful grid search over scale and interval, the best CFG result remains inferior to AutoGuidance.
FLOPs and Computational Equivalence
The authors emphasize that RAE-based diffusion introduces effectively no extra compute cost compared to VAE-based DiTs, since the sequence length is 256 tokens in both cases (RAE with patch size 1 processes 256 tokens directly; SD-VAE with patch size 2 processes 1024 VAE latents as 256 patches). The input projection from 768 channels to the DiT hidden dimension is negligible relative to the backbone cost. For the autoencoder itself, the RAE encoder requires 22 GFLOPs (DINOv2-B) vs. SD-VAE's 135 GFLOPs; the RAE decoder (ViT-XL) requires 106 GFLOPs vs. SD-VAE's 310 GFLOPs—both substantially more efficient. At generation time, the decoder is the primary additional cost beyond the diffusion model, and RAE's decoder is ~3× more efficient.
Ablation Studies and Robustness Checks
Encoder choice for generation (Table 15a): DINOv2-B achieves the best overall generation performance (gFID 4.28 with noise augmentation), while MAE-B performs substantially worse (8.38) despite having the best reconstruction (rFID 0.16 without augmentation, 0.28 with). SigLIP2-B is intermediate (4.93). This demonstrates that reconstruction quality and generation quality ("diffusability") are decoupled properties—a lower rFID does not guarantee better downstream generation. DINOv2-B is selected as the default encoder for all subsequent experiments.
Noise-augmented decoding across encoders (Table 15a): The noise augmentation (τ = 0.8) improves gFID for all encoders: DINOv2-B (4.81 → 4.28), SigLIP2-B (6.69 → 4.93), and MAE-B (16.14 → 8.38). The improvement is largest for MAE-B, suggesting its latent space is particularly sensitive to the distribution mismatch between training and inference. rFID degrades mildly for all encoders: DINOv2-B (0.49 → 0.57), SigLIP2-B (0.53 → 0.82), MAE-B (0.16 → 0.28).
Noise-augmented decoding across DINOv2 sizes (Table 15b): The benefit holds across encoder scales: DINOv2-S (3.83 → 3.50), DINOv2-B (4.81 → 4.28), DINOv2-L (6.77 → 6.09). The absolute improvement is larger for larger encoders, but the relative improvement is consistent.
Noise magnitude τ for DINOv2-B (Table 15c): Sweeping τ = 0.0, 0.5, 0.8, 1.0 shows monotonic improvement in gFID (4.81, 4.39, 4.28, 4.20) and monotonic degradation in rFID (0.49, 0.54, 0.57, 0.60). The default τ = 0.8 is selected as a balance point; τ = 1.0 would further improve gFID but was not chosen (presumably to limit rFID degradation).
DDT head depth vs. width (Table 16): A 2-layer, 2048-dim (G) head achieves gFID 2.16 with 26.78 GFLOPs in the head. A 6-layer, 1152-dim (XL) head achieves 2.36 at 25.65 GFLOPs—similar cost, worse performance. A 4-layer, 2048-dim head achieves 2.31 at 53.14 GFLOPs—double the cost, slightly worse performance. The conclusion is that width dominates depth for the DDT head: shallower and wider is better than deeper and narrower at comparable FLOPs.
DDT head width vs. RAE encoder size (Table 17): For DINOv2-S, the optimal head width is 2048 (gFID 2.42), with 2688 giving 2.43 (diminishing returns). For DINOv2-B, 2048 is optimal (2.16), with 1536 giving 2.24 and 2688 giving 2.22. For DINOv2-L, 2688 gives the best result (2.64), with 2048 at 2.73 and 1536 at 2.95. The optimal head width grows with encoder size, consistent with the width-matching principle: larger encoders produce higher-dimensional tokens, requiring wider heads for optimal performance. The default 2-layer, 2048-dim head is chosen as a general-purpose configuration.
DiT$^\text{DH}$ on VAE latents (Table 10): DiT$^\text{DH}$-XL on SD-VAE latents achieves gFID 11.70 at 80 epochs—significantly worse than DiT-XL on SD-VAE at 7.13. On DINOv2-B latents, the ordering reverses: DiT$^\text{DH}$-XL achieves 2.16, DiT-XL achieves 4.28. This negative result is critical: the DDT head is specifically beneficial for high-dimensional latent spaces and actually degrades performance in low-dimensional (VAE) spaces. It demonstrates that DiT$^\text{DH}$ is not a generally better architecture but rather a dimension-aware one.
Pixel-space diffusion with matched dimensionality (Table 11): On raw pixels at patch size 16 (producing tokens of dimension 16 × 16 × 3 = 768, matching DINOv2-B), DiT-XL achieves gFID 51.09 and DiT$^\text{DH}$-XL achieves 30.56. While the DDT head provides improvement (consistent with the width bottleneck argument), both results are dramatically worse than RAE-based diffusion (4.28 and 2.16). This demonstrates that high dimensionality alone is not sufficient—the structured semantic representation of RAE is essential. Pixel tokens are unstructured (raw RGB patches with no semantic organization), making the diffusion task fundamentally harder.
Encoder scaling for DiT vs. DiT$^\text{DH}$ (Table 6): The DiT$^\text{DH}$-XL advantage over DiT-XL grows monotonically with encoder dimension: 1.08 FID gap for DINOv2-S (dim 384), 2.12 for DINOv2-B (dim 768), 3.36 for DINOv2-L (dim 1024). This is consistent with the width bottleneck becoming more severe as token dimension increases, and DiT$^\text{DH}$'s wide head providing proportionally greater benefit.
Decoder scaling (Table 1b): ViT-B (22.2 GFLOPs) → rFID 0.58; ViT-L (78.1 GFLOPs) → 0.50; ViT-XL (106.7 GFLOPs) → 0.49. The improvement from L to XL is marginal (0.50 to 0.49), suggesting diminishing returns. ViT-XL is chosen as the default for maximum quality.
FID evaluation protocol (Table 14): Class-balanced sampling (50 images per class) consistently yields ~0.1 lower FID than uniform random sampling across methods. For DiT$^\text{DH}$-XL at 800 epochs: random sampling gives gFID 1.60, balanced gives 1.51 (without guidance); with guidance, random gives 1.28, balanced gives 1.13. Re-evaluated baselines: SiT-XL random 8.61 vs. balanced 8.54; REPA random 5.90 vs. balanced 5.78; DDT random 1.40 vs. balanced 1.26; REPA-E random 1.83 vs. balanced 1.70. The paper raises this as a community awareness issue, noting that as FID values approach lower ranges, these sampling details have greater impact.
Unconditional generation (Table 18): DiT$^\text{DH}$-XL on DINOv2-B with AutoGuidance (200 epochs) achieves gFID 4.96 (IS 123.12), compared to DiT-XL on VAE at 30.68 (IS 32.73) and RCG + DiT-XL at 4.89 (IS 143.2). The RAE-based approach is competitive with RCG, a method specifically designed for unconditional generation, while being substantially simpler (single-stage generation, no separate representation generator).
Critical Assessment
Claim 1: Frozen representation encoders can reconstruct at quality comparable to or better than SD-VAE
Status: Strongly supported for the tested encoders and decoder configuration. Table 1 provides clear evidence: DINOv2-B, SigLIP2-B, and MAE-B all achieve lower rFID than SD-VAE (0.49, 0.53, 0.16 vs. 0.62). The ViT-XL decoder's efficiency advantage (106.7 GFLOPs vs. 310.4 for SD-VAE) is also well-demonstrated.
Qualification: The claim is supported specifically for the combination of these encoders, the ViT-XL decoder, and the training recipe (L1 + LPIPS + GAN with DINO-S/8 discriminator). The paper does not ablate the importance of the GAN loss or the LPIPS term individually for reconstruction quality—an ablation showing what reconstruction looks like with L1 alone would strengthen the claim that all three components are necessary. Additionally, the MAE-B result (0.16 rFID) is notably better than SD-VAE, but the text acknowledges that this same encoder performs worst for generation (gFID 16.14 without noise augmentation). This reveals an important caveat: reconstruction quality and downstream generation quality are not equivalent, so the demonstration of "good reconstruction" alone does not constitute a demonstration of "good latent space for diffusion." The paper is transparent about this decoupling, but the claim as stated in Section 1 ("reconstructions superior to SD-VAE") is technically true only in aggregate; the best-reconstructing encoder (MAE) is the worst-generating one.
Claim 2: Standard DiT fails on RAE latents because model width must match or exceed token dimension
Status: Supported by the single-image overfitting experiment and Theorem 1. Extended to full-dataset training but with less direct evidence of causality.
The single-image experiment (Figure 3) provides a clean, controlled demonstration: width matters categorically, depth does not, and the theoretical bound (n-d)/n matches empirical losses. Table 3 extends this across multiple width/dimension pairs, and the convergence condition (d >= n) holds precisely. Theorem 1 provides a formal justification grounded in the representational capacity of low-rank bottlenecks when the target function has full support.
However, the extrapolation from single-image overfitting to full-dataset training warrants caution. The full-dataset result (Table 2: DiT-XL on RAE achieves 23.08 gFID without schedule shift) is cited as evidence, but this is a single data point: one model size on one encoder. The paper does not show a controlled experiment on the full dataset where only width is varied (e.g., DiT-S, B, L, XL all trained on DINOv2-B for 80 epochs) to demonstrate that the d >= n threshold produces a sharp phase transition in gFID. The single-image result is a necessary demonstration of mechanism, but the full-dataset failure could be partly attributable to other factors (poor optimization dynamics, schedule mismatch) that the width constraint alone does not fully explain. The fact that the schedule shift provides such a dramatic improvement (23.08 → 4.81, Table 4) suggests that width matching is necessary but not sufficient—a model with d >= n still fails badly without the schedule fix.
The paper's argument would be strengthened by showing that, with all other fixes in place (schedule shift, noise-augmented decoding), a DiT with d < n (e.g., DiT-L width 1024 on DINOv2-L dim 1024 is borderline, but DiT-B width 768 on DINOv2-L dim 1024) still fails while one with d >= n succeeds. Table 6 provides some evidence along these lines (DiT-XL width 1152 performs worse as encoder dimension increases: DINOv2-S gFID 3.50, DINOv2-B 4.28, DINOv2-L 6.09), but this is confounded by the fact that larger encoders produce different latent spaces, not just higher-dimensional ones.
Claim 3: The combination of width matching, schedule shift, and noise-augmented decoding enables stable diffusion training in RAE latent space
Status: Supported cumulatively, but the ablation is sequential rather than factorial. Each fix is shown to improve performance: width matching (necessary, per Section 4.1), schedule shift (Table 4: 23.08 → 4.81 gFID for DiT-XL), noise-augmented decoding (Table 5: 4.81 → 4.28 gFID). The full combination yields 4.28 gFID after 80 epochs.
However, a full factorial ablation is not presented. The order of application matters, and the paper does not test, for example, whether noise-augmented decoding helps when width matching is satisfied but schedule shift is absent, or whether schedule shift helps when noise-augmented decoding is absent. The sequential presentation makes it difficult to assess interactions between the fixes—do they address orthogonal problems, or do they overlap? The 23.08 → 4.81 → 4.28 progression suggests diminishing returns from the later fixes, but without a factorial design, the marginal contribution of each fix cannot be precisely isolated. The claim that all three are necessary is plausible but not rigorously tested; it is possible that schedule shift + noise-augmented decoding, without explicit width matching (using a model that happens to satisfy d >= n), would produce similar results, making the width-matching principle a diagnostic insight rather than a required implementation step.
Claim 4: RAE-based training converges substantially faster than VAE-based methods (47× faster than SiT-XL, 16× faster than REPA-XL)
Status: Supported in direction but the exact multipliers should be treated as approximate. Figure 4 shows RAE-based DiT-XL reaching gFID 4.28 at 80 epochs and ~2.39 at 720 epochs. The SiT-XL curve is taken from the literature and plotted at its final performance (8.61 at 1400 epochs). The "47× faster" figure is derived from visual comparison: RAE DiT-XL reaches ~8.26 gFID at approximately 30 epochs, which is ~47× fewer epochs than 1400. However, the SiT-XL curve is not measured at intermediate points in this paper—it is a single data point from prior work—so the exact crossover epoch is approximate. Additionally, comparing "epochs" across methods is not strictly a compute-matched comparison; different methods may have different per-epoch FLOPs due to architectural differences. The paper recognizes this and also reports FLOPs-matched comparisons (Figure 6b), which are more rigorous.
The "16× faster than REPA-XL" claim is subject to the same caveat. Figure 4 shows RAE-based DiT-XL reaching ~5.78 gFID (REPA-XL's 800-epoch performance) at approximately 50 epochs—a 16× factor. Again, this is visual comparison against a single literature data point.
The FLOPs-matched comparison in Figure 6b is stronger: DiT$^\text{DH}$-XL requires ~40× less compute than prior methods at equivalent FID. However, this measures the DiT$^\text{DH}$ architecture on RAE, not the RAE + standard DiT combination. The 47× and 16× figures are for DiT-XL on RAE (no DDT head). Figure 6a confirms that DiT$^\text{DH}$ converges faster than DiT on RAE, so the DiT-XL-on-RAE speedups are lower bounds on what DiT$^\text{DH}$ achieves.
Claim 5: DiT$^\text{DH}$ achieves state-of-the-art ImageNet generation (1.51 FID without guidance, 1.13 with guidance at 256×256)
Status: Supported by the numbers in Table 8, with important evaluation protocol caveats. The 1.51 without guidance and 1.13 with guidance are the best reported in the table. However, the paper's own analysis of FID evaluation protocols (Section 5.1, Table 14) reveals that class-balanced sampling consistently produces ~0.1 lower FID than random sampling. The paper re-evaluates several baselines with class-balanced sampling and reports improved numbers: REPA-E drops from 1.83 (random) to 1.70 (balanced), DDT drops from 1.40 to 1.26, SiT drops from 8.61 to 8.54. This makes cross-method comparisons dependent on knowing which sampling method was used. The paper handles this by reporting both random and balanced numbers for its own method and for re-evaluated baselines. However, several methods in Table 8 are taken from their original papers (VAR, MAR, xAR, ADM, PixelFlow, PixNerd, SiD2, MaskDiT, MDTv2) and may use either protocol—the paper does not re-evaluate all of them, so the comparison is not fully controlled.
The 1.13 with guidance result uses AutoGuidance with a DiT$^\text{DH}$-S checkpoint trained for only 14 epochs (0.05% of the guided model's compute). This is an extremely efficient guidance method, but it means the "1.13 FID" does not reflect the cost of a single training run—it requires training a weaker model as well. The paper is transparent about this, and the cost is negligible in relative terms, but it is a two-model result, not a single-model result.
Claim 6: DiT$^\text{DH}$ is specifically beneficial for high-dimensional latent spaces and degrades on low-dimensional VAE latents
Status: Supported by Table 10, but with a single data point. DiT$^\text{DH}$-XL on SD-VAE achieves 11.70 gFID vs. DiT-XL at 7.13 (both 80 epochs). This is a clear degradation and a strong demonstration that the DDT head is not universally beneficial. However, only XL scale is tested. It is possible that at different scales or with different DDT head configurations, the degradation disappears. The paper does not explore whether a narrower or shallower DDT head could improve VAE-based DiT performance, or whether the degradation is due to the specific head configuration (2-layer, 2048-dim) being aggressively overparameterized for the low-dimensional VAE space. A sweep of DDT head sizes on VAE latents would strengthen this negative result and provide insight into when the DDT head becomes beneficial.
Claim 7: Structured representation is essential—high dimensionality alone is insufficient
Status: Supported by Table 11, with appropriate caveats. Pixel-space diffusion at matched dimensionality (768-dim tokens) achieves gFID 30.56 for DiT$^\text{DH}$-XL vs. 2.16 for RAE. This is a 14× difference, clearly demonstrating that something beyond dimensionality matters. The paper attributes this to the "structured semantic representation" of RAE. However, there are other differences between pixel patches and DINOv2 features beyond "structure": the distribution of pixel patch values is very different from the distribution of DINOv2 features (pixels have bounded range [0, 255], are highly correlated across channels, and have a very different eigenvalue spectrum). These distributional differences could affect diffusion training dynamics independently of "semantic structure." An intermediate experiment—training on PCA-projected or whitened pixel features with matched eigenvalue spectrum—would isolate the effect of semantic structure from the effect of the latent distribution's statistical properties.
General Strengths of the Experimental Design
- The paper uses a consistent compute accounting framework (GFLOPs) and emphasizes that RAE-based DiTs have identical sequence length to VAE-based DiTs, making comparisons fair.
- The ablation of FID evaluation protocol (class-balanced vs. random sampling, Table 14) is unusually rigorous and provides important context for interpreting state-of-the-art comparisons.
- The single-image overfitting experiment (Figure 3) is an elegant diagnostic that isolates the width bottleneck from dataset complexity.
- The negative results—DiT
$^\text{DH}$on VAE (Table 10) and pixel diffusion (Table 11)—are crucial for establishing that the contributions are specific to the RAE + DiT$^\text{DH}$combination, not universal improvements.
General Weaknesses and Missing Experiments
-
Single model family for the diffusion backbone. All experiments use LightningDiT. The paper does not test whether the width-matching principle and the benefits of RAE generalize to other diffusion architectures (U-Net-based diffusion, MDT, etc.). Given that the width bottleneck is argued to be architectural (not specific to DiT), testing on at least one non-DiT architecture would support the claim of generality.
-
No ablation of decoder losses. The decoder is trained with L1 + LPIPS + GAN. The paper does not ablate the contribution of each loss term to reconstruction quality or to downstream gFID. Understanding whether the GAN loss is essential (e.g., for recovering sharp textures that L1 blurs) would inform whether the RAE concept requires adversarial training or could work with purely perceptual losses.
-
Limited exploration of decoder initialization. The decoder is trained from scratch for 16 epochs. The paper doesn't explore whether pretrained weights (e.g., from MAE pretraining) would improve reconstruction or generation, though this is partly justified by the goal of showing that frozen encoders work without additional pretraining.
-
Single dataset. All experiments are on ImageNet-1K. The paper doesn't explore whether RAEs transfer to other image domains (e.g., CIFAR, FFHQ, LSUN) or to non-image modalities (though this is acknowledged as future work in the discussion about "other domains and modalities").
-
The 47× speedup comparison relies on a single literature data point for SiT-XL and visual extrapolation. A head-to-head training of SiT-XL on SD-VAE latents alongside RAE-based DiT-XL, at matched FLOP budgets, would be a more rigorous comparison. The current approach compares against reported final performance rather than training curves.
-
No systematic study of the optimal
d/nratio. The width-matching principle statesd >= n, but the paper uses DiT-XL (width 1152) with DINOv2-B (dim 768), givingd/n = 1.5. It is unclear whether exactly matching (d = n) would work as well, or whether a margin is needed. The DiT scaling results (Table 6) show that DiT-XL (d=1152) achieves gFID 4.28 on DINOv2-B (n=768) but 6.09 on DINOv2-L (n=1024)—hered/ndrops from 1.5 to 1.125, and performance degrades significantly. Whether this degradation is due to the lowerd/nratio or the different latent space properties of DINOv2-L is confounded. -
The noise-augmented decoding ablation sweeps
τbut does not test alternative noise distributions (e.g., uniform, Laplacian). The choice of Gaussian noise with stochastic magnitude (σ ~ |N(0, τ²)|) is justified by reference to Normalizing Flows literature but not compared against simpler alternatives like fixedσor learned per-dimension noise scales.
6. Limitations and Trade-offs
The Difficulty Estimation Cost Is Unaccounted For in the Headline Efficiency Gains
The assumption or constraint. The compute-optimal framework for test-time allocation rests on the ability to estimate prompt difficulty before deciding strategy. The paper's method—generating 2048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted)—is extremely expensive. The authors explicitly acknowledge this in Section 3.2:
"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"
The 2048 samples needed for difficulty estimation alone consume more compute than the largest test-time budgets studied (256–512 generations). The paper reports ~4× efficiency gains in the headline claims (Figures 4 and 8), but these are computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment, the total cost would be difficulty estimation + strategy execution, and the former could dominate the latter. The paper acknowledges this as a key avenue for future work but does not develop or evaluate a cheaper estimator.
The consequence. A practitioner deploying this method would find that the ~4× efficiency gain is an upper bound on achievable savings, not a realized reduction in total compute. If difficulty estimation costs, say, 2048 generations per query, then for a query that receives a test-time budget of 64 generations, the total cost is 2112 generations—making the ~4× gain over best-of-256 essentially vanish when amortized. The approach is viable only if difficulty can be estimated cheaply, either by training a direct difficulty predictor (which the authors suggest) or by amortizing the estimation cost over many similar queries. For one-off or low-volume inference, the overhead is prohibitive. Furthermore, the difficulty bins are computed once and treated as static: there is no mechanism for dynamically adjusting strategy mid-computation based on the model's initial outputs, which could subsume difficulty estimation into the solving process.
What evidence exists in the paper. The paper acknowledges this explicitly (Section 3.2) and shows that predicted (non-oracle) difficulty bins perform nearly as well as oracle bins (Figures 4 and 8, curves largely overlap), which is encouraging but only addresses the "do we need ground truth?" question, not the "how do we get difficulty cheaply?" question. There is no experiment that measures total cost including estimation, and no ablation on how few samples are sufficient for useful difficulty binning.
Mitigation status. Not addressed. The paper flags cheap difficulty prediction as "a key avenue for future work" (Section 8) but provides no evidence that a lightweight estimator can preserve the gains. The predicted bins are still based on 2048 samples, just using PRM scores instead of ground truth—so the cost is unchanged. A more practical approach (e.g., using only 4–8 initial samples and the verifier's score distribution as a difficulty signal) is not tested.
Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Create Capability
The assumption or constraint. The compute-optimal framework assumes that the base model already produces correct solutions at some non-trivial rate. Across all methods studied—PRM search, iterative revisions, and their compute-optimal combinations—the hardest questions (difficulty bin 5) show near-zero improvement regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets up to 256 generations. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%.
The authors are candid about this (Section 7 takeaway box): "test-time compute amplifies existing capability but does not create it from nothing." If the base model's pass@1 is near zero on a problem class, there are essentially no correct solutions in the proposal distribution to find through search or to refine through revisions. Additional compute allocated to these problems is wasted—it neither improves accuracy nor provides useful signal.
The consequence. For deployment on genuinely novel or out-of-distribution reasoning tasks that exceed the base model's training distribution, this approach offers no path forward. The method is effective only on problems within the base model's rough capability range—what the authors call "easy-to-medium" difficulty. For hard problems (bin 5, and to a lesser extent bin 4), pretraining a larger model remains the only viable path. This establishes a sharp boundary condition that limits the scope of the contribution: test-time compute scaling is a complement to pretraining scaling, not a substitute for it, when capabilities are fundamentally missing. A practitioner who needs to solve hard problems where the current model scores near zero must invest in better pretraining, not smarter inference.
What evidence exists in the paper. The bin 5 results are consistently near-zero across all experiments: Figure 3 (right, search), Figure 7 (right, revisions), Figure 9 (FLOPs-matched, bottommost blue lines). The FLOPs-matched analysis (Section 7) explicitly quantifies this: on hard problems at R ≫ 1 (high inference volume), test-time compute shows a −52.9% relative disadvantage compared to the ~14× larger model.
Mitigation status. Not mitigated. The paper identifies this as a fundamental limitation and does not attempt to address it. The boundary is clear and well-characterized: the effectiveness of test-time compute scales with the base model's pass@1 on the problem class. This is a useful characterization, but it means the method cannot be applied to problems that are "too hard" in an absolute sense.
All Results Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*)
The assumption or constraint. Every experiment in the paper uses the MATH benchmark (Hendrycks et al., 2021) with PaLM 2-S* as the base model. The test set contains only 500 questions, which when split into five difficulty quintiles yields ~100 questions per bin, further split by two-fold cross-validation to ~50 questions per fold for strategy selection. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is unverified.
Several aspects of the findings could be model- or domain-specific. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution, including its typical error patterns, calibration properties, and step-by-step reasoning style. A model with different failure modes—e.g., one that produces subtly wrong reasoning rather than clearly identifiable step errors—might yield different PRM behavior and different difficulty-dependent scaling curves. The MATH benchmark consists exclusively of competition-level math problems requiring multi-step symbolic reasoning. It is unclear whether the key findings—beam search hurting easy problems due to over-optimization (Figure 3 right), revisions dominating on easy problems but requiring parallel exploration on hard ones (Figure 7 right)—generalize to other reasoning domains such as code generation (HumanEval, MBPP), logical reasoning (ARC, FOLIO), scientific QA, or factual recall tasks where different cognitive processes are involved.
The consequence. The difficulty-dependent strategy recommendations (best-of-N for easy, beam search for medium, revision chains for easy, balanced parallel-sequential for hard) may be specific to the MATH-PaLM 2-S* combination and may not transfer to other domains or models. A practitioner using a different model (e.g., GPT-4, Llama, Gemini) on a different task (e.g., code generation, planning) would need to re-derive the difficulty bins, strategy allocations, and possibly even the qualitative patterns. The small test set (500 questions, ~50 per fold per bin for strategy selection) further means that the compute-optimal policies selected in this paper may have high variance and may not robustly generalize even within the MATH domain.
What evidence exists in the paper. None for generalization—the paper acknowledges this as a scope limitation by stating that the model is "representative" (Section 4) but does not replicate on any other model, dataset, or task family. The FLOPs-matched comparison uses a second PaLM 2-derived model (the ~14× larger variant) but this is still within the same architecture family.
Mitigation status. Not addressed. The paper does not suggest this as future work explicitly, though the broader direction of extending the framework to other domains is implicit in the conclusion. The lack of cross-model or cross-domain validation is a significant gap for a paper that makes prescriptive claims about how practitioners should allocate inference compute.
Revisions and Search Are Studied Independently, Never Combined
The assumption or constraint. The paper studies two complementary axes for test-time compute—modifying the proposal distribution via iterative revisions (Section 6) and optimizing selection via PRM-guided search (Section 5)—but never combines them. Section 8 explicitly acknowledges:
"we did not experiment with PRM tree-search techniques in combination with revisions"
This is significant because the two mechanisms have complementary strengths: revisions improve the proposal distribution (generating better candidates by conditioning on previous mistakes), while PRM search improves candidate selection (finding the best among generated candidates via step-level scoring). The paper demonstrates that revisions are most effective on easy problems (where sequential refinement is sufficient) and that beam search is most effective on medium-hard problems (where broader exploration is needed). A combined system could potentially: (1) use the revision model as the proposal distribution within PRM-guided beam search, generating higher-quality candidate steps that incorporate correction history; (2) use the PRM to guide which revision paths to pursue, pruning unpromising revision chains early; or (3) sequence the two mechanisms—revise to produce a set of high-quality candidates, then search to select among them.
The consequence. The current results represent a lower bound on what a fully integrated system could achieve. The paper's difficulty-dependent findings for each mechanism independently (search helps on medium problems, revisions help on easy ones) suggest natural combination strategies—but these are not tested. A practitioner would need to determine experimentally whether combining the two mechanisms yields additive gains, subadditive gains (because they address overlapping bottlenecks), or even interference (if the revision model's output distribution confuses the PRM, or if search over-optimization corrupts the revision chain). The paper's stated ~4× improvement over best-of-N may be conservative relative to what combined approaches could achieve, but this is speculative without evidence.
What evidence exists in the paper. None for combinations. The two mechanisms are studied entirely independently, with separate training pipelines (separate PRM trained on base model outputs, separate revision model fine-tuned with edit-distance pairing), separate evaluations, and separate compute-optimal policies. The revision model's PRM transfer experiment (Appendix J, Figure 15a) shows that the base-model PRM underperforms on revision model outputs due to distribution shift, which hints at a practical challenge for combination—the verifier would need to be trained on or adapted to revision model outputs—but this challenge is not directly studied.
Mitigation status. Acknowledged as future work (Section 8) but no mitigation is attempted. The paper presents this as a natural next step rather than a limitation per se, but for a practitioner evaluating whether to adopt this framework, the absence of combined results means the full potential of the approach is unknown, and the recommended per-mechanism strategies may be suboptimal if the mechanisms interact.
The ~14× Larger Model Baseline Is Not Compute-Optimal—Weakens the Pretraining-Vs-Inference Comparison
The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023) rather than the Chinchilla-optimal paradigm (Hoffmann et al., 2022) where both parameters and data are scaled equally. The authors explicitly acknowledge this choice:
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."
Additionally, the ~14× larger model is evaluated with greedy decoding only—no majority voting, no best-of-N, no test-time compute augmentation of any kind. This is a deliberately weak inference-time baseline for the larger model, while the smaller model is given the full benefit of compute-optimal test-time scaling.
The consequence. The headline finding—that a smaller model with test-time compute can outperform a ~14× larger model—is shown against a baseline that may be substantially weaker than what a practitioner would actually deploy. A Chinchilla-optimal model trained with ~14× more total FLOPs (scaling both data and parameters) would likely outperform a parameter-only-scaled model, making the pretraining baseline stronger. Furthermore, giving the larger model even a modest test-time compute budget—say, best-of-8 with majority voting or a simple ORM—would create a much more competitive baseline that is never tested. The reported advantages of test-time compute over pretraining (+27.8% on easy questions at R ≪ 1, Figure 1) may shrink or even reverse against a properly compute-optimal pretraining baseline with basic test-time augmentation.
The paper's framing in Section 7 emphasizes that the comparison operates "in the realistic setting where the correct answer is unknown," which justifies not giving the larger model oracle access to answers—but this does not justify denying the larger model any test-time compute strategy of the kind the paper itself advocates. A fairer comparison would allocate some of the total FLOP budget to test-time compute for both models, with the smaller model receiving more (since it saves pretraining FLOPs) but the larger model receiving some.
What evidence exists in the paper. The FLOPs accounting is clearly documented (Section 7, Equations for X and Y), and the choice to scale parameters only is explicitly acknowledged. The limitation is visible in the results: on medium-to-hard problems at moderate-to-high R values, the test-time compute advantage narrows substantially or reverses (Figure 9), suggesting that the favorable comparisons are concentrated in regimes where the pretraining baseline is weakest. The paper does not include any ablation with a stronger pretraining baseline or with test-time compute allocated to the larger model.
Mitigation status. Acknowledged as future work (Section 7) but not addressed. The paper states that compute-optimal pretraining comparison is left to future work. For the zero-test-time-compute issue, the paper does not discuss this as a potential weakness of the baseline, framing the comparison instead as "small model with inference compute vs. large model without." This is a valid comparison for establishing an upper bound on the substitution rate between pretraining and inference compute, but it may overstate the practical advantage of test-time compute for deployment decisions.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper reshapes how the field thinks about autoencoder design for latent diffusion, shifting the conversation from compression as the primary goal toward representation quality as the organizing principle. Since Rombach et al. (2022) established the latent diffusion paradigm, the autoencoder component has been treated as settled infrastructure: compress pixels into a low-dimensional space, train the best possible diffusion model in that space, and accept the autoencoder's limitations as fixed. The paper's central reframing is that this division of labor—autoencoder handles compression, diffusion model handles everything else—is not optimal and not necessary. A frozen pretrained representation encoder, even one explicitly designed to discard low-level detail for semantic invariance (DINOv2), can serve as a superior latent space when paired with a properly trained decoder, producing both better reconstruction (rFID 0.49 vs. SD-VAE's 0.62, Table 1a) and dramatically faster diffusion convergence (47× faster than SiT-XL, Figure 4).
This is not an incremental improvement to the VAE recipe. It is a different category of autoencoder—one that inherits its latent structure from representation learning rather than learning it from scratch through reconstruction. The practical consequence is that the autoencoder design problem changes from "how do we train a better VAE?" to "which pretrained representation encoder produces the best latent space for diffusion, and how do we adapt the diffusion model to exploit it?" The paper shows that this question was hidden behind two assumptions that turned out to be false—that semantic encoders cannot reconstruct faithfully, and that high-dimensional spaces are inherently hostile to diffusion—and that removing these assumptions unlocks a substantially more efficient and higher-quality pipeline.
The work also introduces a new axis of model scaling: the relationship between latent token dimension and diffusion model width. Section 4.1's width-matching principle (d ≥ n) is a crisp, falsifiable design rule that had not been articulated before because the field never needed it—VAE latent tokens have dimension 4, and every DiT used in practice was already wide enough to satisfy d ≫ 4. The single-image overfitting experiment (Figure 3) and Theorem 1 provide both empirical and theoretical evidence that this is not a capacity guideline but a representational necessity: when the width is insufficient, the loss cannot drop below (n-d)/n regardless of training data or optimization, because the model's output is constrained to a d-dimensional subspace of the n-dimensional target space. This changes how practitioners should select model sizes when working with high-dimensional latents—it is not a smooth scaling law but a hard threshold below which training will fail.
The DiT^DH architecture introduces a dimension-aware design principle: when a representational bottleneck exists at a specific point in the model (the output projection to n-dimensional velocity), add capacity at that point rather than scaling uniformly. The experiment showing that DiT^DH improves generation on RAE latents (2.16 vs. 4.28 gFID, Table 6) but degrades on VAE latents (11.70 vs. 7.13, Table 10) demonstrates that this is not a generally better architecture—it specifically addresses the width bottleneck that only appears in high-dimensional spaces. This negative result is as important as the positive ones: it establishes that architectural innovations must be matched to the properties of the latent space, and that "add a wide head" is not a universal recipe.
The work reconciles two conflicting threads in recent literature. On one side, representation-aligned methods (REPA, DDT, REG) demonstrated that injecting representation knowledge into diffusion training accelerates convergence, but they did so indirectly through auxiliary losses in VAE space. On the other side, attempts to use representation encoders directly for generation (VFM-Tok, DiGIT) applied vector quantization to compress the high-dimensional features, reintroducing an information bottleneck. This paper shows that neither auxiliary alignment nor compression is necessary—the representation encoder's space can be used directly, provided the diffusion model is properly configured. This makes representation-aligned methods look like an intermediate solution to a problem that had a more direct answer all along: simply diffuse in the representation space. Future work on aligning VAE-based diffusion models with external encoders now faces a higher bar—it must demonstrate benefits over the simpler approach of using the encoder's space directly.
The paper also shifts the conversation around decoder robustness. The noise-augmented decoder training (Section 4.3) identifies a specific failure mode—decoder fragility to slightly off-distribution latents generated by imperfect diffusion—and provides a simple fix (adding Gaussian noise during decoder training). This is not a contribution to diffusion model design per se, but it shows that autoencoder-decoder integration is a first-class design problem, not an afterthought. The VAE literature has long incorporated noisy latent training implicitly through the Gaussian posterior, but the connection to diffusion output robustness had not been made explicit. The ablation showing that noise augmentation improves gFID across all encoder types (Table 15a, e.g., MAE-B: 16.14 → 8.38) while mildly degrading rFID (0.16 → 0.28) reveals a trade-off that future autoencoder designs should explicitly optimize: reconstruction fidelity versus downstream generation robustness.
On the scaling side, the result that DiT^DH-XL with DINOv2-B achieves 1.51 gFID without guidance and 1.13 with AutoGuidance (Table 8) places representation-based latent spaces at the top of the ImageNet generation leaderboard. More important than the absolute number is the convergence speed: reaching gFID 4.28 in 80 epochs (Figure 4) versus SiT-XL requiring ~1400 epochs to reach 8.61. This suggests that the semantic structure of RAE latents provides a curriculum effect—the diffusion model does not need to learn semantic organization from scratch, so it can converge faster to a better optimum. The practical implication is that teams training large diffusion models should strongly consider whether their autoencoder is the bottleneck, not just their diffusion backbone.
The finding that DiT^DH degrades on VAE latents (Table 10) and that pixel-space diffusion performs far worse than RAE-based diffusion even at matched dimensionality (Table 11) establishes boundary conditions for the approach. RAE is not a universal replacement for VAE—it is a specific combination of high-dimensional semantic latents with dimension-aware architectures. The pixel-space result (30.56 gFID for DiT^DH on raw pixels with 768-dim tokens vs. 2.16 on DINOv2-B) is particularly important because it isolates the contribution of semantic structure from that of dimensionality. High dimensionality alone does not help; the latent space must carry organized, structured information that the diffusion model can exploit.
Follow-Up Research This Work Enables
Reconciling reconstruction quality and diffusability as separate optimization targets. The paper shows that MAE-B achieves the best reconstruction (rFID 0.16) but the worst generation (gFID 16.14 without noise augmentation, 8.38 with), while DINOv2-B has worse reconstruction (0.49) but much better generation (4.28). This decoupling is observed but not explained. A follow-up study could systematically measure "diffusability" (cf. Skorokhodov et al., 2025) across a wide range of pretrained encoders—DINOv2, SigLIP, CLIP, MAE, JEPA, and their variants at different scales—with a fixed decoder architecture and diffusion recipe. The goal would be to identify which properties of a latent space predict downstream diffusion performance: eigenvalue spectrum of the latent covariance? Smoothness of the velocity field (measured by Lipschitz constant)? Cluster structure? Mutual information with class labels? The paper provides the experimental template (fixed decoder, fixed DiT^DH, standard training recipe) to run this study at scale, and the theoretical framework from Theorem 1 provides tools for analyzing the relationship between latent geometry and diffusion training dynamics.
Training a direct difficulty—or here, "diffusability"—predictor for latent spaces. The paper uses gFID after full training as the evaluation metric for latent space quality, but this requires training an entire diffusion model for each candidate encoder—prohibitively expensive for screening large numbers of encoders. A natural follow-up is to train a lightweight predictor that takes a latent space (or a small sample of encoded images) as input and predicts downstream diffusion performance, analogous to how the earlier paper on compute-optimal scaling needed a difficulty predictor for prompts. Candidate predictors: (1) linear probe accuracy on the latent space (Table 1d shows correlation—DINOv2: 84.5% accuracy, 4.28 gFID; SD-VAE: 8.0%, 7.13 gFID—but MAE at 68.0% with 8.38 gFID breaks the monotonic relationship); (2) the intrinsic dimension of the latent manifold; (3) the Frobenius norm of the velocity field's Jacobian estimated from finite differences on encoded training images; (4) a small "diffusion proxy" trained for only 10 epochs and correlated with final gFID. This would directly address the encoder selection problem that the paper currently solves by empirical sweep (Section G.1).
Combining RAE with on-policy decoder fine-tuning during diffusion training. The noise-augmented decoder training (Section 4.3) treats the diffusion model's output distribution as unknown during decoder training—it adds isotropic Gaussian noise as a generic regularizer. But once a diffusion model is partially trained, one can sample latents from it (at various timesteps, with various numbers of denoising steps, at various guidance scales) and use those samples as additional decoder training data. This would make the decoder robust specifically to the errors the diffusion model actually makes, rather than to generic perturbations. A concrete experiment: train the decoder with the standard noise augmentation for 16 epochs, then train a diffusion model for 40 epochs, then generate "pseudo-clean" latents from the diffusion model (fully denoised outputs at t=0), add these to the decoder training set, and fine-tune the decoder for 2 more epochs. Measure whether gFID improves beyond the 4.28 baseline (DiT-XL on DINOv2-B) and whether the optimal τ for initial noise augmentation changes when on-policy data is available.
Stress-testing the width-matching principle across architectures and modalities. Theorem 1 makes a strong claim: any model with a low-rank bottleneck d < n cannot represent the velocity field. The empirical validation is on DiT architectures with specific projection structures. A stress-test would: (1) test whether U-Net-based diffusion models (which process latents of shape N × d with convolutional rather than attention-based operations) exhibit the same bottleneck—does a U-Net with bottleneck channel count c < d fail to overfit a single RAE latent image, and does the loss plateau at a predictable value? (2) test whether the principle extends to text or audio modalities, where token dimensions differ—do discrete diffusion models on high-dimensional text embeddings (e.g., from a language model's hidden states) exhibit the same width constraint? (3) test whether architectural modifications that avoid explicit low-rank projection (e.g., factorized attention, cross-attention to a wider latent) can circumvent the bound while maintaining d < n internal representations. Negative results—finding architectures or modalities where the bound does not hold—would refine our understanding of when Theorem 1 applies.
RAE for video and multi-modal generation. The paper demonstrates RAE on ImageNet image generation. The most natural extension is to video, where the efficiency argument is even stronger: video VAEs must compress both spatial and temporal dimensions, and the SD-VAE's reconstruction burden scales poorly. A frozen video representation encoder (e.g., a pretrained VideoMAE or a spatial-temporal DINOv2 variant) could serve as the RAE encoder, with a space-time ViT decoder trained with the same L1 + LPIPS + GAN recipe. The key question is whether the N = 256 token count (spatial only) scales gracefully to N × T tokens for temporal length T, or whether the DDT head's width requirement becomes prohibitive. A concrete experiment: train a DINOv2-B RAE decoder on individual video frames, freeze it, and train a space-time DiT^DH on frame sequences with causal or bidirectional temporal attention. Measure gFID-equivalent metrics (FVD, IS) and compare convergence speed against VAE-based video diffusion models like Stable Video Diffusion. The decoder upsampling trick (Section 6.1, p_d = 2p_e for 2× upsampling) becomes even more valuable for video, where training the diffusion model at low spatial resolution and upsampling with the decoder could save ~4× in the temporal dimension as well.
Understanding why DINOv2 latents are more "diffusable" than MAE latents through representation dynamics. The paper identifies the empirical fact—DINOv2 achieves gFID 4.28, MAE achieves 8.38 (both with noise augmentation, Table 15a)—but does not explain the mechanism. DINOv2 is trained with self-distillation (student-teacher matching across augmented views), which encourages invariance to nuisance transformations and produces features that cluster by semantic category. MAE is trained with masked reconstruction, which encourages each patch token to be individually decodable and produces features with less global semantic organization. A hypothesis: the flow matching velocity field v(x_t, t) = E[ε - x | x_t] is "smoother" in DINOv2's space because similar images map to nearby latents, making the conditional expectation easier to learn, while MAE's space is more "fragmented" with no semantic clustering, making the velocity field more irregular. A follow-up could test this by: (1) measuring the local Lipschitz constant of a trained velocity predictor in DINOv2-B vs. MAE-B latent space; (2) measuring the alignment between the velocity field and the principal curvature directions of the data manifold; (3) interpolating between DINOv2 and MAE latents (linearly in latent space) and measuring whether the interpolated velocity fields are well-behaved (indicating a smoother manifold). Understanding this mechanism would enable principled design of representation objectives that optimize for diffusability alongside downstream task performance.
Practical Applications and Downstream Use Cases
Training budget reallocation: train smaller diffusion models longer on RAE latents rather than scaling model size. The convergence results in Figure 4 and Figure 6b imply that, given a fixed compute budget, a practitioner gets substantially better generation quality by adopting RAE + DiT^DH and training for more epochs than by scaling up a standard DiT on VAE latents. Concretely: DiT^DH-XL on DINOv2-B reaches gFID 2.16 in 80 epochs (Figure 6a). Training it for 800 epochs yields 1.51 without guidance (Table 8). The total training FLOPs for 800 epochs of DiT^DH-XL are substantial, but the alternative—training a larger model on VAE latents to reach comparable FID—would require scaling DiT beyond XL size (e.g., to XXL or H, Table 13) at quadratic FLOPs cost, and the paper's Figure 6c shows that VAE-based methods at larger FLOPs (REPA-XL, MDTv2-XL) do not match RAE-based DiT^DH performance. For teams with a fixed GPU budget, the recommendation is: adopt a pretrained DINOv2-B encoder (freely available), train a ViT-XL decoder for 16 epochs (106 GFLOPs per image, 3× cheaper than SD-VAE decoder), and invest the bulk of the compute budget in training DiT^DH-XL for 400–800 epochs rather than experimenting with larger architectures on VAE latents. The decoder upsampling trick (Section 6.1) further extends this: the same diffusion model trained at 256×256 can generate 512×512 outputs at competitive quality (gFID 1.61 vs. 1.13 for direct training, Table 9) without retraining the diffusion model, saving ~4× in the diffusion training budget for high-resolution applications.
Pretrained representation encoders as a new axis of model customization. The paper demonstrates that three different frozen encoders—DINOv2 (self-supervised), SigLIP2 (language-supervised), and MAE (reconstruction-supervised)—all produce viable RAE latent spaces with different generation characteristics (Table 15a). This suggests a practical workflow for domain-specific generation: select a pretrained encoder that is strong on the target domain (e.g., a medical imaging encoder like BiomedCLIP for radiology image generation, a satellite imagery encoder for remote sensing, or a domain-specific DINOv2 fine-tune), freeze it, train a decoder on domain data, and then train a DiT^DH on domain-specific latents. The encoder provides domain knowledge "for free" through its pretraining; the decoder and diffusion model adapt it to generation. This is substantially cheaper than training a domain-specific VAE from scratch (which would require both encoder and decoder training) and benefits from the convergence speedup demonstrated on ImageNet. A concrete deployment: a team with 50K domain-specific images could train an RAE decoder in ~1 GPU-day (16 epochs at batch size 512), then train a DiT^DH-S (smallest variant, gFID 6.07 on ImageNet at 80 epochs, Figure 6c) in a few GPU-days, achieving generation quality that would require a much larger model if trained from scratch in pixel space or with a domain-specific VAE.
Efficient high-resolution generation via decoder-only scaling. The decoupling of encoder, diffusion model, and decoder in the RAE pipeline enables a pragmatic resolution-scaling strategy that is not possible with standard VAEs (where the encoder-decoder compression ratio is fixed). Train the diffusion model once at a moderate resolution (256×256, 256 tokens). For different output resolutions, swap in different decoder configurations: a 1× decoder (p_d = p_e) for native resolution, a 2× decoder (p_d = 2p_e) for 2× upsampling, and potentially a 4× decoder for 4× upsampling. Each decoder is trained separately (16 epochs each, ~106 GFLOPs per image at ViT-XL size) but the diffusion model is shared. For applications that need both thumbnail and high-resolution outputs (e.g., an image generation API with multiple quality tiers), this avoids training separate diffusion models for each resolution. The quality trade-off documented in Table 9 (gFID 1.61 for 2× upsampling vs. 1.13 for native 512×512, a 0.48 FID gap) quantifies the cost of this flexibility—acceptable for many applications where perfect fidelity is not required, and substantially cheaper than training a 1024-token diffusion model for native 512×512 generation.
Accelerating research iteration on diffusion model architectures. The ~16× convergence speedup over REPA-XL (Figure 4) changes the economics of research experimentation. An ablation that previously required 800 epochs of training (e.g., comparing two diffusion backbone variants on ImageNet) can now produce meaningful results in 50 epochs. Teams developing new diffusion architectures, training objectives, or conditioning mechanisms can use RAE latents as a fast evaluation platform: train on DINOv2-B RAE for 50–80 epochs, compare gFID across variants, and only scale the winning configuration to long training runs. The paper provides the complete recipe (decoder training in Appendix C, diffusion training in Section D, DDT head configuration in Table 16–17) and demonstrates that relative performance in the 80-epoch regime correlates with final performance at 800 epochs (compare DiT^DH-XL at 80 epochs: 2.16 gFID, vs. 800 epochs: 1.51; the improvement continues but the ranking across methods is preserved). This use case matters especially for academic groups with limited compute, who can now contribute to diffusion architecture research without needing industrial-scale GPU clusters.
When to Prefer This Method
The paper's experimental results—particularly the negative results in Tables 10 and 11—provide clear conditions under which the RAE + DiT^DH combination is advantageous versus when standard VAE-based approaches remain preferable.
-
Prefer RAE + DiT^DH when the target domain has a strong pretrained representation encoder available (DINOv2, SigLIP2, or a domain-specific encoder with ImageNet-comparable quality), when training compute is the primary bottleneck rather than inference latency, and when the target latent token dimension is high (≥384 for DINOv2-S, ≥768 for DINOv2-B). The width-matching principle (
d ≥ n) must be satisfied either through a wide-enough DiT backbone or the DDT head—if total model width is forced below token dimension by hardware constraints, the approach will fail (Section 4.1). The convergence advantage (47× faster than SiT-XL, Figure 4) is most valuable when long training runs are impractical. -
Prefer standard VAE + DiT when using a VAE latent space that is already well-optimized for the domain (SD-VAE for natural images), when the latent token dimension is small (≤16 channels), or when inference latency is critical and the decoder cost matters—though RAE's ViT-XL decoder is already 3× more efficient than SD-VAE's (106 vs. 310 GFLOPs, Table 1b). The negative result in Table 10 (DiT^DH on VAE: 11.70 gFID vs. DiT-XL: 7.13) shows that the DDT head is specifically harmful in low-dimensional spaces, so standard DiT architectures should be used with VAEs.
-
Prefer pixel-space diffusion only when the domain lacks any pretrained encoder and training a decoder is impractical—but pixel diffusion requires substantially more training compute and achieves far worse quality (Table 11: DiT^DH-XL on pixels: 30.56 gFID vs. RAE: 2.16), making it a fallback when representation encoders are unavailable, not a competitive alternative.
