LiteVAE: Autoencoders Variacionais Leves e Eficientes para Modelos de Difusão Latente

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

May 23, 2024

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

cs.AI

Resumo

Os avanços nos modelos de difusão latente (LDMs) revolucionaram a geração de imagens de alta resolução, mas o espaço de design do autoencoder, que é central para esses sistemas, permanece pouco explorado. Neste artigo, apresentamos o LiteVAE, uma família de autoencoders para LDMs que aproveita a transformada discreta de wavelets 2D para melhorar a escalabilidade e a eficiência computacional em relação aos autoencoders variacionais (VAEs) padrão, sem sacrificar a qualidade da saída. Também investigamos as metodologias de treinamento e a arquitetura do decodificador do LiteVAE, propondo várias melhorias que aprimoram a dinâmica de treinamento e a qualidade da reconstrução. Nosso modelo base LiteVAE iguala a qualidade dos VAEs estabelecidos nos LDMs atuais com uma redução de seis vezes nos parâmetros do codificador, resultando em treinamento mais rápido e menores requisitos de memória GPU, enquanto nosso modelo maior supera VAEs de complexidade comparável em todas as métricas avaliadas (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).