LiteVAE: Легковесные и эффективные вариационные автоэнкодеры для моделей латентной диффузии

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

May 23, 2024

Авторы: Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, Romann M. Weber

cs.AI

Аннотация

Прогресс в моделях латентной диффузии (LDM) революционизировал генерацию изображений высокого разрешения, однако пространство проектирования автоэнкодера, центрального элемента этих систем, остается недостаточно исследованным. В данной статье мы представляем LiteVAE, семейство автоэнкодеров для LDM, которые используют 2D дискретное вейвлет-преобразование для улучшения масштабируемости и вычислительной эффективности по сравнению со стандартными вариационными автоэнкодерами (VAE) без ущерба качеству вывода. Мы также исследуем методики обучения и архитектуру декодера LiteVAE, предлагая несколько улучшений, которые повышают динамику обучения и качество восстановления. Наша базовая модель LiteVAE соответствует качеству установленных VAE в текущих LDM с шестикратным сокращением параметров энкодера, что приводит к более быстрому обучению и меньшим требованиям памяти GPU, в то время как наша более крупная модель превосходит VAE сопоставимой сложности по всем оцениваемым метрикам (rFID, LPIPS, PSNR и 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).