PaperlayerView exact revision on arXiv

Source-linked AI summary

Improved Baselines with Representation Autoencoders

Jaskirat Singh, Boyang Zheng, Zongze Wu, Richard Zhang, Eli Shechtman, Saining Xie

arXiv:2605.18324v2cs.CVcs.AIcs.GRcs.LGstat.ML

TL;DR

RAE still has reconstruction, guidance, and representation-design limitations despite replacing VAEs with pretrained vision encoders. This paper simplifies RAE by aggregating encoder layers, combining RAE with REPA, and reformulating REPA as latent-space x-prediction, yielding faster convergence and strong ImageNet results across additional settings.

  • Problem

    RAE adoption is limited by weaker reconstruction, incompatibility with traditional CFG, and reliance on final-layer encoder features.

  • Method

    RAEv2 aggregates the last k encoder layers, uses the same representation for RAE encoding and REPA targets, and reformulates REPA as x-prediction for internal guidance.

  • Results

    RAEv2 achieves over 10× faster convergence, state-of-the-art gFID 1.06 and FDr6 2.17 in 80 epochs, and EPFID@2 of 35 epochs versus 177 for original RAE.

  • Takeaways & Limitations

    RAE and REPA are complementary rather than redundant, and their combination improves reconstruction, generation, guidance, and training efficiency across evaluated settings.

  • Takeaways & Limitations

    The paper evaluates only simple feature aggregation methods and selects representations through empirical search rather than directly optimizing them end-to-end.

Abstract

from arXiv · show

Representation Autoencoders (RAE) replace traditional VAE with pretrained vision encoders. In this paper, we systematically investigate several design choices and find three insights which simplify and improve RAE. First, we study a generalized formulation where the representation is defined as sum of the last k encoder layers rather than solely the final layer. This simple change greatly improves reconstruction without encoder finetuning or specialized data (e.g., text, faces). Second, we study the prevalent assumption that RAE (using pretrained representation as encoder) replaces representation alignment (REPA), which distills the same representation to intermediate layers instead. Through large-scale empirical analysis, we uncover a surprising finding: RAE and REPA exhibit complementary working mechanisms, allowing the same representation to be used as both encoder and target for intermediate diffusion layers. Finally, the original RAE struggles with classifier-free guidance (CFG) and requires training a second, weaker diffusion model for AutoGuidance (AG). We show that REPA itself can be viewed as x-prediction in RAE latent space. By simply re-parameterizing the output of the DiT model, it can provide guidance for "free". Overall, RAEv2 leads to more than 10x faster convergence over the original RAE, achieving a state-of-the-art gFID of 1.06 in just 80 epochs on ImageNet-256. On FDr6, RAEv2 achieves a state-of-the-art 2.17 at just 80 epochs compared to the previous best 3.26 (800 epochs) without any post-training. This motivates EPFID@k (epochs to reach unguided gFID < k) as a measure of training efficiency. RAEv2 attains an EPFID@2 of 35 epochs, versus 177 for the original RAE. We also validate our approach across diverse settings for text-to-image generation and navigation world models, showing consistent improvements. The code is available at https://raev2.github.io.

1. Introduction

RAEv2 addresses RAE’s reconstruction, guidance, and representation-design challenges through three complementary insights: aggregating encoder layers, combining RAE with REPA, and reusing REPA for internal guidance. These changes improve training efficiency while reducing reliance on extra models and specialized data.

  • Motivation: RAE faces reconstruction gaps, CFG incompatibility, and limited exploration of representations beyond pretrained encoders’ final layers.Traditional RAE also requires a separate weaker diffusion model for AutoGuidance, increasing compute and complexity.
  • Generalized Representation Autoencoder: Summing the last k encoder layers provides training-free control over reconstruction quality and yields Pareto-optimal reconstruction-generation performance.The formulation uses pretrained encoder features distributed across layers rather than only the final layer.
  • RAE and REPA: RAE and REPA have complementary mechanisms: RAE supplies semantically rich latents, while REPA improves intermediate diffusion features’ spatial structure.This supports using the same representation as both the RAE encoder and the REPA target.
  • Guidance: REPA can serve as x-prediction in RAE latent space, enabling internal guidance from the existing prediction head instead of a separately trained weaker model.Re-parameterizing the output head also avoids CFG’s additional unconditional forward pass and effectively halves NFEs.
  • Training efficiency: RAEv2 converges over 10× faster than original RAE, reaching gFID 1.06 in 80 epochs and EPFID@2 of 35 epochs versus 177 for RAE.It also reaches FDr6 2.17 in 80 epochs versus the previous best 3.26 in 800 epochs without post-training.

