后训练 扩散修剪 Transformer

Diffusion Transformer (DiTs) 在图像生成方面表现出了令人印象深刻的性能，但受到大量计算开销和资源消耗的影响。 后训练 修剪提供了一个有前途的解决方案；然而，由于DiTs独特的架构设计和参数分布，传统的剪枝方法不适用，导致性能显着下降。具体来说，为 LLM 开发的先前方法通过一系列近似推导度量，放大了权重在显着性度量中的相对贡献。此外，DiT 中的权重比 LLM 中的权重显着更大。此外，现有的修剪粒度忽略了模型结构的变化。在本文中，我们提出了 DiT-Pruning，它通过引入定制的显着性标准和剪枝粒度来提高剪枝性能。我们设计了一种新颖的指标，从基于能量的角度平衡权重和激活的贡献，从而能够更有效地识别重要元素。此外，我们在二维权重空间中观察到不同的聚类模式。因此，我们采用集群感知的修剪粒度，实现有效的稀疏分配。对各种 DiT 的广泛评估表明，我们的方法始终保持图像质量，特别是在高稀疏度下。对于 MJHQ 上分辨率为 512x512 的 FLUX.1-dev，DiT-Pruning 在稀疏度为 50% 时仅损失了 0.001 的 CLIP 分数，显着优于最近的剪枝方法。