RoAR-VAE: Improving Latent Diffusability through Robust Prediction

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

cs.CVOct 19, 2026

RoAR-VAE: Improving Latent Diffusability through Robust Prediction

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

Organizations: Tsinghua University · Microsoft Research

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.

Figures & tables

Figure 1: Reconstruction versus generation. RoAR-VAE achieves a favorable trade-off, combining low reconstruction LPIPS with strong diffusability.

Figure 2: From REPA-E to RoAR. (a) REPA-E jointly trains the tokenizer and diffusion model, but only semantic alignment backpropagates through corrupted latents to the tokenizer encoder. (b) RobustAlign-VAE isolates this semantic-prediction pathway by removing the diffusion objective and its associated modules. (c) RoAR further introduces RobustRecon, combining semantic and image reconstruction targets under latent corruption alongside the standard VAE objective on clean latents.

Method	4k	8k	16k	32k
T2I modeling w. Vanilla-KL-VAE	11.74 $\pm$ 0.30	23.34 $\pm$ 0.57	35.47 $\pm$ 0.53	47.60 $\pm$ 0.31
T2I modeling w. REPA-E-Tokenizer	26.33 $\pm$ 0.25	36.42 $\pm$ 0.56	49.46 $\pm$ 0.31	60.60 $\pm$ 0.90
T2I modeling w. RobustAlign-VAE	24.56 $\pm$ 0.27	39.47 $\pm$ 0.70	51.12 $\pm$ 0.35	60.89 $\pm$ 0.46

Table 1: Controlled comparison with REPA-E. We train Vanilla-KL-VAE, REPA-E-Tokenizer, and RobustAlign-VAE using the same VAE architecture. For each tokenizer, we train a separate downstream text-to-image diffusion Transformer using the same backbone and training data, and report GenEval ( Ghosh et al., 2023 ) performance after 4k, 8k, 16k, and 32k downstream optimization steps. Both REPA-E and RobustAlign substantially improve over the reconstruction-only Vanilla-KL baseline. Results are reported as mean

\pm

standard deviation over four sampling seeds.

Figure 3: PCA visualization of latent representations. For the input images in (a), we visualize the top three principal components of normalized spatial latent features from (b) KL-VAE, (c) RobustAlign-VAE, (d) RobustRecon-VAE, and (e) RoAR-VAE. RobustAlign emphasizes object-level organization, RobustRecon preserves richer local variation, and RoAR combines both properties. PCA is fitted jointly across the four images for each tokenizer.

Method	Reconstruction				Generation
Method	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	GenEval	GenEval2	DPG-Bench	Avg.
SD3-VAE	0.2038	29.58	0.0688	0.8565	52.46	38.03	73.27	54.59
FLUX.1-VAE	0.1711	31.08	0.0510	0.8938	51.79	34.82	72.08	52.90
VA-VAE	0.2892	26.26	0.1052	0.7625	57.43	51.28	76.30	61.67
VTP	0.3457	25.75	0.1339	0.7383	53.79	50.31	76.08	60.06
PAE	0.2643	25.90	0.1123	0.7487	56.60	49.77	76.82	61.06

Table 2: Reconstruction and downstream generation results. We report reconstruction metrics evaluated on the ImageNet validation set at

256^{2}

resolution. Generation performance is evaluated at 32k downstream updates across three benchmarks: GenEval, GenEval2, and DPG-Bench. Avg. denotes the unweighted mean of the three generation scores.

Figure 4: Qualitative comparison of text-to-image generation. Rows compare FLUX.1-VAE, VA-VAE, and RoAR-VAE tokenizers; each subject is shown after 32k and 120k downstream updates.

Method	Text (%) $\uparrow$				Face $\uparrow$
Method	T-ACC s	T-ACC m	T-NED s	T-NED m	F-Sim s	F-Sim m
SD3-VAE	37.29	68.32	59.86	81.36	0.43	0.71
FLUX.1-VAE	52.07	76.57	71.41	86.96	0.53	0.77
VA-VAE	7.05	37.15	25.12	56.30	0.22	0.49
VTP	0.14	6.88	5.87	20.03	0.10	0.24
PAE	3.31	28.97	16.73	46.97	0.19	0.45

Table 3: Text and face reconstruction on TokBench, using text accuracy (T-ACC), text normalized edit distance (T-NED), and face similarity (F-Sim), evaluated at an image resolution of

256^{2}

. Subscripts

s

and

m