2. Improved Representation Autoencoders

RAEv2 improves RAE through three design insights: aggregating encoder layers, combining RAE with REPA, and using REPA for self-guidance. Together, these changes improve reconstruction and generation while simplifying guidance.

  • 2.1. Generalized Representation Encoder: Aggregating the last k encoder-layer features improves reconstruction without encoder finetuning or specialized data.The generalized formulation preserves the original latent shape and adds no learned parameters.
  • 2.2. RAE and REPA exhibit Complementary Working Mechanisms: RAE and REPA consistently improve generation when used together, despite the assumption that their shared representation creates a wasteful skip connection.The same pretrained representation can serve as both the RAE encoder and the REPA intermediate target.
  • 2.2. RAE and REPA exhibit Complementary Working Mechanisms: RAE supplies semantic information while REPA improves the spatial structure of diffusion features.Across 27 encoders, RAE performance correlates more with global semantics, whereas REPA performance correlates more with spatial structure.
  • 2.2. RAE and REPA exhibit Complementary Working Mechanisms: Stronger encoders such as DINOv3-L perform best with RAEv2 because they combine strong global semantics and spatial structure.The combined method benefits from encoders that perform well on both representation properties.
  • 2.3. Reformulating REPA as x-prediction with RAE: Reformulating the DiT output as x-prediction lets the REPA head provide internal guidance without a separate AutoGuidance model or an additional CFG forward pass.The REPA head predicts the clean RAE latent from early-layer features and runs in the same forward pass as the main model.

3. Experiments

Experiments show that RAEv2 improves reconstruction, guided generation, encoder selection, guidance efficiency, convergence, and evaluation across scales and metrics. The gains arise from generalized layer aggregation, complementary RAE–REPA mechanisms, and a faster training recipe.

  • Encoder selection: DINOv3-L gives the best RAEv2 generation performance among tested encoders, despite stronger encoders performing worse under the original RAE recipe.The result is attributed to representations combining strong global semantics and spatial performance.
  • Generalized RAE: MLS consistently outperforms MLR for Stage-2 generation, so RAEv2 uses simple addition of the last K encoder layers as its default aggregation.The two parameter-free schemes are effectively tied on Stage-1 reconstruction.
  • Generalized RAE: At K=23, reconstruction reaches rFID 0.18 and PSNR 27.03 versus the standard baseline’s rFID 0.60 and PSNR 18.93; guided gFID is best at K=7 with 1.06.Unguided gFID is best near K=1 at 1.50 after 80 epochs.
  • Generalized RAE: The generalized formulation improves reconstruction and guided generation while preserving ImageNet linear-probing performance and enabling unified understanding and generation tokenization.This preserves the representation’s global semantics while improving reconstruction-generation performance.
  • Guidance: Internal guidance with the REPA head achieves the best gFID at no extra inference cost, whereas CFG provides little improvement and AutoGuidance requires an additional model and forward pass.This uses the x-prediction formulation described for RAEv2 guidance.
  • Reconstruction-generation trade-off: RAEv2 achieves a Pareto-optimal reconstruction-generation trade-off without encoder finetuning or specialized data.Qualitative comparisons also report competitive reconstruction against proprietary VAEs.

4. Generalization to Other Tasks

RAEv2 generalizes beyond ImageNet to text-to-image generation and action-conditioned navigation world models. It improves prompt adherence and video prediction quality while converging faster and preserving scene structure over long rollouts.

  • Scope: The paper validates RAEv2 on text-to-image generation and action-conditioned future-frame prediction.These experiments test generalization across distinct conditioning and rollout settings.
  • Text-to-image generation: Text-to-image evaluation uses GenEval, DPG-Bench, and GenAI-Bench with 50-step Euler ODE sampling from the EMA model.Training uses JourneyDB, BLIP3o subsets, and BLIP3o-60k finetuning.
  • Text-to-image generation: In text-to-image generation, RAEv2 improves GenEval from 41.7 to 62.4 during pretraining and reaches 82.7 after finetuning.The finetuned score exceeds Flux-VAE’s 78.3 and original RAE’s 81.5.
  • Navigation world models: RAEv2-NWM achieves FVD 105.61 on RECON, outperforming DIAMOND at 762.73, NWM at 200.97, and RAE at 312.01.The same ordering holds for FID and LPIPS across horizons from 1 to 16 seconds.
  • Navigation world models: Generalized RAE improves temporal consistency because earlier encoder layers retain low-level texture and geometry useful for future-state prediction.Qualitative rollouts show less flickering and better preservation of scene structure.
  • Navigation world models: RAEv2-NWM matches RAE’s final FID within 10K iterations and reaches FID 7.5 versus 18.0 and LPIPS 0.24 versus 0.29 within about 30K iterations.The comparison uses single-shot predictions with random target offsets of 1–8 frames at 4 FPS.

