LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models (2405.14477v1)

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).

• The paper presents LiteVAE, a novel VAE that leverages 2D discrete wavelet transform to reduce encoder parameters by six-fold for efficient latent diffusion.
• Methodologically, LiteVAE processes images using multi-level DWT and self-modulated convolution to extract robust multi-scale features for high-quality reconstructions.
• LiteVAE improves scalability and training by enabling low-resolution pretraining and incorporating a UNet-based adversarial setup for fine-grained pixel-wise discrimination.

LiteVAE: Efficient Autoencoders for Latent Diffusion Models

Introduction

Large Diffusion Models (LDMs) have been making waves in the world of high-resolution image generation. These models are robust, scalable, and generate high-quality images from latent representations. But there's an under-explored element in this system: the autoencoder, specifically the Variational Autoencoder (VAE) component. The paper we're discussing introduces LiteVAE, a family of autoencoders designed to enhance both scalability and computational efficiency without sacrificing the quality of the generated images.

The Core Innovation: LiteVAE

LiteVAE uses the 2D discrete wavelet transform (DWT) to process images before feeding them into a simplified, lightweight encoder. Let's break down what makes this approach unique:

1. Wavelet Processing: An image is first processed using a multi-level DWT to extract multi-scale features from the image.
2. Feature Extraction and Aggregation: The extracted features are then processed via specialized modules that combine them into a unified latent code.
3. Efficient Training and Decoding: The encoded latent representation is then used to reconstruct the image through a decoder network, ensuring high-quality output.

Computational Efficiency

One of the standout features of LiteVAE is its computational efficiency. The base model of LiteVAE achieves similar reconstruction quality to established VAEs in LDMs but with a six-fold reduction in encoder parameters. This means faster training times and lower GPU memory requirements. For instance, in the experiments, LiteVAE’s base model requires significantly less computational cost for training autoencoders and offers higher throughput during the latent diffusion model training.

Performance and Scalability

The paper showcases that LiteVAE doesn't just match but often exceeds the performance of standard VAEs across various datasets, including popular benchmarks like FFHQ and CelebA-HQ. Here are some key points from their results:

• Maintaining Quality with Fewer Parameters: LiteVAE maintains high reconstruction quality with significantly fewer encoder parameters compared to standard VAEs.
• Scalability: When increasing the complexity of the feature-extraction and feature-aggregation blocks, LiteVAE's performance improves, surpassing standard VAEs of equivalent complexity.
• Training Resolution: Most of LiteVAE's training can be effectively conducted at a lower resolution (e.g., 128x128) and then fine-tuned at full resolution (256x256). This reduced overall training compute substantially.

Enhancements in Training Dynamics

The authors introduced several training enhancements to boost the performance and quality of LiteVAE:

1. Self-Modulated Convolution (SMC): This avoids imbalanced feature maps by using learnable parameters to modulate convolutional weights, improving final reconstruction quality.
2. Improved Adversarial Setup: Replacing the PatchGAN discriminator with a UNet-based model for more fine-grained, pixel-wise discrimination.
3. Training Resolution: Pretraining at lower resolutions speeds up the learning process without compromising on final image quality.
4. High-Frequency Loss Terms: Introducing wavelet-based and Gaussian-filter-based loss terms ensures that the model captures and reconstructs high-frequency details better.

Practical and Theoretical Implications

• Practical Use: Given the reduced computational cost, LiteVAE can make high-resolution image generation more accessible, especially in resource-constrained environments. This opens up new possibilities for applications in generative art, medical imaging, and even video generation.
• Future Directions: The notion of simplifying complex encoder networks with pre-processing steps like wavelet transforms can inspire other fields. Further research can explore combining LiteVAE with other feature extraction techniques, such as Fast Fourier Transforms (FFT).

Conclusion

LiteVAE brings an efficient, scalable approach to the VAE component of Latent Diffusion Models. By leveraging the DWT and introducing thoughtful enhancements to both architecture and training processes, LiteVAE manages to enhance computational efficiency significantly while maintaining, or even improving, image reconstruction quality.

This streamlined autoencoder design could set the stage for even more sophisticated models and techniques, potentially broadening the scope and accessibility of high-resolution image generative models in the future.