Technical Summary: Improved Baselines with Representation Autoencoders (RAEv2)

Problem Statement

Representation Autoencoders (RAE) have emerged as a promising framework for replacing traditional Variational Autoencoders (VAEs) in diffusion transformer training, aiming to unify tokenization for both visual understanding and generation. However, the original RAE formulation faces three critical bottlenecks hindering practical adoption:

1. Suboptimal Reconstruction: Reconstruction performance lags behind specialized VAEs, particularly when trained without specialized data (e.g., text or faces).
2. Incompatibility with Guidance: RAE struggles with traditional Classifier-Free Guidance (CFG), necessitating the training of a secondary, weaker diffusion model for AutoGuidance (AG). This adds significant computational overhead and complexity.
3. Underutilized Encoder Representations: Prior work defaults to using only the final layer of a pretrained vision encoder as the latent representation, ignoring the rich, diverse abstractions present in intermediate layers.

Methodology

The paper proposes RAEv2, an improved baseline that addresses these issues through three key insights and design changes, all while maintaining a training-free approach for the encoder.

1. Generalized Representation Encoder

Instead of relying solely on the final layer ($K=1$) of a pretrained vision encoder, the authors propose a generalized formulation where the latent representation is the sum of the last $K$ encoder layers.

• Mechanism: The output is defined as $\mathbf{x} = \sum_{\ell=L-K+1}^{L} \mathbf{z}_\ell$, where $\mathbf{z}_\ell$ represents the feature map at layer $\ell$.
• Rationale: Different layers capture complementary features; earlier-to-middle layers retain finer spatial structures, while later layers emphasize global semantics. Aggregating them preserves geometric structure without increasing latent dimensionality or requiring additional learned parameters.
• Implementation: The authors compare simple addition (MLS) against random-matrix projection (MLR) and find that simple addition consistently yields better generation performance while improving reconstruction.

2. Complementary Mechanisms of RAE and REPA

The paper challenges the prevailing assumption that RAE eliminates the need for Representation Alignment (REPA). REPA distills the same pretrained representation into intermediate diffusion layers.

• Finding: Empirical analysis across 27 vision encoders reveals that RAE and REPA operate through complementary mechanisms.
  • RAE provides a semantically rich latent space (leveraging global semantics, measured by Linear Probing accuracy).
  • REPA regularizes the spatial structure of intermediate diffusion features (measured by Local Distance Similarity, LDS).
• Result: Using the same representation as both the encoder (RAE) and the target for intermediate layers (REPA) consistently improves generation performance. This synergy allows stronger encoders (e.g., DINOv3-L), which excel in both global and spatial metrics, to achieve superior generation results.

3. REPA as Self-Guidance (x-prediction)

The original RAE requires a separate weak model for guidance. The authors reformulate the REPA head to enable internal guidance.

• Insight: In the RAE latent space, the REPA projection head (which maps early-layer features to the clean latent) effectively performs x-prediction.
• Mechanism: By re-parameterizing the full model output to also be x-prediction, the REPA head's prediction ($\hat{\mathbf{x}}_{repa}$) serves as a "weaker" baseline compared to the full model's prediction ($\hat{\mathbf{x}}_{full}$).
• Guidance Formula: Guidance is applied internally within a single forward pass: $\hat{\mathbf{x}}_{guided} = \hat{\mathbf{x}}_{full} + w \cdot (\hat{\mathbf{x}}_{full} - \hat{\mathbf{x}}_{repa})$.
• Benefit: This eliminates the need for a separate AutoGuidance model and avoids the extra forward pass required by CFG, effectively halving the Number of Function Evaluations (NFEs) for guidance.

Key Results

RAEv2 demonstrates significant improvements in convergence speed, generation quality, and reconstruction across ImageNet, text-to-image, and navigation world models.

• Training Efficiency: RAEv2 achieves 10$\times$ faster convergence compared to the original RAE. It reaches a state-of-the-art guided gFID of 1.06 in just 80 epochs on ImageNet-256, whereas the original RAE required 800 epochs to reach 3.26.
• New Efficiency Metric (EPFID): The paper introduces EPFID@k (epochs to reach unguided gFID $\le k$) as a more practical metric than absolute gFID. RAEv2 achieves an EPFID@2 of 35 epochs, a massive improvement over the original RAE (177 epochs) and prior works (>480 epochs).
• Generation Quality: On the recently proposed Representation Fréchet Distance ($FD_r$), RAEv2 achieves a score of 2.17 in 80 epochs, surpassing the previous best of 3.26 (trained for 800 epochs) without any post-training.
• Reconstruction: The generalized formulation (specifically with $K=23$) achieves a Pareto-optimal trade-off between reconstruction (rFID 0.18) and generation, outperforming proprietary VAEs (e.g., Flux-VAE, SD3.5) despite being trained only on ImageNet.
• Generalization:
  • Text-to-Image: RAEv2 shows consistent improvements over Flux-VAE and original RAE on GenEval, DPG-Bench, and GenAI-Bench.
  • Navigation World Models: In future-frame prediction tasks (RECON dataset), RAEv2-NWM achieves an FVD of 105.61, significantly outperforming NWM (200.97) and RAE (312.01), with reduced flickering in autoregressive rollouts.

Significance and Claims

The paper claims that RAEv2 provides a simplified, highly efficient baseline for representation autoencoders. Its primary contributions are:

1. Simplification: It removes the need for encoder fine-tuning, specialized training data, or separate guidance models.
2. Efficiency: It drastically reduces the training time required to reach state-of-the-art performance, shifting the focus from absolute metric improvements to training efficiency (EPFID).
3. Unified Tokenization: By preserving the global semantic understanding capabilities of the frozen encoder (verified via linear probing) while enhancing generation and reconstruction, it reinforces the viability of using pretrained vision encoders as a unified tokenization scheme for both understanding and generation tasks.

The authors hope these insights facilitate the practical adoption of representation autoencoders in diverse generative modeling scenarios.