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

초록

잠재 확산 모델(LDMs)의 발전은 고해상도 이미지 생성에 혁명을 일으켰지만, 이러한 시스템의 핵심인 오토인코더의 설계 공간은 아직 충분히 탐구되지 않았습니다. 본 논문에서는 2D 이산 웨이블릿 변환을 활용하여 표준 변분 오토인코더(VAEs) 대비 확장성과 계산 효율성을 향상시키면서도 출력 품질을 희생하지 않는 LiteVAE라는 오토인코더 패밀리를 소개합니다. 또한 LiteVAE의 훈련 방법론과 디코더 아키텍처를 조사하고, 훈련 역학과 재구성 품질을 개선하는 여러 가지 개선 사항을 제안합니다. 우리의 기본 LiteVAE 모델은 현재 LDMs에서 사용되는 기존 VAEs와 동등한 품질을 유지하면서 인코더 매개변수를 6분의 1로 줄여 더 빠른 훈련과 더 낮은 GPU 메모리 요구 사항을 달성하며, 더 큰 모델은 모든 평가 지표(rFID, LPIPS, PSNR, SSIM)에서 비슷한 복잡도의 VAEs를 능가합니다.

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