5. Related Work

The paper builds on work using pretrained vision encoders as latent spaces, representation alignment, and reconstruction improvements. It argues that RAEv2 combines these directions through complementary mechanisms while avoiding encoder finetuning or specialized data.

  • Pretrained encoders as latent spaces: Pretrained vision encoders increasingly replace VAE latents in diffusion, and RAEv2 improves the original recipe with more than 10× faster convergence.This work positions representation autoencoders as a step toward unified tokenization for understanding and generation.
  • Representation alignment: Representation alignment distills pretrained representations into intermediate diffusion layers, while RAE uses those representations as the latent-space encoder.The paper tests the assumption that using RAE eliminates the need for REPA.
  • Representation alignment: RAE and REPA work through complementary mechanisms, so their combination supports stronger encoders and simplifies guidance with RAE.The paper’s large-scale analysis motivates using the same representation as encoder and intermediate target.
  • Reconstruction quality: Prior reconstruction approaches use specialized data or finetune pretrained encoders, whereas RAEv2 obtains Pareto-optimal reconstruction-generation performance from frozen encoders.The approach relies on low-level details already present in pretrained vision representations.

6. Conclusion

RAEv2 simplifies RAE through generalized layer aggregation, complementary RAE–REPA use, and simpler guidance. Across ImageNet, text-to-image, and world-model experiments, it improves reconstruction, convergence, and generation efficiency.

  • Conclusion: RAEv2 aggregates the last K encoder layers, combines RAE with REPA, and uses their interaction to simplify guidance.These changes improve reconstruction without encoder finetuning or specialized data.
  • Conclusion: RAEv2 achieves 10× faster convergence, state-of-the-art gFID and FDr6 in 80 epochs, and supports stronger representations such as DINOv3-L.The paper validates the approach on ImageNet, text-to-image generation, and world models.

A. Implementation Details

RAEv2 uses pretrained vision encoders and DiT-based flow matching across ImageNet, text-to-image, and navigation settings. Its implementation combines self-REPA, in-context conditioning, and efficient sampling and evaluation protocols.

  • Encoder and architecture: The DiT_DH-XL backbone uses a 28-block transformer encoder and a 2-block decoder, producing 256 image-latent tokens with latent patch size 1.The model uses RMSNorm, SwiGLU, RoPE, and absolute positional embeddings.
  • Encoder and architecture: The default DINOv3-L encoder produces a 1024 × 16 × 16 latent representation from 256 × 256 images, while text-to-image uses SiGLIP2-B with 768-dimensional features.Both configurations use a 16 × 16 patch grid.
  • Encoder and architecture: RAEv2 evaluates eight pretrained vision-encoder families through a unified interface that extracts normalized patch tokens while discarding [CLS] and register tokens.The evaluated families include DINOv2, DINOv3, WebSSL, Perception Encoders, MoCov3, CLIP, I-JEPA, and MAE.
  • Training and guidance: Self-REPA aligns intermediate features at encoder block depth 8 using a linear projection and coefficient λ=0.5, with the target encoder matching the RAE encoder.This configuration is reported to consistently improve generation across pretrained encoders.
  • Training and guidance: At inference, flow matching uses velocity prediction, while self-guidance converts outputs to x-prediction and applies guidance through the REPA head.The reported evaluation uses Euler-discretized ODE sampling, gFID, Inception Score, FD_r, FD_6^r, and EPFID@k.
  • Multimodal settings: The implementation extends to text-to-image and navigation world models by replacing class-conditioning tokens with text or past-frame and action tokens.Navigation conditions on four past frames, four action tokens, and one Fourier-embedded time token, totaling 1029 conditioning tokens.
  • Compute and efficiency: RAEv2 reaches gFID 1.06 in roughly 12 hours on a 4 × 8 H100 setup, compared with over a week for original RAE training.The same comparison is reported for the original RAE at 800 epochs.

C.1. Comparisons with original RAE

