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）的進步已經徹底改變了高解析度圖像生成，但是這些系統核心的自編碼器的設計空間仍然未被充分探索。在本文中，我們介紹了LiteVAE，這是一個針對LDMs的自編碼器家族，利用2D離散小波變換來提高可擴展性和計算效率，而不會犧牲輸出質量。我們還研究了LiteVAE的訓練方法和解碼器架構，並提出了幾個增強措施，以改善訓練動態和重建質量。我們的基本LiteVAE模型在當前LDMs中與已建立的VAEs相匹配，同時減少了六倍的編碼器參數，從而實現更快的訓練速度和更低的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).