cs.CVMar 22, 2026

Taming Sampling Perturbations with Variance Expansion Loss for Latent Diffusion Models

Authors: Qifan Li, Xingyu Zhou, Jinhua Zhang, Weiyi You, Shuhang Gu

Organizations: University of Electronic Science and Technology of China

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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 β\beta -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).
TokenizerEpochs.Spec.Reconstruction PerformanceGeneration Performance (FID-10K) ↓\downarrow
rFID ↓\downarrowPSNR ↑\uparrowLPIPS ↓\downarrowSSIM ↑\uparrowDiT-BLightningDiT-BLightningDiT-B +
LDM10f16d160.5526.050.1410.71031.9322.25-
LDM+VE loss100.6025.230.1470.69129.0319.70-
VAVAE16f16d320.3527.430.1040.7722.2719.85-
VAVAE500.2827.960.0960.79--15.82
VAVAE+VE loss160.4526.540.1180.7419.4215.5012.89
Table 1: Comprehensive comparisons show that the VE loss consistently improves generative performance across architectures while maintaining competitive reconstruction quality. ↓\downarrow and ↑\uparrow indicate whether lower or higher values are better, respectively. ( LightningDiT-B + denotes the experiment with training extended to 160 epochs.)
Methodσ2\bm{\sigma^{2}}ReconstructionGeneration
rFID ↓\downarrowPSNR ↑\uparrowFID-10K ↓\downarrow
++ KL β=10−6\beta=10^{-6}10−810^{-8}0.3927.1223.12
++ KL β=10−4\beta=10^{-4}10−710^{-7}0.3927.0723.03
++ KL β=10−2\beta=10^{-2}10−510^{-5}0.4426.7122.87
++ KL β=10−1\beta=10^{-1}10−310^{-3}0.5026.2323.16
++ KL β=0.5\beta=0.510−210^{-2}0.5226.1622.99
Table 2: System-level comparison on ImageNet 256×256 without classifier-free guidance (CFG). ↓\downarrow and ↑\uparrow indicate whether lower or higher values are better, respectively.
Figure 3: Visualization Results. Examples of class-conditional generation on ImageNet 256 ×\times 256
MethodTraining Epoches#paramsGeneration w/ CFG
gFID ↓\downarrowsFID ↓\downarrowIS ↑\uparrowPre. ↑\uparrowRec. ↑\uparrow
LlamaGen [ 23 ]3003.1B2.185.97263.30.810.58
VAR [ 25 ]3502.0B1.80-365.40.830.57
MVAR [ 31 ]-1.0B2.155.62298.90.840.56
MAR [ 16 ]800945M1.55-303.70.810.62
MaskDiT [ 34 ]1600675M2.285.67276.60.800.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. ↓\downarrow and ↑\uparrow indicate whether lower or higher values are better, respectively.
MethodTraining EpochesReconstruction Performance
rFID ↓\downarrowPSNR ↑\uparrowLPIPS ↓\downarrowSSIM ↑\uparrow
VA-VAE [ 29 ]1300.2827.710.0970.779
++ VE Loss++ 100.2628.310.0900.792
Table 4: Reconstruction performance of our finetuned model on ImageNet 256×256. ↓\downarrow and ↑\uparrow 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 ×\times 256
Figure 5: Visualization Results. Examples of class-conditional generation on ImageNet 256 ×\times 256
More in this topic

Latent-Space Regularization

See all 15 papers

Explore similar work

CardsList

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

    Jun 2, 2026Tianxiong Zhong, Xingye Tian, Xuebo Wang +2Latent Diffusion ModelsImage Tokenization

    Kling Team, Kuaishou Technology

  2. RoAR-VAE: Improving Latent Diffusability through Robust Prediction

    Oct 19, 2026Yang Yue, Fangyun Wei, Yitong Wang +4Image TokenizationVariational Autoencoders

    Tsinghua University · Microsoft Research

  3. What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion

    May 8, 2026Zhengrong Yue, Taihang Hu, Mengting Chen +8Latent Diffusion ModelsAutoencoders

    Shanghai Jiao Tong University · Alibaba Group · Beihang University +4