PaperlayerView exact revision on arXiv

Source-linked AI summary

RoAR-VAE: Improving Latent Diffusability through Robust Prediction

Yang Yue, Fangyun Wei, Yitong Wang, Jinjing Zhao, Henry Zheng, Dong Chen, Gao Huang

arXiv:2610.26216v1cs.CV

TL;DR

Visual tokenizers must preserve reconstruction fidelity while producing latents that downstream diffusion models can learn efficiently. RoAR-VAE predicts semantic and image targets from corrupted latents, improving the average downstream benchmark score by 2.19 points over the strongest competing learnability-enhanced tokenizer while maintaining high-fidelity reconstruction.

  • Problem

    Visual tokenizers must balance faithful image reconstruction with latent representations that downstream diffusion models can learn efficiently.

    Source: §1 · p. 1

  • Method

    RoAR-VAE combines RobustAlign, which predicts pretrained semantic features from corrupted latents, with RobustRecon, which reconstructs the clean image from corrupted latents.

    Source: §1 · p. 2

  • Results

    RoAR-VAE exceeds the strongest competing learnability-enhanced tokenizer by 2.19 points on average downstream benchmark score while maintaining high-fidelity reconstruction.

    Source: §4.2 · p. 8, §5 · p. 10

  • Takeaways & Limitations

    Robust prediction offers a simple principle for designing visual tokenizers that better balance reconstruction and downstream generation.

    Source: §5 · p. 10

    The tokenizer is general-purpose and is neither trained specifically on human-centric datasets nor optimized for human-related attributes, despite evaluation on a reconstruction benchmark that includes human faces.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

An effective visual tokenizer must preserve image fidelity while producing latents that downstream diffusion models can learn efficiently. Prior work improves latent diffusability through semantic alignment or joint tokenizer--diffusion training. Revisiting REPA-E, a representative approach, we find that removing its diffusion objective and dedicated generative blocks while retaining semantic prediction from corrupted latents yields comparable downstream generation performance under a matched setting. This observation motivates a general principle of robust prediction: generation-relevant information should remain recoverable from corrupted latents. We instantiate this principle with two complementary objectives. RobustAlign predicts pretrained semantic features from corrupted latents to preserve high-level structure, while RobustRecon reconstructs the clean image to retain fine-grained spatial and appearance information. Together, they form Robust Alignment and Reconstruction VAE (RoAR-VAE). Our analysis suggests that robust prediction improves diffusability by making useful information more accessible in high-variance latent directions, which diffusion models tend to learn earlier. Empirically, RoAR-VAE improves average generation benchmark performance by 2.19 points over the strongest baseline tokenizer under the same training budget, while improving small-text reconstruction accuracy by 16.93 percentage points over FLUX.1-VAE. These results demonstrate that RoAR-VAE achieves a favorable balance between reconstruction fidelity and downstream generative learnability.

1 INTRODUCTION

Visual tokenizers must balance faithful image reconstruction with latents that downstream diffusion models can learn efficiently. The introduction shows that semantic prediction from corrupted latents can match REPA-E's generation performance and motivates combining it with image reconstruction.

1 INTRODUCTION

  • Visual tokenizers must balance faithful reconstruction with representations that downstream diffusion models can learn efficiently (Rombach et al., 2022; Esser et al., 2024).Higher-capacity or weakly regularized latents may retain more detail yet be harder to model.Source: §1 · p. 1
  • Prior work improves diffusability through semantic supervision (Yao et al., 2025; 2026) or joint tokenizer–diffusion optimization (Heek et al., 2026; Duggal et al., 2026), while REPA-E combines both.The introduction asks whether REPA-E's improvement comes from joint diffusion training or semantic supervision; only semantic alignment directly updates the tokenizer encoder.Source: §1 · p. 1
  • Under matched evaluation, RobustAlign achieves comparable downstream generation performance to REPA-E after removing its diffusion objective and dedicated generative blocks.The result motivates semantic prediction from corrupted latents as an effective learning signal for a diffusion-friendly latent space.Source: §1 · p. 2
  • RoAR combines RobustAlign's semantic prediction and RobustRecon's clean-image reconstruction from corrupted latents, while the standard VAE objective preserves reconstruction from clean latents.The objectives are complementary: RobustAlign emphasizes object-level organization and foreground–background separation, RobustRecon richer local variation, and their combination exhibits both.Source: §1 · p. 2
  • Under Gaussian corruption, higher-variance latent directions retain more signal, and robust prediction may make useful information more accessible in these directions, which diffusion models tend to learn earlier.The analysis suggests robust prediction may strengthen principal modes or enrich the image information they encode.Source: §1 · p. 2

