Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
| Autor: | Chen, Junyu, Cai, Han, Chen, Junsong, Xie, Enze, Yang, Shang, Tang, Haotian, Li, Muyang, Lu, Yao, Han, Song |
|---|---|
| Rok vydání: | 2024 |
| Předmět: | |
| Druh dokumentu: | Working Paper |
| Popis: | We present Deep Compression Autoencoder (DC-AE), a new family of autoencoder models for accelerating high-resolution diffusion models. Existing autoencoder models have demonstrated impressive results at a moderate spatial compression ratio (e.g., 8x), but fail to maintain satisfactory reconstruction accuracy for high spatial compression ratios (e.g., 64x). We address this challenge by introducing two key techniques: (1) Residual Autoencoding, where we design our models to learn residuals based on the space-to-channel transformed features to alleviate the optimization difficulty of high spatial-compression autoencoders; (2) Decoupled High-Resolution Adaptation, an efficient decoupled three-phases training strategy for mitigating the generalization penalty of high spatial-compression autoencoders. With these designs, we improve the autoencoder's spatial compression ratio up to 128 while maintaining the reconstruction quality. Applying our DC-AE to latent diffusion models, we achieve significant speedup without accuracy drop. For example, on ImageNet 512x512, our DC-AE provides 19.1x inference speedup and 17.9x training speedup on H100 GPU for UViT-H while achieving a better FID, compared with the widely used SD-VAE-f8 autoencoder. Our code is available at https://github.com/mit-han-lab/efficientvit. Comment: Preprint. First two authors contributed equally to this work. Update: add diffusion model scaling results |
| Databáze: | arXiv |
| Externí odkaz: |
|