Light-DiT: An Importance-Aware Dynamic Compression Framework for Diffusion Transformers
摘要
Diffusion Transformers (DiTs) demonstrate remarkable generative abilities in AI. However, the iterative denoising process inherent in diffusion incurs substantial computational cost and memory overhead, impeding fast and energy-efficient edge inference. To mitigate the overhead of denoising, we propose a post-training framework that jointly utilizes pruning and quantization for hardware-efficient DiT inference, which is based on the observation that not all denoising blocks within a model are equally important during image generation. We introduce metrics to assess the importance of DiTs’ blocks and layers. To achieve importance-aware dynamic compression, we unify mixed-sparsity pruning and mixed-precision quantization based on importance metrics. Experiments show that our approach achieves a 1.41 \(\times \) inference speedup through pruning with a mixed precision of W3.2A4.9 while incurring minimal accuracy loss. Furthermore, the evaluation of bit-flexible DNN accelerators demonstrates up to 2.78 \(\times \) performance improvement, and 1.99 \(\times \) better energy efficiency can be achieved compared to W8A8 quantization without pruning.