Adnan KhanAdnanProduct & AI
San Francisco Bay Area
Get in touch

Source note2 min read

Post-Training Pruning for Diffusion Transformers

Source

arXiv

arxiv.org/abs/2607.00927v1

Read the source

Source

Plain-Language Brief

The paper proposes DiT-Pruning, a method to improve post-training pruning performance for Diffusion Transformers (DiTs), which are used in image generation but have high computational overhead. The new method uses customized saliency criteria and pruning granularity, including a novel energy-based metric and clustering-aware pruning, to more effectively identify important elements and preserve image quality, especially under high sparsity.

What the Document Says

  • The paper discusses Diffusion Transformers (DiTs) and their high computational overhead.
  • Post-training pruning is proposed as a solution for DiTs, but traditional methods have significant performance degradation due to DiTs' unique architectural design and parameter distribution.
  • The paper highlights issues with prior methods developed for LLMs and their amplification of weight contributions in the saliency metric.
  • Weights in DiTs exhibit larger magnitudes than those in LLMs.
  • Existing pruning granularity overlooks variations in model structures.
  • The authors propose DiT-Pruning, which introduces customized saliency criteria and pruning granularity.
  • A novel energy-based metric is designed to balance the contributions of weights and activations.
  • Distinct clustering patterns are observed in the two-dimensional weight space.
  • Clustering-aware pruning granularity is adopted for effective sparse allocation.
  • Extensive evaluations show that DiT-Pruning preserves image quality, especially under high sparsity.
  • DiT-Pruning significantly outperforms recent pruning methods on FLUX.1-dev at 512x512 resolution on MJHQ.

Key Concepts Mentioned

  • Diffusion Transformers (DiTs)
  • post-training pruning
  • computational overhead
  • customized saliency criteria
  • pruning granularity
  • energy-based metric
  • clustering-aware pruning
  • image quality
  • high sparsity

Claims Made by the Document

  • Traditional pruning methods have significant performance degradation for DiTs.
  • DiT-Pruning improves pruning performance for DiTs.
  • DiT-Pruning preserves image quality, especially under high sparsity.

Evidence Used

  • Performance degradation of traditional pruning methods for DiTs is inferred from the text.
  • Improved pruning performance and preserved image quality of DiT-Pruning are stated based on evaluation results.

Methods / Approach

UNKNOWN

Limitations Stated by the Authors

UNKNOWN

Limitations Noted by the Agent

Agent interpretation — not from the source.

UNKNOWN

Open Questions

UNKNOWN

Notes for Future Q&A

May include interpretation. Labeled where it does.

UNKNOWN

Raw Abstract / Excerpt

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.