Latent Diffusion Models with Masked AutoEncoders

TL;DR

This work analyzes autoencoder design choices for Latent Diffusion Models and identifies three critical properties—latent space smoothness, perceptual compression, and reconstruction quality—that jointly influence generation performance. It proposes Variational Masked AutoEncoders (VMAEs) that combine probabilistic encoding with hierarchical, masked autoencoding to realize these properties, and integrates them into Latent Diffusion Models as LDMAEs. Through comprehensive experiments on ImageNet-1K and CelebA-HQ, VMAEs demonstrate superior latent-space behavior, perceptual reconstruction, and overall generation quality, while also offering considerable efficiency gains. The results establish VMAEs as a principled autoencoder design for LDMs with practical implications for scalable, high-fidelity image synthesis.

Abstract

In spite of the remarkable potential of Latent Diffusion Models (LDMs) in image generation, the desired properties and optimal design of the autoencoders have been underexplored. In this work, we analyze the role of autoencoders in LDMs and identify three key properties: latent smoothness, perceptual compression quality, and reconstruction quality. We demonstrate that existing autoencoders fail to simultaneously satisfy all three properties, and propose Variational Masked AutoEncoders (VMAEs), taking advantage of the hierarchical features maintained by Masked AutoEncoders. We integrate VMAEs into the LDM framework, introducing Latent Diffusion Models with Masked AutoEncoders (LDMAEs). Our code is available at https://github.com/isno0907/ldmae.

Figures (19)

• Figure 1: Class-conditional generation samples of our proposed LDMAE. Trained on ImageNet-1K, resolution of $256 \times 256$.
• Figure 2: Comparison of autoencoding methods across key metrics: latent smoothness, perceptual compression quality, reconstruction quality (at pixel and perceptual levels), and overall generation quality. Each score is linearly rescaled so that the outermost and innermost grids indicate the highest and lowest scores, respectively.
• Figure 3: Sparse latent points vs. smooth latent space. Deterministic decoders allow only the exact points to be decodable due to its sparse latent points. In contrast, probabilistic decoders, with a smooth latent space, ensure that all latent points in the vicinity of an encoded point remain within the decodable space.
• Figure 4: Separability of latent features. Our VMAE achieves a more highly separable latent space while maintaining strong semantic grouping, indicating hierarchical compression. In contrast, SD-VAE primarily exhibits semantic clustering without clearly differentiating features within each cluster.
• Figure 5: Key factors for reconstruction quality. Although all autoencoders achieve a comparable pixel-level reconstruction (in MSE), they differ in perceptual details especially within the boxed region (in rFID), suggesting that the autoencoders should be optimized for both pixel-level accuracy and perceptual restoration.