Boosting Latent Diffusion Models via Disentangled Representation Alignment

TL;DR

The Semantic disentangled VAE (Send-VAE) is explicitly optimized for disentangled representation learning through aligning its latent space with the semantic hierarchy of pre-trained VFMs, facilitating effective guidance for VAE learning.

Abstract

Latent Diffusion Models (LDMs) generate high-quality images by operating in a compressed latent space, typically obtained through image tokenizers such as Variational Autoencoders (VAEs). In pursuit of a generation-friendly VAE, recent studies have explored leveraging Vision Foundation Models (VFMs) as representation alignment targets for VAEs, mirroring the approach commonly adopted for LDMs. Although this yields certain performance gains, using the same alignment target for both VAEs and LDMs overlooks their fundamentally different representational requirements. We advocate that while LDMs benefit from latents retaining high-level semantic concepts, VAEs should excel in semantic disentanglement, enabling encoding of attribute-level information in a structured way. To address this, we propose the Semantic disentangled VAE (Send-VAE), explicitly optimized for disentangled representation learning through aligning its latent space with the semantic hierarchy of pre-trained VFMs. Our approach employs a non-linear mapper network to transform VAE latents, aligning them with VFMs to bridge the gap between attribute-level disentanglement and high-level semantics, facilitating effective guidance for VAE learning. We evaluate semantic disentanglement via linear probing on attribute prediction tasks, showing strong correlation with improved generation performance. Finally, using Send-VAE, we train flow-based transformers SiTs; experiments show Send-VAE significantly speeds up training and achieves a state-of-the-art FID of 1.21 and 1.75 with and without classifier-free guidance on ImageNet 256x256.

Figures (4)

• Figure 1: Our Send-VAE enhances VAEs by aligning their latent representations with the semantically rich representations of pre-trained vision foundation models using a specialized mapper network. Unlike the direct alignment methods typically used during diffusion model training, this mapper network efficiently bridges the representation gap, enabling a seamless integration of semantic information. Notably, the usage of Send-VAE results in significantly more efficient and effective training of diffusion models.
• Figure 2: We conduct experiments with three recently proposed evaluation methods for VAE latent space, and show their correlation with down stream generation performance (g-FID). Experimental results on four VAEs with identical specifications indicate that these metrics do not accurately reflect the impact of VAEs on downstream generative performance. Conversely, we find that the ability of VAEs regarding low-level attributes is the key factor.
• Figure 3: Qualitative comparisons among VA-VAE, E2E-VAE, and Send-VAE. Results for both methods are sampled using the same seed, noise and class label. The classifier-free guidance scale is set to 4.0.
• Figure 4: Qualitative Results on ImageNet 256 × 256 using Send-VAE and SiT-XL.