LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent
  Diffusion Models
Authors: name: Seyedmorteza Sadat; name: Jakob Buhmann; name: Derek Bradley; name: Otmar Hilliges; name: Romann M. Weber
Published: 2024-05-23 12:06:00

Abstract: Advances in latent diffusion models (LDMs) have revolutionized
high-resolution image generation, but the design space of the autoencoder that
is central to these systems remains underexplored. In this paper, we introduce
LiteVAE, a family of autoencoders for LDMs that leverage the 2D discrete
wavelet transform to enhance scalability and computational efficiency over
standard variational autoencoders (VAEs) with no sacrifice in output quality.
We also investigate the training methodologies and the decoder architecture of
LiteVAE and propose several enhancements that improve the training dynamics and
reconstruction quality. Our base LiteVAE model matches the quality of the
established VAEs in current LDMs with a six-fold reduction in encoder
parameters, leading to faster training and lower GPU memory requirements, while
our larger model outperforms VAEs of comparable complexity across all evaluated
metrics (rFID, LPIPS, PSNR, and SSIM).