LiteVAE: Autoencoder Variazionali Leggeri ed Efficienti per Modelli di Diffusione Latente

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

May 23, 2024

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

cs.AI

Abstract

I progressi nei modelli di diffusione latente (LDMs) hanno rivoluzionato la generazione di immagini ad alta risoluzione, ma lo spazio di progettazione dell'autoencoder, centrale in questi sistemi, rimane ancora poco esplorato. In questo articolo, introduciamo LiteVAE, una famiglia di autoencoder per LDMs che sfrutta la trasformata wavelet discreta 2D per migliorare la scalabilità e l'efficienza computazionale rispetto ai tradizionali autoencoder variazionali (VAEs), senza compromettere la qualità dell'output. Investigiamo inoltre le metodologie di addestramento e l'architettura del decoder di LiteVAE, proponendo diversi miglioramenti che ottimizzano le dinamiche di addestramento e la qualità della ricostruzione. Il nostro modello base LiteVAE eguaglia la qualità dei VAEs consolidati negli attuali LDMs con una riduzione di sei volte nei parametri dell'encoder, portando a un addestramento più veloce e a minori requisiti di memoria GPU, mentre il nostro modello più grande supera i VAEs di complessità comparabile in tutte le metriche valutate (rFID, LPIPS, PSNR e 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).