Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
打开互动全文版(逐段中英对照 + 图/公式 + 论文问答)→我们提出了深度压缩自编码器(DC-AE),这是一类用于加速高分辨率扩散模型的新型自编码器模型。现有的自编码器模型在中等空间压缩比(例如 8 倍)下表现出色,但在高空间压缩比(例如 64 倍)下无法保持令人满意的重建精度。我们通过引入两种关键技术来解决这一挑战:(1)残差自编码,其中我们设计模型基于空间到通道变换的特征学习残差,以缓解高空间压缩自编码器的优化困难;(2)解耦的高分辨率适应,一种高效的三阶段解耦训练策略,用于减轻高空间压缩自编码器的泛化惩罚。通过这些设计,我们将自编码器的空间压缩比提高到 128 倍,同时保持重建质量。将我们的 DC-AE 应用于潜在扩散模型,我们实现了显著的加速而精度不下降。例如,在 ImageNet 512x512 上,与广泛使用的 SD-VAE-f8 自编码器相比,我们的 DC-AE 在 UViT-H 上实现了 19.1 倍的推理加速和 17.9 倍的训练加速,同时获得了更好的 FID。我们的代码可在 https://github.com/mit-han-lab/efficientvit 获取。
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.