Source-linked AI summary
Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
Junyu Chen, Han Cai, Junsong Chen, Enze Xie, Shang Yang, Haotian Tang, Muyang Li, Yao Lu, Song Han
TL;DR
High spatial-compression autoencoders offer a route to reduce token costs in high-resolution diffusion models but suffer reconstruction and resolution-generalization problems. DC-AE addresses these issues with Residual Autoencoding and Decoupled High-Resolution Adaptation, reaching compression ratios up to 128 and substantially improving diffusion-model throughput while improving FID. The paper concludes that autoencoder design can accelerate both diffusion-model training and inference.
Problem
High spatial-compression autoencoders are difficult to optimize and suffer reconstruction-quality drops, limiting efficient high-resolution diffusion models.
Method
DC-AE combines Residual Autoencoding, which learns residuals after space-to-channel transformation, with Decoupled High-Resolution Adaptation for resolution generalization.
Results
17.9× higher H100 training throughput and 19.1× higher H100 inference throughput improve ImageNet 512 × 512 FID from 3.55 to 3.01 on UViT-H versus SD-VAE-f8.
Takeaways & Limitations
DC-AE supports spatial compression ratios up to 128 while maintaining good reconstruction accuracy and improving latent-diffusion efficiency.
Takeaways & Limitations
The reported application setting includes 1024 × 1024 input resolution and batch size 12.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
DC-AE targets the reconstruction and efficiency limits of high spatial-compression autoencoders for high-resolution diffusion models. Residual Autoencoding and Decoupled High-Resolution Adaptation enable compression ratios up to 128 while improving diffusion-model efficiency and image-generation quality.
- Motivation: High-resolution diffusion transformers benefit from larger autoencoder compression ratios because their computation grows quadratically with token count.The common f8 autoencoder is adequate for lower-resolution synthesis, but higher compression is critical at high resolutions.
- Motivation: SD-VAE reconstruction rFID worsens from 0.90 at f8 to 28.3 at f64 on ImageNet 256 × 256.This reconstruction-accuracy drop is the main bottleneck to using high spatial-compression autoencoders.
- Method: Residual Autoencoding adds non-parametric shortcuts so network modules learn residuals from space-to-channel-transformed features, alleviating optimization difficulty.The paper identifies high spatial-compression autoencoders as harder to optimize than low-compression models.
- Method: Decoupled High-Resolution Adaptation uses high-resolution latent adaptation and low-resolution local refinement to mitigate resolution-generalization penalties at low training cost.The strategy separates adaptation and refinement phases rather than relying on a single training phase.
- Results: 17.9× higher H100 training throughput and 19.1× higher H100 inference throughput improve ImageNet 512 × 512 FID from 3.55 to 3.01 on UViT-H versus SD-VAE-f8.DC-AE-f64 takes over the token-compression task, allowing the diffusion model to focus on denoising.
- Results: DC-AE improves reconstruction accuracy at high compression ratios and delivers training and inference speedups for diffusion models compared with prior autoencoders.The reported family supports spatial compression ratios of 32, 64, and 128 while maintaining good reconstruction accuracy.
2 RELATED WORK
Prior latent-diffusion work largely retained f8 autoencoders while improving their reconstruction quality or accelerating diffusion-model sampling. DC-AE instead studies increasing autoencoder spatial compression as a distinct route to accelerate both training and inference.
- Autoencoder for Diffusion Models: Latent diffusion models reduce computational cost by operating on compressed representations produced by pretrained autoencoders.The widely adopted baseline uses an 8× spatial compression ratio and four latent channels.
- Autoencoder for Diffusion Models: Prior autoencoder work mainly improves f8 reconstruction through more latent channels, heavier decoders, or task-specific priors.These approaches enhance reconstruction quality without increasing the spatial compression ratio.
- Position of DC-AE: DC-AE focuses on the underexplored direction of increasing autoencoder spatial compression, such as f64, rather than modifying only the diffusion model.The paper describes this as distinct from prior autoencoder and diffusion-acceleration work.
- Diffusion Model Acceleration: Diffusion-model acceleration methods commonly reduce inference sampling steps while leaving the autoencoder unchanged.DC-AE addresses the autoencoder side and is positioned as a complementary acceleration direction.
3 METHOD
The method diagnoses accuracy gaps in high spatial-compression autoencoders and addresses them with residual shortcuts and decoupled high-resolution adaptation. DC-AE then shifts token compression toward the autoencoder for more efficient latent diffusion models.
- Residual Autoencoding: Residual Autoencoding adds non-parametric shortcuts so neural modules learn residuals relative to space-to-channel mappings rather than identity mappings.The shortcuts are added around downsample and upsample blocks, with space-to-channel or channel-to-space operations and channel matching.
- Residual Autoencoding: Residual Autoencoding improves reconstruction accuracy for high spatial-compression autoencoders on ImageNet 256 × 256.The design explicitly targets the optimization difficulty identified in the ablation analysis.
- Motivation: High spatial-compression autoencoders degrade despite greater capacity because their optimization is harder and their accuracy generalizes poorly across resolutions.The f64 autoencoder’s rFID rises from 0.50 to 7.40 from 256 × 256 to 1024 × 1024, while f8 improves from 0.51 to 0.19.
- Decoupled High-Resolution Adaptation: Decoupled High-Resolution Adaptation separates high-resolution latent adaptation from GAN-based local refinement to mitigate generalization penalties at lower training cost.Middle layers are tuned during high-resolution adaptation, while decoder head layers alone are tuned for local refinement; memory cost decreases from 153.98 GB to 67.81 GB under the stated setting.
- Application to Latent Diffusion Models: Directly reaching the target spatial compression ratio in the autoencoder performs best, and shifting compression from the diffusion model to the autoencoder consistently improves FID.Applying DC-AE to diffusion transformers requires changing only the patch size; larger patch size corresponds to additional space-to-channel token compression.
4 EXPERIMENTS
Experiments show that DC-AE improves reconstruction quality at high compression ratios and accelerates diffusion-transformer training and inference while maintaining or improving generation quality. The gains are demonstrated across image reconstruction, class-conditional generation, text-to-image generation, and model-scale comparisons.
- Image Compression and Reconstruction: DC-AE improves reconstruction rFID over SD-VAE across evaluated compression settings.On ImageNet 512 × 512, rFID improves from 16.84 to 0.22 for f64c128 and from 100.74 to 0.23 for f128c512.
- Image Compression and Reconstruction: DC-AE reconstructions retain better visual quality than SD-VAE, including small text and human faces at f64 and f128.
- Class-Conditional Image Generation: 4.5× higher H100 training throughput and 4.8× higher H100 inference throughput are achieved by DC-AE-f32p1 versus SD-VAE-f8p2 for DiT-XL.DC-AE-f32p1 also has 4× fewer tokens and better FID than SD-VAE-f8p2 across the evaluated diffusion transformers.
- USiT Image Generation: 1.72 FID is achieved by DC-AE-f32+USiT-2B on ImageNet 512×512, outperforming EDM2-XXL, MAGVIT-v2, and MAR-L.
- Text-to-Image Generation: DC-AE-f32p1 provides better FID and CLIP Score than SD-VAE-f8p2 in text-to-image generation.The models are trained for 100K iterations from scratch, and samples demonstrate high-quality synthesis with improved efficiency.
5 CONCLUSION
The paper concludes that DC-AE enables efficient high-resolution diffusion by combining deep token compression with reconstruction-preserving autoencoder designs. The resulting family reaches spatial compression ratios up to 128 and improves latent-diffusion efficiency.
- DC-AE reaches spatial compression ratios up to 128 while maintaining satisfactory reconstruction accuracy.
- Applying DC-AE to latent diffusion models produces significant training and inference efficiency improvements.
A DC-AE ARCHITECTURE AND TRAINING DETAILS
This appendix describes the architectures and optimization settings used for SD-VAE and DC-AE training. It specifies AdamW-based three-phase training with phase-specific learning rates, losses, and optimizer parameters.
- Figure 10 details the architectures of SD-VAE, DC-AE, and the DC-AE encoder and decoder stages.
- Optimization: AdamW is used as the optimizer for all training phases.
- Phase 1: Low-Resolution Full Training: Phase 1 uses low-resolution full training with a 6.4e-5 learning rate, weight decay 0.1, and L1 plus LPIPS losses.
- Phase 2: High-Resolution Latent Adaptation: Phase 2 uses high-resolution latent adaptation with a 1.6e-5 learning rate, weight decay 0.001, and the phase-1 loss combination.
- Phase 3: Low-Resolution Local Refinement: Phase 3 uses low-resolution local refinement with a 5.4e-5 learning rate and L1, LPIPS, and PatchGAN losses.
B ABLATION STUDY ON TRAINING DIFFERENT NUMBERS OF LAYERS
The ablation study examines how many layers are trained during the high-resolution latent-adaptation and low-resolution local-refinement phases.
- Figure 11 ablates different numbers of trained layers in phases 2 and 3.
C ADDITIONAL IMAGE RECONSTRUCTION RESULTS
Under low spatial compression, DC-AE slightly outperforms SD-VAE in reconstruction results; the section also documents related ablation and implementation settings.
- DC-AE delivers slightly better reconstruction results than SD-VAE under the low spatial-compression ratio setting.
- Figure 11 studies training different numbers of layers in Phase 2 and Phase 3.
- Latent features are normalized using the inverse root mean square as a scaling factor, without a shifting factor.
- The UViT models use DPMSolver with 30 sampling steps and guidance scale 1.5, while Figure 12 reports robustness across sampling hyperparameters.
G HIGH-RESOLUTION IMAGE GENERATION RESULTS
At higher generation resolutions, DC-AE-f32p1 achieves better FID than SD-VAE-f8p2 across all reported cases.
- DC-AE-f32p1 achieves better FID than SD-VAE-f8p2 for all tested 1024×1024 and 2048×2048 image-generation cases.
H IMAGE GENERATION RESULTS WITH OTHER EVALUATION METRICS
On ImageNet 512×512, DC-AE improves efficiency while maintaining generation performance across multiple evaluation metrics.
- DC-AE delivers significant efficiency improvements while maintaining generation performance across FID, IS, precision, recall, and CMMD.
I ADDITIONAL SAMPLES
Additional reconstruction and generation samples provide qualitative comparisons for DC-AE and SD-VAE across high-compression settings and 512×512 datasets.
- DC-AE reconstructions demonstrate better visual quality than SD-VAE reconstructions, especially for f64 and f128 autoencoders.
- Additional reconstruction samples include crops highlighting details such as human faces and small texts.
- Randomly generated samples are shown for ImageNet 512×512 and MJHQ-30K 512×512 using diffusion models with DC-AE.
- Table 7 reports class-conditional ImageNet 512×512 results with additional metrics and annotations for training settings and NFE.
- Figures 13 and 14 contain additional autoencoder image reconstruction samples.
- Figures 15 and 16 contain random 512×512 text-to-image and ImageNet-generated samples.