LiteVAE: Autoencoders Variacionales Ligeros y Eficientes para Modelos de Difusión 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

Resumen

Los avances en los modelos de difusión latente (LDMs, por sus siglas en inglés) han revolucionado la generación de imágenes de alta resolución, pero el espacio de diseño del autoencoder, que es fundamental para estos sistemas, sigue siendo poco explorado. En este artículo, presentamos LiteVAE, una familia de autoencoders para LDMs que aprovechan la transformada discreta de wavelets en 2D para mejorar la escalabilidad y la eficiencia computacional en comparación con los autoencoders variacionales (VAEs) estándar, sin sacrificar la calidad de la salida. También investigamos las metodologías de entrenamiento y la arquitectura del decodificador de LiteVAE, y proponemos varias mejoras que optimizan la dinámica del entrenamiento y la calidad de la reconstrucción. Nuestro modelo base de LiteVAE iguala la calidad de los VAEs establecidos en los LDMs actuales con una reducción de seis veces en los parámetros del codificador, lo que resulta en un entrenamiento más rápido y menores requisitos de memoria en la GPU, mientras que nuestro modelo más grande supera a los VAEs de complejidad comparable en todas las métricas evaluadas (rFID, LPIPS, PSNR y 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).