Robust prediction improves image tokenizers for downstream diffusion models
RoAR-VAE: Improving Latent Diffusability through Robust Prediction
Computer Vision and Pattern Recognition
Summary
Generating high-quality images using AI often involves turning images into a simpler form called latents. These latents need to keep important details and be easy for later models to learn from. The authors found that making these latents recoverable even when slightly corrupted helps models generate better images. They created a method called RoAR-VAE that uses two tasks: predicting important features and reconstructing the original image from corrupted latents. This approach strikes a good balance between keeping image details and making latents easier to use for image generation.
What this means in practice
- •For image generation teams: Improve image generation quality by using RoAR-VAE latents that balance detail preservation and ease of learning for diffusion models.
- •For computer vision engineers: Build better image compression or representation systems that maintain important image details while supporting downstream generative tasks.
Authors
Yang Yue, Fangyun Wei, Yitong Wang, Jinjing Zhao, Henry Zheng, Dong Chen, Gao Huang
Abstract
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.