2 RELATED WORK

Related work improves tokenizer diffusability through semantic representations, joint tokenizer–diffusion training, and perturbation-based regularization. RoAR revisits REPA-E and frames robust prediction as an alternative source of supervision, with diffusion learning dynamics offering a possible explanation for its effect.

2 RELATED WORK

  • Figure 2 traces the transition from REPA-E's joint training to RobustAlign's isolated semantic-prediction pathway and RoAR's added image-reconstruction target under corruption.The standard VAE objective remains on clean latents.Source: Figure 2 · p. 3
  • Prior approaches improve tokenizer representations through pretrained-feature alignment, frozen pretrained encoders, joint tokenizer–diffusion optimization, equivariance, or perturbation-based regularization.The cited work also addresses fine-grained information loss when pretrained encoders are repurposed for generation.Source: §2 · pp. 2–3
  • Wang & Pehlevan (2025) derive an inverse-variance learning timescale for tractable denoisers, showing that higher-variance covariance modes are learned earlier.This learning-dynamics result provides context for RoAR's focus on making useful information accessible in dominant modes.Source: §2 · p. 3
  • RoAR revisits REPA-E (Leng et al., 2025a), which combines semantic enhancement with joint tokenizer–diffusion training, and identifies robust prediction as an effective supervision source for latent learnability.Its proposed explanation links information accessibility in dominant latent modes to diffusion models' tendency to learn those modes earlier.Source: §2 · p. 3

3 METHOD

The method tests whether semantic prediction from corrupted latents can shape a diffusion-friendly tokenizer without joint diffusion training, then combines semantic prediction and image reconstruction. RoAR-VAE is analyzed as making generation-relevant information accessible in dominant latent modes.

3.1 PRELIMINARIES

  • The VAE encodes images as latents for reconstruction, and downstream latent diffusion operates on normalized latents.The VAE objective combines pixel reconstruction, perceptual, posterior-regularization, and adversarial losses.Source: §3.1 · p. 4

3.2 ROBUST ALIGNMENT

  • RobustAlign removes REPA-E’s flow-matching objective and diffusion-only components, retaining the VAE objective and semantic prediction from corrupted latents.Its semantic predictor recovers pretrained teacher features, and the alignment loss backpropagates through corrupted latents to the encoder.Source: §3.2 · pp. 4–5, §3.1 · p. 4
  • RobustAlign-VAE performs comparably to REPA-E in matched downstream generation evaluation, while both improve over Vanilla-KL-VAE.Table 1 reports GenEval after 4k, 8k, 16k, and 32k downstream optimization steps.Source: Table 1 · p. 4, §3.2 · p. 5

3.3 ROBUST RECONSTRUCTION AND THE ROAR FRAMEWORK

  • RobustRecon predicts the clean image from corrupted latents, encouraging fine-grained image information to remain recoverable across a broad range of noise levels.Its auxiliary decoder uses pixel and perceptual reconstruction losses, has independent parameters, and is discarded after tokenizer training.Source: §3.3 · p. 6
  • RoAR-VAE combines RobustAlign’s high-level semantic structure with RobustRecon’s fine-grained image details in one tokenizer objective.Its latent representations retain both object-level organization and richer local structure, including boundaries, textures, and intra-object variations.Source: §3.3 · p. 6

