Source-linked AI summary
Boosting Latent Diffusion Models via Disentangled Representation Alignment
John Page, Xuesong Niu, Kai Wu, Kun Gai
TL;DR
The paper addresses the mismatch between LDM-oriented VFM alignment and the fine-grained information-preservation needs of VAEs. It proposes Send-VAE, which uses non-linear mapping to bridge VAE and VFM representations, and reports state-of-the-art ImageNet generation quality with 1.21 FID.
Problem
Prior VAE alignment methods reuse LDM alignment targets despite VAEs needing to preserve fine-grained, disentangled visual structure.
Method
Send-VAE uses a non-linear mapping architecture to distill VFM semantics while preserving the VAE’s inherent structural integrity.
Results
Send-VAE achieves a state-of-the-art 1.21 FID with classifier-free guidance on ImageNet 256 × 256 generation.
Takeaways & Limitations
Linear probing of low-level attribute separability provides an intrinsic way to evaluate VAE latent quality.
Takeaways & Limitations
The Send-VAE latent space prevents capturing excessive fine-grained low-level details.
Abstract
from arXiv · showhide
Latent Diffusion Models (LDMs) rely heavily on the compressed latent space provided by Variational Autoencoders (VAEs) for high-quality image generation. Recent studies have attempted to obtain generation-friendly VAEs by directly adopting alignment strategies from LDM training, leveraging Vision Foundation Models (VFMs) as representation alignment targets. However, such alignment paradigms overlook the fundamental differences in representational requirements between LDMs and VAEs. Simple feature mapping from local patches to high-dimensional semantics can induce semantic collapse, leading to the loss of fine-grained attributes. In this paper, we reveal a key insight: unlike LDMs that benefit from high-level global semantics, a generation-friendly VAE must possess strong semantic disentanglement capabilities to preserve fine-grained, attribute-level information in a structured manner. To address this discrepancy, we propose the Semantic-Disentangled VAE (Send-VAE). Deviating from previous shallow alignment approaches, Send-VAE introduces a non-linear mapping architecture to effectively bridge the local structures of VAEs and the dense semantics of VFMs, thereby encouraging emergent disentangled properties in the latent space without explicit regularization. Extensive experiments establish a new paradigm for evaluating VAE latent spaces via low-level attribute separability and demonstrate that Send-VAE achieves state-of-the-art generation quality (FID of 1.21) on ImageNet 256x256.
1 Introduction
The paper argues that generation-friendly VAEs require semantic disentanglement rather than direct inheritance of LDM alignment targets, and introduces Send-VAE to preserve fine-grained structure while leveraging VFM semantics.
- VAEs are critical image tokenizers for LDMs because their compressed latent spaces reduce computational demands for high-resolution generation.
- Directly aligning VAE latents with VFM representations overlooks the different semantic requirements of VAEs and LDMs.The paper identifies this shared-target assumption as a conceptual flaw in prior alignment approaches.
- Linear separability of low-level attributes is presented as an intrinsic metric for evaluating VAE latent quality and generation friendliness.The metric is motivated by experiments linking attribute prediction with downstream generation quality.
- Send-VAE uses a non-linear mapper to align VAE representations with VFM semantics while preserving structured semantics and encouraging disentanglement without explicit regularization.Its design bridges VAE local structures and VFM dense semantics.
- 1.21 FID is achieved with classifier-free guidance on ImageNet 256 × 256 generation, establishing a state-of-the-art result.The paper also reports a 1.75 FID result without classifier-free guidance.
2 Related work
Related work situates VAEs as continuous image tokenizers for diffusion models and reviews alignment methods that inject VFM semantics, while noting that these approaches overlook the VAE representation bottleneck.
- VAEs are continuous image tokenizers that provide compact latent representations for diffusion-based generation.
- Prior methods align VAE latents with VFM semantics or use multiple target features to enrich tokenizer representations.Examples include VA-VAE and MAETok, alongside related strategies for discrete tokenizers.
- Existing alignment and regularization approaches can improve generation but often treat the VAE latent space as a static representation without addressing its bottleneck.The paper redirects attention from denoising-centric alignment to semantic disentanglement within the VAE.
3 Method
This section rethinks VAE latent-space evaluation around semantic disentanglement and introduces Send-VAE, which uses non-linear VFM alignment to preserve structured semantics. The method evaluates low-level attribute separability, bridges VAE and VFM representations, and retains reconstruction objectives during training.
- Rethinking VAE Evaluation: Existing latent-space metrics do not consistently characterize downstream generative performance across VAEs.Experiments use semantic gap, uniformity, and discrimination evaluations, but their relationship with generation quality is inconsistent.
- Rethinking VAE Evaluation: Low-level attribute separability shows a positive correlation with generation performance, supporting semantic disentanglement as a key VAE property.The analysis uses linear probing on attribute prediction tasks, including Top5 Recall on DeepFashion.
- Semantic-Disentangled VAE: Send-VAE uses a non-linear mapper to bridge VAE local structures and dense VFM semantics while preserving structured latent representations.The mapper replaces direct shallow alignment and is designed as a semantic buffer for contextual knowledge.
- Semantic-Disentangled VAE: Send-VAE aligns noisy VAE latents with VFM features rather than aligning only clean latent codes.Gaussian noise is injected into the VAE latent to simulate downstream diffusion-model inputs, and the mapper output is compared patch-wise with VFM features.
- Semantic-Disentangled VAE: The training objective combines the alignment loss with the original VAE objective to retain reconstruction capabilities.The original objective includes reconstruction, adversarial GAN, and KL-divergence losses, with θ and ϕ denoting trainable VAE and mapper parameters.
4 Experiments
Experiments on ImageNet 256×256 evaluate Send-VAE’s design, generation quality, convergence, reconstruction trade-offs, and robustness across configurations. Send-VAE improves downstream generation and accelerates diffusion training, while slightly sacrificing reconstruction fidelity.
- System-level comparison: Send-VAE achieves gFID scores of 2.88 without CFG and 1.41 with CFG after 80 training epochs.The comparison is reported for ImageNet 256×256 generation.
- System-level comparison: Send-VAE reaches generation FIDs of 1.21 with CFG and 1.75 without CFG, outperforming E2E-VAE on ImageNet 256×256.The results use the same DINOv2 foundation model across the compared VAEs.
- Reconstruction analysis: Send-VAE’s reconstruction performance is slightly inferior to VA-VAE because its disentangled latent space captures fewer fine-grained low-level details.Across reconstruction metrics, E2E-VAE and Send-VAE have similar profiles around PSNR≈27.6 and SSIM≈0.77, slightly below Naive VAE.
- System-level comparison: Send-VAE accelerates diffusion convergence, producing more structurally meaningful images during early training than VA-VAE and E2E-VAE.Qualitative comparisons use matched labels and initial noise across checkpoints trained for 10, 20, and 80 epochs.
- Ablation studies: A one-layer ViT mapper gives the best ablation result at gFID=8.42, while both shallower and deeper configurations perform worse.The paper attributes this balance to sufficient representation bridging without excessively weakening foundation-model influence.
- Ablation studies: Send-VAE improves generation across VAE specifications and at 512×512 resolution, while noise injection and DINO-family features further support disentangled representations.Attribute-prediction performance is positively correlated with downstream generation across CelebA, DeepFashion, and AwA.
5 Conclusion
The paper argues that generation-friendly VAEs require semantic disentanglement rather than identical alignment targets to LDMs. Send-VAE uses nonlinear VFM alignment and achieves strong ImageNet generation results while preserving structured semantics.
- Conclusion: VAEs require semantic disentanglement to preserve fine-grained structured visual elements, whereas LDMs benefit from abstract semantics.The paper supports this distinction through a correlation between low-level attribute separability and downstream generation quality.
- Conclusion: Send-VAE uses a nonlinear mapper to absorb contextual VFM guidance while safeguarding the VAE’s native structured semantics.This design bridges the representation gap without forcing rigid shallow alignment.
- Conclusion: Send-VAE demonstrates empirical superiority on ImageNet 256×256 generation in the paper’s extensive experiments.The conclusion presents this as evidence for the proposed generation-friendly VAE design.