LiteVAE : Autoencodeurs variationnels légers et efficaces pour les modèles de diffusion latente

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

Résumé

Les avancées dans les modèles de diffusion latente (LDMs) ont révolutionné la génération d'images haute résolution, mais l'espace de conception de l'autoencodeur central à ces systèmes reste peu exploré. Dans cet article, nous présentons LiteVAE, une famille d'autoencodeurs pour LDMs qui exploite la transformée en ondelettes discrète 2D pour améliorer l'évolutivité et l'efficacité computationnelle par rapport aux autoencodeurs variationnels (VAEs) standards, sans sacrifier la qualité de sortie. Nous étudions également les méthodologies d'entraînement et l'architecture du décodeur de LiteVAE, et proposons plusieurs améliorations qui optimisent la dynamique d'entraînement et la qualité de reconstruction. Notre modèle LiteVAE de base atteint la qualité des VAEs établis dans les LDMs actuels avec une réduction de six fois des paramètres de l'encodeur, permettant un entraînement plus rapide et des besoins en mémoire GPU réduits, tandis que notre modèle plus large surpasse les VAEs de complexité comparable sur toutes les métriques évaluées (rFID, LPIPS, PSNR et 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).