针对扩散模型剪枝难题,提出新型剪枝方法提升稀疏度下的图像质量。
Post-Training Pruning for Diffusion Transformers

- 基于能量视角设计新重要性度量,平衡权重与激活贡献。
- 发现权重空间聚类规律,采用聚类感知剪枝粒度,提升稀疏分配效率。
- 在50%稀疏度下,CLIP分数仅降0.001,显著优于现有方法,适合高稀疏部署。
扩散Transformer(DiTs)在图像生成中表现优异,但计算开销大、资源消耗高。训练后剪枝是潜在解决方案,但因其独特的架构和参数分布,传统剪枝方法不适用,易导致性能大幅下降。以往针对大语言模型的方法通过近似推导重要性度量,放大了权重在度量中的相对贡献;而DiTs的权重幅度远高于大语言模型。此外,现有剪枝粒度未考虑模型结构差异。本文提出DiT-Pruning,通过定制化重要性度量和剪枝粒度改进剪枝效果。设计一种基于能量的新型度量,平衡权重与激活的贡献,更准确识别关键组件。同时观察到二维权重空间中的明显聚类模式,采用聚类感知剪枝粒度,实现高效稀疏分配。在多个DiTs上的大量实验表明,该方法在高稀疏度下仍能保持图像质量。例如,在MJHQ数据集上,FLUX.1-dev于512x512分辨率下达到50%稀疏度时,CLIP分数仅损失0.001,显著优于近期剪枝方法。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) have demonstrated impressive performance in image generation but suffer from substantial computational overhead and resource consumption. Post-training pruning offers a promising solution; however, due to DiTs' unique architectural design and parameter distribution, traditional pruning methods are inapplicable, leading to significant performance degradation. Specifically, prior methods developed for LLMs, which derive metrics through a series of approximations, amplify the relative contribution of weights in the saliency metric. In addition, weights in DiTs exhibit significantly larger magnitudes than those in LLMs. Moreover, existing pruning granularity overlooks variations in model structures. In this paper, we propose DiT-Pruning, which improves pruning performance by introducing customized saliency criteria and pruning granularity. We design a novel metric that balances the contributions of weights and activations from an energy-based perspective, enabling more effective identification of important elements. Furthermore, we observe distinct clustering patterns in the two-dimensional weight space. Accordingly, we adopt a clustering-aware pruning granularity, enabling effective sparse allocation. Extensive evaluations on various DiTs show that our method consistently preserves image quality, especially under high sparsity. For FLUX.1-dev at 512x512 resolution on MJHQ, DiT-Pruning achieves only a 0.001 loss in CLIP score at 50% sparsity, dramatically outperforming recent pruning methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。