提出DC-AE自动编码器,实现128倍压缩下高保真图像重建。
Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
- 通过残差编码与解耦训练策略,解决高压缩比下的优化难题。
- 在ImageNet 512x512上实现19.1倍推理加速和17.9倍训练加速。
- 适合追求高效高分辨率生成的开发者和研究者使用。
我们提出Deep Compression Autoencoder(DC-AE),一种用于加速高分辨率扩散模型的新一代自动编码器。现有自动编码器在中等空间压缩比(如8x)下表现优异,但在高压缩比(如64x)时重建质量显著下降。为此,我们引入两项关键技术:(1) 残差自动编码,使模型基于空间到通道变换特征学习残差,缓解高压缩比自动编码器的优化困难;(2) 解耦高分辨率适配,采用高效的三阶段解耦训练策略,减轻高压缩比自动编码器的泛化损失。通过这些设计,我们将自动编码器的空间压缩比提升至128倍,同时保持高质量重建。将我们的DC-AE应用于潜在扩散模型,在不牺牲精度的前提下实现显著加速。例如,在ImageNet 512x512数据集上,相对于广泛使用的SD-VAE-f8自动编码器,我们的方法在H100 GPU上为UViT-H模型带来19.1倍推理加速和17.9倍训练加速,并获得更优的FID得分。代码已公开于https://github.com/mit-han-lab/efficientvit。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。