LiteVAE: Leichte und effiziente Variational Autoencoder für latente Diffusionsmodelle

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

May 23, 2024

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

cs.AI

Zusammenfassung

Fortschritte in latenten Diffusionsmodellen (LDMs) haben die Generierung hochauflösender Bilder revolutioniert, aber der Gestaltungsraum des Autoencoders, der für diese Systeme zentral ist, bleibt untererforscht. In diesem Artikel stellen wir LiteVAE vor, eine Familie von Autoencodern für LDMs, die die 2D-diskrete Wavelet-Transformation nutzen, um die Skalierbarkeit und Rechenleistung im Vergleich zu herkömmlichen Variationalen Autoencodern (VAEs) zu verbessern, ohne dabei die Ausgabequalität zu beeinträchtigen. Wir untersuchen auch die Trainingsmethoden und die Decoder-Architektur von LiteVAE und schlagen mehrere Verbesserungen vor, die die Trainingsdynamik und die Rekonstruktionsqualität verbessern. Unser Basis-LiteVAE-Modell erreicht die Qualität der etablierten VAEs in aktuellen LDMs bei einer sechsfachen Reduzierung der Encoder-Parameter, was zu schnellerem Training und geringeren GPU-Speicheranforderungen führt, während unser größeres Modell VAEs vergleichbarer Komplexität in allen bewerteten Metriken (rFID, LPIPS, PSNR und SSIM) übertrifft.

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