denote metrics averaged over small-scale and all instances, respectively. Higher is better for every column. Bold denotes the best displayed value.

Figure 5: Image reconstruction visualization. Columns show the original image and reconstructions from baseline tokenizers and RoAR-VAE. Images are selected from TokBench ( Wu et al., 2025 ) .

	Noisy Latent	$\mathcal{L}_{\mathrm{RobustAlign}}$	$\mathcal{L}_{\mathrm{RobustRecon}}$	GenEval	GenEval2	DPG-Bench	Avg.
❶	✗	✗	✗	47.60	33.68	70.22	50.50
❷	✓	✗	✓	57.63	45.20	75.32	59.38 ( $\uparrow$ 8.88)
❸	✓	✓	✗	60.89	51.11	75.48	62.49 ( $\uparrow$ 11.99)
❹	✓	✓	✓	62.22	51.79	77.56	63.86 ( $\uparrow$ 13.36)
❺	✗	✓	✓	57.54	44.01	73.86	58.47 ( $\downarrow$ 5.39)

Table 4: Ablation of robust prediction objectives and latent corruption. “Noisy Latent” indicates whether the auxiliary predictors receive noise-corrupted latents as input. Scores are measured after 32k downstream updates; Avg. denotes the mean over the three benchmarks.

Figure 6: Downstream learning curves. RoAR-VAE consistently outperforms using RobustAlign or RobustRecon alone across different training budgets.

Variant	Avg. Score
RobustAlign w. simple MLP $F_{\eta}$	52.70
RobustAlign-VAE	62.49
RoAR w. shared $G_{\xi}$ and $D_{\psi}$	62.13
RoAR-VAE	63.86

Table 5: Predictor architecture choices.

Image Tokenization

Image tokenization converts images into compact discrete or continuous visual units for efficient generation, compression, and multimodal reasoning.

See all 53 papers

Diffusing in the Right Space: A Systematic Study of Latent Diffusability
Jun 2, 2026Tianxiong Zhong, Xingye Tian, Xuebo Wang +2Latent Diffusion Models Image Tokenization
Kling Team, Kuaishou Technology
DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment
Mar 23, 2026Xin Cai, Zhiyuan You, Zhoutong Zhang +1Diffusion-Based Image Generation Diffusion Model Compression
Multimedia Laboratory, The Chinese University of Hong Kong · Shanghai AI Laboratory · Adobe NextCam +1
RPiAE: A Representation-Pivoted Autoencoder Enhancing Both Image Generation and Editing
Mar 19, 2026Yue Gong, Hongyu Li, Shanyuan Liu +8T2I Diffusion Models Diffusion-Based Image Generation
Beihang University · AI Research · The Chinese University of Hong Kong

Jun 2, 2026J/K move · Enter open · S save

Diffusing in the Right Space: A Systematic Study of Latent Diffusability

Tianxiong Zhong, Xingye Tian, Xuebo Wang, Xin Tao+1

Kling Team, Kuaishou Technology

Latent Diffusion Models Image Tokenization Image Generation Evaluation

Open paper AskPDF arXiv

Latent diffusion models leverage visual tokenizers to compress images into latent spaces for efficient generative modeling. However, better reconstruction quality of a tokenizer does not necessarily translate into better generation quality, suggesting that latent representations should be evaluated not only by fidelity but also by their diffusability. Recent studies have proposed diverse explanations for diffusion-friendly latent spaces, including semantic separability, affine equivariance, distribution uniformity, spatial structure, spectral smoothness, and manifold continuity. Yet these properties are often validated on a limited set of tokenizers, leaving it unclear which factors are most predictive of downstream generation quality and whether such conclusions hold beyond the specific settings in which they are introduced. In this work, we conduct a systematic study of latent diffusability by training a large collection of tokenizers with diverse regularization strategies, architectures, and latent configurations, and evaluating them with multiple downstream diffusion backbones. Our analysis identifies several latent properties that consistently correlate with generation quality and exhibit strong generalization across experimental settings. Beyond existing metrics, we introduce Velocity Irreducible Variance (VIV), a measure of velocity ambiguity induced by trajectory crossings. Extensive experiments show that VIV is one of the most stable predictors of generation quality.

Abstract

Image Tokenization

Explore similar work

Diffusing in the Right Space: A Systematic Study of Latent Diffusability

DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment

RPiAE: A Representation-Pivoted Autoencoder Enhancing Both Image Generation and Editing