Latent diffusion models have emerged as the dominant framework for high-fidelity and efficient image generation, owing to their ability to learn diffusion processes in compact latent spaces. However, while previous research has focused primarily on reconstruction accuracy and semantic alignment of the latent space, we observe that another critical factor, robustness to sampling perturbations, also plays a crucial role in determining generation quality. Through empirical and theoretical analyses, we show that the commonly used β-VAE-based tokenizers in latent diffusion models, tend to produce overly compact latent manifolds that are highly sensitive to stochastic perturbations during diffusion sampling, leading to visual degradation. To address this issue, we propose a simple yet effective solution that constructs a latent space robust to sampling perturbations while maintaining strong reconstruction fidelity. This is achieved by introducing a Variance Expansion loss that counteracts variance collapse and leverages the adversarial interplay between reconstruction and variance expansion to achieve an adaptive balance that preserves reconstruction accuracy while improving robustness to stochastic sampling. Extensive experiments demonstrate that our approach consistently enhances generation quality across different latent diffusion architectures, confirming that robustness in latent space is a key missing ingredient for stable and faithful diffusion sampling.
Figures & tables
Figure 1: Toy example on a fractal-like 2D distribution following [ 12 ] . In this toy example, we observe a seemingly counterintuitive phenomenon: a tokenizer with better reconstruction and a diffusion model trained in its latent space that achieves lower diffusion loss still produces visually inferior generations (left). To understand this behavior, we visualize the learned latent spaces. The latent space on the left, produced by a commonly used β -VAE tokenizer, collapses into an overly compact region, which allows even small diffusion sampling perturbations to push samples outside the data manifold. In contrast, the latent space on the right, learned with our Variance Expansion loss, remains sufficiently spread out to resist stochastic perturbations during diffusion sampling, leading to much more faithful generations.
Figure 2: Toy example on a fractal-like 2D distribution following [ 12 ] . KL regularization severely degrades reconstruction quality (left), whereas our method maintains high reconstruction quality (right).
Tokenizer
Epochs.
Spec.
Reconstruction Performance
Generation Performance (FID-10K) ↓
rFID ↓
PSNR ↑
LPIPS ↓
SSIM ↑
DiT-B
LightningDiT-B
LightningDiT-B +
LDM
10
f16d16
0.55
26.05
0.141
0.710
31.93
22.25
-
LDM+VE loss
10
0.60
25.23
0.147
0.691
29.03
19.70
-
VAVAE
16
f16d32
0.35
27.43
0.104
0.77
22.27
19.85
-
VAVAE
50
0.28
27.96
0.096
0.79
-
-
15.82
VAVAE+VE loss
16
0.45
26.54
0.118
0.74
19.42
15.50
12.89
Table 1: Comprehensive comparisons show that the VE loss consistently improves generative performance across architectures while maintaining competitive reconstruction quality. ↓ and ↑ indicate whether lower or higher values are better, respectively. ( LightningDiT-B + denotes the experiment with training extended to 160 epochs.)
Method
σ2
Reconstruction
Generation
rFID ↓
PSNR ↑
FID-10K ↓
+ KL β=10−6
10−8
0.39
27.12
23.12
+ KL β=10−4
10−7
0.39
27.07
23.03
+ KL β=10−2
10−5
0.44
26.71
22.87
+ KL β=10−1
10−3
0.50
26.23
23.16
+ KL β=0.5
10−2
0.52
26.16
22.99
Table 2: System-level comparison on ImageNet 256×256 without classifier-free guidance (CFG). ↓ and ↑ indicate whether lower or higher values are better, respectively.
Figure 3: Visualization Results. Examples of class-conditional generation on ImageNet 256 × 256
Method
Training Epoches
#params
Generation w/ CFG
gFID ↓
sFID ↓
IS ↑
Pre. ↑
Rec. ↑
LlamaGen [ 23 ]
300
3.1B
2.18
5.97
263.3
0.81
0.58
VAR [ 25 ]
350
2.0B
1.80
-
365.4
0.83
0.57
MVAR [ 31 ]
-
1.0B
2.15
5.62
298.9
0.84
0.56
MAR [ 16 ]
800
945M
1.55
-
303.7
0.81
0.62
MaskDiT [ 34 ]
1600
675M
2.28
5.67
276.6
0.80
0.61
Table 3: System-level comparison on ImageNet 256×256 with guidance. All models use classifier-free guidance (CFG) except RAE, which uses auto-guidance. ↓ and ↑ indicate whether lower or higher values are better, respectively.
Method
Training Epoches
Reconstruction Performance
rFID ↓
PSNR ↑
LPIPS ↓
SSIM ↑
VA-VAE [ 29 ]
130
0.28
27.71
0.097
0.779
+ VE Loss
+ 10
0.26
28.31
0.090
0.792
Table 4: Reconstruction performance of our finetuned model on ImageNet 256×256. ↓ and ↑ indicate whether lower or higher values are better, respectively.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Visualization Results. Examples of class-conditional generation on ImageNet 256 × 256
Figure 5: Visualization Results. Examples of class-conditional generation on ImageNet 256 × 256
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.
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.
Tokenizers are a crucial component of latent diffusion models, as they define the latent space in which diffusion models operate. However, existing tokenizers are primarily designed to improve reconstruction fidelity or inherit pretrained representations, leaving unclear what kind of latent space is truly friendly for generative modeling. In this paper, we study this question from the perspective of latent manifold organization. By constructing controlled tokenizer variants, we identify three key properties of a diffusion-friendly latent manifold: coherent spatial structure, local manifold continuity, and global manifold semantics. We find that these properties are more consistent with downstream generation quality than reconstruction fidelity. Motivated by this finding, we propose the Prior-Aligned AutoEncoder (PAE), which explicitly shapes the latent manifold instead of leaving diffusion-friendly manifold to emerge indirectly from reconstruction or inheritance. Specifically, PAE leverages refined priors derived from VFMs and perturbation-based regularization to turn spatial structure, local continuity, and global semantics into explicit training objectives. On ImageNet 256x256, PAE improves both training efficiency and generation quality over existing tokenizers, reaching performance comparable to RAE with up to 13x faster convergence under the same training setup and achieving a new state-of-the-art gFID of 1.03. These results highlight the importance of organizing the latent manifold for latent diffusion models.
Zhengrong Yue, Taihang Hu, Mengting Chen +8
Shanghai Jiao Tong University · Alibaba Group · Beihang University +4