RoAR-VAE: Improving Latent Diffusability through Robust Prediction
Yang Yue, Fangyun Wei, Yitong Wang, Jinjing Zhao, Henry Zheng, Dong Chen, Gao Huang
cs.CV
2026-10-19
By predicting clean targets from corrupted latents, RoAR-VAE reaches 63.86 avg on three T2I benchmarks, +2.19 over the best baseline, plus 16.93pp small-text accuracy over FLUX.1-VAE.
Latent diffusion models never touch pixels directly. A visual tokenizer, usually a VAE, compresses images into latents, and the diffusion model learns to generate inside that space. This split creates a tension the paper names directly: latents must preserve enough detail for faithful reconstruction, yet stay organized so a downstream diffusion model can learn them quickly. That second property is called diffusability. High-capacity, weakly regularized latents tend to reconstruct well and model poorly.
Two recent lines attack this. One adds semantic supervision during VAE training; VA-VAE aligns latents with pretrained visual features. The other jointly trains the tokenizer with a diffusion model, with REPA-E as the representative. REPA-E has an awkward detail: its flow-matching loss is gradient-detached from the VAE encoder, because those gradients were found to collapse latent variance. Only the semantic alignment term actually updates the encoder. So what is joint diffusion training buying? That question starts the paper.
The authors subtract first. RobustAlign deletes REPA-E's diffusion objective and its generative blocks, keeping one pathway: predict frozen DINOv2-B features from noise-corrupted latents. Under matched VAE architecture and matched downstream training, RobustAlign-VAE scores 60.89 on GenEval after 32k steps against 60.60 for the REPA-E tokenizer. Jointly training a full diffusion model for the encoder contributes roughly nothing.
That result becomes a general principle: generation-relevant information should stay predictable from heavily corrupted latents. RoAR-VAE implements it with two objectives.
Both backpropagate through the corrupted latent into the encoder, while the standard VAE loss keeps operating on clean latents. The final objective is LVAE + 1.5·RobustAlign + 0.5·RobustRecon.
The proposed account of why this works: under isotropic Gaussian corruption, the signal-to-noise ratio along the i-th covariance eigenvector is (1-t)²νi/t², so high-variance directions stay cleaner. Robust prediction pushes useful information into those dominant modes, and prior analysis (Wang & Pehlevan, 2025) shows diffusion models learn high-variance modes earlier. PCA visualizations agree: RobustAlign latents show clearer object-level organization, RobustRecon exposes richer local variation, and the combination has both.
The tokenizer is an 84M-parameter convolutional autoencoder, 8x downsampling, 32 channels, trained only on ImageNet at 256² (200k-step VAE warm-up, then 250k steps of RoAR refinement). For evaluation each tokenizer is frozen and the same roughly 0.8B MMDiT is trained from scratch on filtered, recaptioned LAION; benchmarks are measured after 32k updates.
| Method | Reconstruction LPIPS↓ | Generation avg (3 benchmarks) |
| FLUX.1-VAE | 0.0510 | 52.90 |
| VA-VAE (strongest baseline) | 0.1052 | 61.67 |
| REPA-E-FLUX | 0.0546 | 61.45 |
| FLUX.2-VAE (reference only) | 0.0398 | 62.73 |
| RoAR-VAE | 0.0349 | 63.86 |
That is 2.19 points over the strongest baseline under the same training budget, with better reconstruction on top: 32.88 dB PSNR (1.80 dB above FLUX.1-VAE) and 0.1309 rFID. The fine-grained gaps are larger. On TokBench, small-text accuracy is 69.00% against 52.07% for FLUX.1-VAE, a 16.93-point margin, and small-face similarity rises from 0.53 to 0.65. Both slightly exceed FLUX.2-VAE (67.33%, 0.63).
Ablations locate the active ingredient. Vanilla-KL averages 50.50; RobustRecon alone reaches 59.38, RobustAlign alone 62.49, both together 63.86. Swap the corrupted latents for clean ones in both objectives and the score falls to 58.47. The corruption itself is what matters. Replacing the Transformer semantic predictor with a simple MLP drops RobustAlign-VAE from 62.49 to 52.70, so the prediction pathway needs capacity.
Tokenizer quality caps both training efficiency and detail fidelity in latent diffusion, and this space has been crowded lately. RoAR contributes two things. It shows the expensive part of REPA-E, jointly training a generator alongside the VAE, can be deleted entirely: no diffusion model to maintain during tokenizer training, and latents become cacheable again, which matters at scale. It also breaks the usual trade-off, where learnability methods wreck reconstruction (VA-VAE's LPIPS is three times worse) and reconstruction-oriented tokenizers learn slowly downstream. Code and models are released; anyone building a T2I system can drop this in.
The paper labels its own mechanistic account, the SNR argument plus spectral bias, a possible explanation, supported by qualitative PCA plots rather than causal verification. Training and evaluation stay at moderate scale: ImageNet at 256² for the tokenizer, a 0.8B downstream model, 32k-step comparisons. Whether the 2.19-point margin survives larger models, longer training, or higher resolution is untested; Figure 4 shows 120k-step qualitative samples but no full benchmark numbers. FLUX.2-VAE, trained on undisclosed data, already sits at 62.73, so industrial recipes may not be far behind. DINOv2 is the only semantic teacher tested. In absolute terms the margin is modest; this is a clean, well-explained incremental improvement.