Post-Training Pruning for Diffusion Transformers

2026-07-01

AI Summary

This work addresses the high computational cost of Diffusion Transformers (DiTs) in image generation, where existing post-training pruning methods suffer significant performance degradation due to neglecting DiT-specific structural characteristics. To overcome this limitation, the authors propose DiT-Pruning, an efficient, training-free compression approach that jointly evaluates the importance of weights and activations from an energy perspective and incorporates two-dimensional weight-space clustering to determine a structure-adaptive pruning granularity. Evaluated on FLUX.1-dev, DiT-Pruning achieves only a 0.001 drop in CLIP score at 50% sparsity, substantially outperforming current pruning techniques and maintaining high-fidelity image generation even under aggressive sparsity levels.

Abstract

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.

Problem

• Diffusion Transformers
• Post-Training Pruning
• Computational Overhead
• Saliency Metric
• Pruning Granularity

Innovation

• Diffusion Transformers
• post-training pruning
• saliency metric
• clustering-aware granularity
• energy-based pruning