Across reconstruction, generation, guidance, and representation quality, RAEv2 improves or matches original RAE while using a training-free multi-layer representation and internal guidance.

  • Generalized RAE: RAEv2’s multi-layer sum is strictly training-free and achieves the best generation quality at K=7 and reconstruction quality at K=23.Prior methods rely on auxiliary losses, encoder finetuning, or architectural modifications.
  • Generalized RAE: Training with additional data and for longer consistently improves RAEv2 decoder reconstruction performance.The reported comparisons use only 16 training epochs with a frozen pretrained vision encoder.
  • Generalized RAE: MLS dominates MLR on Stage-2 gFID for every K, while the two parameter-free aggregation methods are essentially tied on Stage-1 reconstruction metrics.MLS sums the last K encoder layers; MLR concatenates them and projects with a fixed random matrix.
  • Guidance: REPA Guidance achieves the best gFID and Inception Score for K=7 and K=23 without a separate model or an extra forward pass.This contrasts with AutoGuidance’s additional model and CFG’s extra forward pass.
  • Guidance: x-prediction outperforms default internal guidance at K=7, confirming its importance for RAEv2 generation.The table identifies x-prediction with using REPA at intermediate layers.
  • Understanding: The generalized formulation improves reconstruction and generation without meaningfully affecting ImageNet linear-probing performance.This supports unified tokenization for understanding and generation.

C.2. Text-to-Image Generation

RAEv2 improves text-to-image generation in both pretraining and finetuning settings. The setup adapts DiT_DH-XL to caption conditioning and evaluates compositional, dense-prompt, and human-preference performance.

  • Training setup: The text-to-image model is pretrained on JourneyDB and BLIP3o caption subsets, then finetuned on BLIP3o-60k using SiGLIP2-B and Qwen3-0.6B text embeddings.Pretraining uses 150K iterations, while finetuning uses 50 epochs with batch size 1024.
  • Evaluation: Evaluation covers GenEval, DPG-Bench, and GenAI-Bench using 50-step Euler sampling with the EMA model.These benchmarks cover compositionality, dense prompts, and human preference.
  • Pretraining results: GenEval improves from 41.7 with Flux-VAE and 58.4 with original RAE to 62.4 with RAEv2 during pretraining.The passage attributes further gains to the improved RAEv2 training recipe.
  • Training efficiency: Across the reported text-to-image settings, RAEv2 trains faster than original RAE and Flux-VAE.The comparison is summarized for both pretraining and finetuning configurations.

C.3. Navigation World Models

RAEv2-NWM applies the generalized RAE to autoregressive navigation world modeling with frame and action conditioning. It produces more accurate, temporally stable rollouts and converges faster than the RAE baseline and prior systems.

  • Task and setup: The navigation model predicts future frames from four past egocentric frames and an action sequence, feeding predicted frames back autoregressively for long-horizon rollout.Each target frame is predicted from a sliding context window and the next ground-truth action.
  • Task and setup: The model encodes four 256 × 256 frames into 1024 context tokens and adds four action tokens plus one time token, totaling 1029 conditioning tokens.Each frame contributes a 16 × 16 patch grid.
  • Evaluation: Evaluation reports FID, LPIPS, PSNR, and DreamSim at horizons of 1, 2, 4, 8, and 16 seconds.Predicted RGB frames are re-encoded and fed back as context during rollout evaluation.
  • Rollout quality: RAEv2-NWM achieves FVD 105.61, compared with 762.73 for DIAMOND, 200.97 for NWM, and 312.01 for RAE.The same ordering holds across all horizons on both FID and LPIPS, with qualitatively less flickering.
  • Representation choice: The generalized representation aggregates the last K encoder layers, whose earlier features retain texture and geometry useful for temporally consistent navigation rollouts.The passage links these features to a substantial convergence-speed improvement.
  • Convergence: Within approximately 30K iterations, RAEv2-NWM reaches lower FID and LPIPS than the RAE baseline, with FID 7.5 versus 18.0 and LPIPS 0.24 versus 0.29.It matches the RAE baseline’s final FID within the first 10K iterations.

D. Qualitative Results

RAEv2 produces strong text-to-image samples despite a relatively short training schedule and small model size, with prompt adherence across diverse subjects.

  • RAEv2 generates 256×256 text-to-image samples with strong prompt adherence across animals, landscapes, and stylized scenes.The 0.9B model uses self-guidance with the REPA head and is evaluated on MJHQ test-set prompts.

E. Discussion and Limitations

The discussion identifies two limitations: simple feature aggregation choices and empirical selection from a discrete encoder set. The paper also documents its use of automated tools for figures, literature search, and writing.

  • The generalized RAE study considers only simple addition and random projection for aggregating pretrained encoder features.The authors suggest that better aggregation optimization could improve both generation and reconstruction.
  • RAEv2 selects its representation through empirical search over a discrete set of pretrained encoders.Future work aims to optimize the representation directly through end-to-end learning.
  • The paper states that figures were generated from experiment logs and checkpoints using Claude Code, with LLM assistance for related-work search and formulation.Cursor was also used for parts of the paper writing.
Loading 2605.18324v2…