3.4 DISCUSSION

  • At a fixed corruption level, higher-variance latent modes retain more signal relative to noise, offering a possible explanation for robust prediction’s diffusability gains.The proposed account is that robust prediction keeps target-relevant semantics and spatial details accessible in dominant modes, which diffusion models learn more rapidly.Source: §3.4 · p. 6

4 EXPERIMENTS

The experiments evaluate reconstruction and downstream generation under controlled tokenizer comparisons. RoAR-VAE balances both objectives, preserves fine-grained details, and benefits from combining robust prediction objectives with latent corruption.

4.1 EXPERIMENTAL SETUP

  • RoAR-VAE uses an ImageNet-trained tokenizer and is evaluated with frozen tokenizers and a common text-conditioned diffusion setup across three generation benchmarks.Tokenizer training uses a 200k-step VAE warm-up followed by 250k refinement steps; comparisons exclude FLUX.2-VAE from baselines because its training data and recipe are undisclosed.Source: §4.1 · pp. 6–7

4.2 MAIN RESULTS

  • RoAR-VAE improves PSNR by 1.80 dB and reduces LPIPS from 0.0510 to 0.0349 versus FLUX.1-VAE, while exceeding the strongest competing learnability-enhanced tokenizer by 2.19 points on average downstream benchmark score.The comparison uses a fixed 32k-step downstream training budget; Table 2 covers reconstruction and generation results.Source: Table 2 · p. 7, §4.2 · p. 8
  • RoAR-VAE produces clearer object shapes and finer visual structure than FLUX.1-VAE, while VA-VAE generations are smoother and less detailed.The qualitative comparison covers downstream updates at 32k and 120k.Source: Figure 4 · p. 7, §4.2 · p. 8
  • RoAR-VAE achieves the best reported text and face reconstruction metrics, improving small-text accuracy by 16.93 percentage points over FLUX.1-VAE and small-face similarity from 0.53 to 0.65.The TokBench results reveal fine-grained information loss in VA-VAE, VTP, and PAE despite their competitive generation scores; Figure 5 shows examples of textual and facial distortions.Source: §4.2 · pp. 8–9, Table 3 · p. 8, Figure 5 · p. 8

4.3 ANALYSIS

  • Combining RobustAlign and RobustRecon raises the average score to 63.86, above RobustRecon alone at 59.38 and RobustAlign alone at 62.49; removing latent corruption lowers it to 58.47.The ablation also finds that either objective alone improves over Vanilla-KL at 50.50, and the combined method outperforms either objective alone across training budgets.Source: §4.3 · p. 9, Table 4 · p. 9, Figure 6 · p. 9
  • Replacing RobustAlign’s Transformer predictor with an MLP causes a substantial performance drop, while sharing RobustRecon predictor and decoder parameters slightly degrades performance.The ablations suggest the semantic prediction pathway benefits from greater expressivity and the pixel predictor and reconstruction decoder serve distinct roles.Source: §4.3 · p. 9

5 CONCLUSION

RoAR-VAE combines RobustAlign for semantic structure with RobustRecon for fine-grained spatial and appearance details, aiming to balance reconstruction fidelity and downstream generative learnability. The section also records AI-use disclosures, ethics guidance, and reproducibility information.

5 CONCLUSION

  • RoAR-VAE combines RobustAlign for semantic structure and RobustRecon for fine-grained spatial and appearance details to balance reconstruction fidelity and downstream generative learnability.The conclusion presents robust prediction as a source of supervision for diffusion-friendly representations.Source: §5 · pp. 9–10

AI USE STATEMENT

  • The authors used generative AI for implementation, dataset preparation, translation, research feedback, literature work, and manuscript drafting, but not to develop the conceptual framework or formulate hypotheses or mathematical claims.Source: AI USE STATEMENT · p. 10
  • The authors state that AI-assisted outputs were reviewed and that they take responsibility for the paper’s final content and artifacts.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

  • The reconstruction benchmark includes human faces, but the tokenizer is general-purpose and neither specifically trained on human-centric datasets nor optimized for human-related attributes; downstream systems should consider safeguards.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The paper points to Section 3 and Appendix A for method and experiment details and says implementation code and trained models will be released.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 TOKENIZER TRAINING

