LiteVAE: Lichtgewicht en efficiënte variational autoencoders voor latente diffusiemodellen

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

May 23, 2024

Auteurs: Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, Romann M. Weber

cs.AI

Samenvatting

Vooruitgang in latente diffusiemodellen (LDMs) heeft een revolutie teweeggebracht in het genereren van hoogwaardige afbeeldingen, maar de ontwerpruimte van de auto-encoder die centraal staat in deze systemen blijft onderbelicht. In dit artikel introduceren we LiteVAE, een familie van auto-encoders voor LDMs die gebruikmaken van de 2D discrete wavelet transformatie om de schaalbaarheid en rekenkundige efficiëntie te verbeteren ten opzichte van standaard variational autoencoders (VAEs) zonder in te leveren op uitvoerkwaliteit. We onderzoeken ook de trainingsmethodologieën en de decoderarchitectuur van LiteVAE en stellen verschillende verbeteringen voor die de trainingsdynamiek en reconstructiekwaliteit verhogen. Onze basis LiteVAE-model evenaart de kwaliteit van de gevestigde VAEs in huidige LDMs met een zesvoudige reductie in encoderparameters, wat leidt tot snellere training en lagere GPU-geheugenvereisten, terwijl ons grotere model VAEs van vergelijkbare complexiteit overtreft op alle geëvalueerde metrieken (rFID, LPIPS, PSNR en SSIM).

English

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