RoAR-VAE uses a convolutional autoencoder trained first with a standard VAE objective and then refined with robust prediction objectives. The appendix specifies its architecture, losses, predictors, and optimizer.

A.1 TOKENIZER TRAINING

  • RoAR-VAE uses an approximately 84M-parameter convolutional autoencoder with 8× spatial downsampling and 32 latent channels, trained on ImageNet at 256 × 256 resolution.The encoder and decoder each have four resolution stages with widths (128, 256, 512, 512) and two residual blocks per stage.Source: §A.1 · p. 14
  • Training begins with 200k steps of VAE-objective warm-up, then adds 250k refinement steps using RoAR objectives with λA = 1.5 and λR = 0.5.Warm-up uses unit L1 and LPIPS weights, KL weight 10−6, and DINOv2-S-based adversarial supervision; corruption level t is sampled uniformly from [0, 1] during refinement.Source: §A.1 · p. 14
  • RobustAlign predicts frozen DINOv2-B features with an 8-layer Transformer and MLP projector, while RobustRecon predicts the clean image from corrupted latents using L1 and LPIPS losses.The RobustRecon pixel predictor is initialized from the warm-up VAE decoder.Source: §A.1 · p. 14
  • The complete RoAR framework is optimized with AdamW using batch size 256 and learning rate 1 × 10−5.Source: §A.1 · p. 14

A.2 DOWNSTREAM GENERATIVE MODELING

Downstream text-conditioned diffusion models are trained from scratch with tokenizer parameters frozen. The setup adjusts latent patch size across tokenizer compression factors and specifies RoAR-VAE’s timestep sampling and inference settings.

A.2 DOWNSTREAM GENERATIVE MODELING

  • For each tokenizer, a text-conditioned MMDiT diffusion Transformer is trained from scratch on filtered, recaptioned LAION data with tokenizer parameters frozen.Models use 256 × 256 resolution, global batch size 1,024, and a frozen Qwen3-0.6B text encoder; latent patch size is adjusted to keep sequence length consistent across compression factors.Source: §A.2 · p. 14
  • RoAR-VAE uses shifted logit-normal diffusion-timestep sampling with exponential shifting and μ = 1.5, while each baseline uses its official or recommended scheduler.All models use AdamW, constant learning rate 2 × 10−4, and caption dropout 0.1; inference uses Euler sampling for 50 steps with guidance weight 5.0.Source: §A.2 · p. 14

A.3 EVALUATION DETAILS

Evaluation compares RoAR-VAE with reconstruction-oriented and diffusion-oriented tokenizer baselines, using image-level and fine-detail reconstruction measures alongside text-to-image benchmarks. Generation scores are measured under a fixed downstream training budget.

A.3 EVALUATION DETAILS

  • The baselines include SD3-VAE and FLUX.1-VAE, alignment- or representation-learning tokenizers VA-VAE and VTP, PAE, REPA-E-FLUX, and FLUX.2-VAE.For REPA-E-FLUX, evaluation uses its resulting tokenizer and retrains the downstream diffusion model from scratch; FLUX.2-VAE is included despite undisclosed training data and recipe.Source: §A.3 · p. 14
  • Reconstruction is evaluated on ImageNet validation images with rFID, PSNR, LPIPS, and SSIM, and on TokBench with text-recognition and face-embedding metrics for small and overall instances.TokBench includes challenging text- and face-centric images; its text measures are T-ACC and T-NED, and its face measure is F-Sim.Source: §A.3 · p. 15
  • Text-to-image alignment is evaluated on GenEval, GenEval2, and DPG-Bench after 32k downstream optimization steps, with scores averaged over four sampling seeds and summarized by their arithmetic mean.The benchmarks assess compositional generation, broader visual concepts, and dense prompts, respectively.Source: §A.3 · p. 15
Loading 2610.26216v1…