1 Introduction↩︎

Diffusion Transformers [1] (DiTs) exhibit remarkable performance in image generation, fully leveraging their advantages in scalability and expressive modeling capabilities. Recent models, such as PixArt  [2] and FLUX [3], have demonstrated the ability to handle complex image distributions and generate high-quality samples. However, their numerous number of parameters and prolonged inference process result in substantial computational overhead and high resource consumption, posing significant challenges for deployment in resource-constrained environments.

Neural network pruning [4]–[7] reduces computational and storage overhead by eliminating redundant weights. Post-training pruning avoids additional training and typically adopts unstructured approaches treating individual weights as sparse units. For DiTs, the massive parameter scale and complex architecture make training-based methods [8], [9] prohibitively costly, whereas post-training pruning is training-free and offers strong scalability and promising potential. Its effectiveness depends on two aspects: the saliency criteria used to measure parameter importance, and the pruning granularity that dictates sparsity allocation across different structural components like channels or layers. The proper design of both factors is intrinsically associated with the architectural design and the parameter distribution.

Early CNN pruning [10] employs Taylor-based saliency metrics. Similar criteria [11] have been adopted to prune attention heads in ViTs. Recently, several methods have been proposed for LLM pruning. For instance, SparseGPT [12] based on Optimal Brain Surgeon(OBS) theory [5] prunes weights that minimally impact output perturbations, while Wanda [13] introduces second-order approximations for computational efficiency. However, applicable pruning methods for Diffusion models remain largely unexplored. Some efforts have targeted U-Net-based diffusion models, for example, Diff-Pruning [14], which proposes dynamic pruning strategies incorporating timestep information and relies on retraining to restore performance. Nevertheless, due to the unique inference paradigm, architectural design, and parameter distribution of DiTs, existing pruning methods are inapplicable, leading to significant performance degradation, especially at high sparsity rates.

To this end, we thoroughly analyze the characteristics of DiTs from the perspectives of saliency criteria and pruning granularity. Regarding saliency criteria, many pruning methods developed for LLMs, represented by Wanda [13], construct efficient sensitivity estimations based on second-order approximations, achieving remarkable performance. However, these methods have the following limitations: Theoretically, the approximations adopted amplify weight contributions, disrupting the balance with activations in sensitivity estimation; empirically, we observe that DiT and LLM exhibit distinct statistical properties in the relative magnitude of weights and activations. As shown in Figure 2, weights in DiTs reach significantly larger magnitudes, leading to a mismatch between weight and activation scales. This discrepancy further exacerbates errors introduced by the approximations. Furthermore, as shown in Figure 3, we observe that the metric demonstrates distinct clustering patterns in the two-dimensional weight space, indicating that DiTs are incompatible with uniform pruning granularity.

In this work, we propose DiT-Pruning, an efficient post-training pruning method for DiTs, which improves pruning performance by introducing customized saliency criteria and pruning granularity. Specifically, we apply a squared transformation on weights to mitigate the amplification effect dominated by weight magnitudes in sensitivity estimation, which alleviates the scale disparity, facilitating accurate identification of the important elements. To address the clustering patterns observed in saliency criteria, we exploit the distributional structure to dynamically adjust pruning granularity. Per-layer and per-channel strategies each demonstrate advantages under different clustering patterns, enabling more effective and robust sparsity allocation. Our main contributions are summarized as follows:

• We propose DiT-Pruning, a post-training pruning method for DiTs, which effectively improves pruning performance at high sparsity utilizing parameter saliency estimation and pruning granularity.

• We modulate the saliency metric from an energy-based modeling perspective, providing a stronger theoretical foundation and better alignment with DiTs’ properties. Moreover, we propose a clustering-aware pruning granularity to facilitate alignment with distribution patterns.

• We conduct extensive evaluations on DiTs, demonstrating that our method consistently preserves image quality under high sparsity. On FLUX.1-dev at 512\(\times\)​512 resolution on MJHQ, we achieve only a 0.001 loss in CLIP score at 50% sparsity, outperforming exsisting methods.

2 Related Works↩︎

Diffusion Transformers. Diffusion Transformers [1] (DiTs) have attracted increasing attention and gained great popularity by replacing convolutional U-NeTs [15] with transformer-based architectures [16] for image generation. DiTs demonstrate state-of-the-art performance and are gradually emerging as a mainstream backbone for generative applications. Models like Pixart [2], Flux [3] and SANA [17] further extend DiTs to text-to-image tasks, highlighting their strong capability and scalability to model high-resolution images under complex semantic conditions. However, these advances come at the cost of rapidly growing parameter sizes and computational complexity, leading to substantial resource consumption.

Post-Training Pruning. Post-training pruning is a widely adopted model compression technique [18], [19] that aims to directly operate on pre-trained models without introducing additional training or fine-tuning costs. By identifying and removing redundant or less important weight parameters, post-training pruning produces a compact model representation, thereby reducing model sizes and computational overhead in practical deployment, while maintaining competitive performance.

Magnitude pruning [6], [20] is a classical model compression method and a strong baseline [9] in network pruning. It prunes parameters with small absolute values by ranking weights based on magnitude. However, in complex network architectures, parameter importance is influenced by multiple factors, further undermining the reliability of magnitude-based criteria.

Prior works like CNN pruning [10] combine gradients when designing saliency metrics. Similar criteria [11] have been adopted to prune attention heads in ViTs. Recent advanced methods have been proposed for LLMs, such as SparseGPT [12] and Wanda [13] , which focus on efficient unstructured pruning. Furthermore, Llm-pruner [21], SlimGPT [22] and SoBP [23] explore structured pruning. In contrast, pruning methods for diffusion models remain underexplored. Diff-Pruning [14] target U-Net-based [15] diffusion models, which introduces a timestep-aware dynamic pruning but requires retraining. Due to DiTs’ unique architectural design and parameter distribution, these methods are inapplicable, and resulting in significant performance degradation. This highlights the need for efficient post-training pruning methods specifically designed for DiTs.

3 Method↩︎

In this section, we thoroughly analyze the characteristics of DiTs from the perspectives of saliency criteria and pruning granularity and propose a novel post-training pruning method which effectively improves pruning performance. The overview of our method is shown in Figure 4.

3.1 Analysis of the OBD-based Criteria↩︎

Classical post-training pruning methods primarily stems from Optimal Brain Damage (OBD) [4] and Optimal Brain Surgeon (OBS) [5], which employ second-order sensitivity analysis to estimate the loss variation caused by pruning individual weights, thereby achieving unstructured sparsity. \[I_i \approx \frac{1}{2} h_{ii} w_i ^2 \label{OBD}\tag{1}\]

OBD are shown in Eq. 1 , where \(I_i\) denotes the importance of the \(i_{th}\) weight parameter, defined as the approximate increase in the loss function when this weight is set to zero. \(w_i\) denotes the magnitude of pruned weights, \(h_{ii}\) denotes the \(i_{th}\) diagonal element of the Hessian matrix(\(\mathbf{H}\)) of the loss function. OBS are shown in Eq. 2 , which further extends OBD. The formulations explicitly show that pruning sensitivity is jointly determined by both the weight magnitude and local curvature information, providing a principled second-order criterion for network sparsification. \[I_i \approx \frac{1}{2} \frac{ w_i ^2}{(\mathbf{H}^{-1})_{ii}} \label{OBS}\tag{2}\]

Based on these principles, a variety of pruning methods have been proposed. However, most of them are tailored to LLM architectures and become ineffective when applied to DiTs. In the following, we provide a detailed analysis from both theoretical and empirical perspectives to illustrate.

Subsequent works, such as SparseGPT [12] and Wanda [13], inherit from OBD and OBS framework and extend to the post-training pruning for LLMs by approximating the Hessian with activation-based second-order statistics under strong independence assumptions. This formulation enables the pruning metric to explicitly account for the combined influence of weight terms and activation terms on pruning sensitivity estimation. To achieve computational efficiency, these methods further adopt approximations, which reduce the accuracy of sensitivity estimation and alter the relative contributions of weight and activation terms. \[I_{ij} \approx \left[ \frac{|\mathbf{W}|^2}{\mathrm{diag}\big((\mathbf{X}^\top \mathbf{X} + \lambda \mathbf{I})^{-1}\big)} \right]_{ij} \label{SparseGPT}\tag{3}\]

Specifically, methods like SparseGPT formulates the pruning sensitivity as shown in Eq. 3 by employing a diagonal Hessian approximation with reconstruction compensation mechanism. In this equation, \(\mathbf{W}\in\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}\) denotes the weight matrix, \(\mathbf{X}\in\mathbb{R}^{N\times d_{\text{in}}}\) denotes the input activation matrix, where each row corresponds to a token embedding, and \(\mathbf{X}^\top\mathbf{X}\) serves as an empirical Hessian approximation. The diagonal operator \(\mathrm{diag}(\cdot)\) retains only the diagonal elements, and the damping term \(\lambda\mathbf{I}\) is introduced for numerical stability and partially mitigates the bias caused by the diagonal approximation, but it remains imperfect. \[I_{ij}^{\lambda \to 0} \approx \left[ \frac{|\mathbf{W}|^2}{\mathrm{diag}((\mathbf{X}^\top \mathbf{X})^{-1})} \right]_{ij}^{\text{diag approx}} = \left( |\mathbf{W}_{ij}| \cdot \|\mathbf{X}_j\|_2 \right)^2 \label{wanda}\tag{4}\]

Whereas Wanda in Eq. 4 further simplifies the estimation by discarding the reconstruction-aware optimization and taking the limit \(\lambda \rightarrow 0\), leading to a diagonally approximated, first-order surrogate. Here, \(|\mathbf{W}_{ij}|\) denotes the weight connecting the \(j\)-th input channel to the \(i\)-th output channel, and \(||\mathbf{X}_j||_2\) represents the \(\ell_2\) norm of the \(j\)-th input activation across tokens. The diagonal approximation of \((\mathbf{X}^\top \mathbf{X})^{-1}\) collapses the curvature information into a channel-wise scaling factor, yielding a multiplicative form composed of weight magnitude and activation norm.

Consequently, these approximations operate primarily on the activation terms, reducing their contribution to a bounded scaling factor. Meanwhile, the weight term is no longer modulated by the inverse Hessian and instead directly dominates sensitivity estimation, thus amplifying its contribution.

3.2 Pruning Metric↩︎

We revisit the design of importance metrics from an energy-based perspective, as illustrated in Figure 5 (Right). We start from the observation that parameter pruning can be interpreted as introducing a perturbation to the model, where the resulting loss increase can be formalized via a Taylor expansion. In classical mechanics, elastic potential energy grows quadratically with displacement, where the stored energy is determined by second-order interactions of the system, which motivates modeling the loss increase as a quadratic energy form governed by the Hessian. In this view, weight importance is treated as an energy allocation problem, where each parameter contributes proportionally to its squared magnitude, weighted by local curvature. This reveals that existing importance metrics, which rely on linear estimations of weight perturbations, inherently ignore second-order effects.

Actually, as shown in Figure 5 (Left), such linearization fails to capture the intrinsic quadratic energy structure of the loss landscape, leading to systematic underestimation of relative large-magnitude weights. In particular, linear estimations remain locally accurate in the small-weight regime but increasingly deviate as weight magnitude grows. Moreover, this discrepancy becomes more pronounced as weight magnitude increases, reflecting the important role of second-order effects in large parameter regimes.

This issue is particularly severe in DiTs due to their distinct scale characteristics compared to LLMs. As shown in Figure 4 (Left), DiT weights exhibit a wider range and overall larger magnitudes, which leads to the breakdown of linear estimations. Besides, \(\mathbf{X}\in\mathbb{R}^{N\times d_{\text{in}}}\) denotes the input activation matrix, \(N\) denotes the number of tokens, since DiTs operate in the latent space with significantly fewer tokens (e.g., \(N=256\) or \(512\)), compared to LLMs with long token sequences (e.g., \(N=2048\), \(4096\) or \(8192\)), the resulting \(\|\mathbf{X}_j\|_2\) values are substantially smaller. Consequently, the scale imbalance further amplifies the relative dominance of weight terms, making the approximation imbalance more apparent and structurally biased in DiTs.

To retain the computational efficiency of existing approximations, we adopt \(|\mathbf{W}_{ij}|\) and \(|\mathbf{X}_j|_2\) as the fundamental weight and activation components. Following classical second-order pruning criteria, we preserve the separable multiplicative structure and introduce exponents \((m,n)\) to parameterize contributions, as shown in Eq. 5 . Since the scaling factors \(a\) and \(b\) do not affect within-layer ranking, they are omitted, resulting in the simplified form in Eq. 6 . \[I_{ij} = \left( a \cdot |\mathbf{W}_{ij}| \right)^{m} \cdot \left( b \cdot \|\mathbf{X}_j\|_2 \right)^{n} \label{eq:parametric95form}\tag{5}\] \[I_{ij} = \left( |\mathbf{W}_{ij}| \right)^{m} \cdot \left( \|\mathbf{X}_j\|_2 \right)^{n} \label{eq:parametric95form951}\tag{6}\] \[I_{ij} = \left( |\mathbf{W}_{ij}| \right)^{2} \cdot \left( \|\mathbf{X}_j\|_2 \right)^{1} \label{eq:parametric95form952}\tag{7}\]

Motivated by this observation, we instantiate our parametric formulation by applying a squared transformation to the weight term (STW), corresponding to \((m=2)\) and \((n=1)\) in Eq. 6 , as defined in Eq. 7 . This design restores the quadratic energy structure of perturbation-induced loss variations and rebalances the relative contributions of weights and activations, while preserving the original power relationship, yielding a more faithful and reliable sensitivity estimation. Meanwhile, STW compresses the dominant DiT weight scale from \(10^{-1}\) to \(10^{-2}\), bringing it closer to the regime where Wanda was originally developed for LLMs (Appendix A shows detailed statistics).

3.3 Pruning Granularity↩︎

In general, pruning removes parameters with low importance scores within predefined structural groups. As shown in Eq. 8 , \(\mathbf{W}\) and \(\widetilde{\mathbf{W}}\) denote the original and pruned weight matrices, respectively, while \(\mathbf{M}_{\mathcal{G}} \in {0,1}^{|\mathbf{W}|}\) is a group-wise binary mask determined by the predefined pruning groups \(\mathcal{G}\). These groups define the comparison scope for ranking importance scores, making the choice of pruning granularity critical to the final sparsification performance. \[\widetilde{\mathbf{W}} = \mathbf{W} \odot \mathbf{M}_{\mathcal{G}} \label{mask}\tag{8}\]

Existing post-training pruning methods typically adopt a uniform granularity, overlooking variations in weight distributions and model structures. In contrast, we observe that our metric exhibits distinct clustering patterns in the two-dimensional weight space of DiTs, as illustrated in Figure 4 (Right). Specifically, some weights form elongated, channel-aligned clusters, while others aggregate into compact, isotropic groups. These observations reveal substantial heterogeneity in importance distributions across layers and channels, suggesting that a globally uniform pruning granularity is inherently suboptimal. \[\bar{I}_i = \frac{1}{C_{\text{in}}} \sum_{j=1}^{C_{\text{in}}} I_{ij} \label{eq:per95output95mean}\tag{9}\] \[\mathcal{G}_{\text{out}} = \frac{ \mathrm{Var}\!\left( \left\{ \bar I_i \right\}_{i=1}^{C_{\text{out}}} \right) }{ \left( \mathrm{Mean}\!\left( \left\{ \bar I_i \right\}_{i=1}^{C_{\text{out}}} \right) \right)^2 } \label{eq:pnv}\tag{10}\]

Motivated by the insight, we propose a clustering-aware pruning granularity (CAG) that dynamically adjusts comparison groups according to the underlying importance distribution. Specifically, we use the averaged channel importance \(\bar{I}_i\) and the normalized heterogeneity metric \(\mathcal{G}_{\text{out}}\) in Eq. 9 and Eq. 10 to characterize layer-wise importance concentration. Layers with small \(\mathcal{G}_{\text{out}}\), such as MLP and self-attention layers, exhibit compact importance patterns and are pruned at the per-layer level. In contrast, layers with large \(\mathcal{G}_{\text{out}}\), notably adaLN layers, show strong channel-wise heterogeneity and are pruned at the per-channel level. This strategy enables more effective sparsity allocation and better preserves model performance under high sparsity.

4 Experiments↩︎

4.1 Experimental Settings↩︎

Models. We evaluate our method primarily focus on three types of diffusion models: the classical pre-trained DiT-XL/2 model, the state-of-the-art PixArt-\(\Sigma\) model [2], and the FLUX.1-dev model [3]. All evaluated diffusion models adopt transformer-based backbones [16].

Datasets. Following previous works [24], [25], we randomly sample prompts from COCO [26] for calibration. We evaluate DiT-XL/2 on ImageNet [27], generating 10,000 images at 256\(\times\)​256 resolution using the official DDPM sampler, consistent with prior studies [28], [29]. Robustness is further assessed with 250, 100, and 50 diffusion steps [30]. For PixArt-\(\Sigma\) and FLUX.1-dev, we evaluate on COCO and MJHQ, generating images at 256\(\times\)​256 and 512\(\times\)​512 resolutions with default samplers. We use the first 1,024 COCO annotation prompts and randomly sample 1,024 prompts from MJHQ-5K, a subset of the MJHQ-30K dataset [31].

\begin{table}[t]

4pt

Timesteps	Sparsity	Method	Params	FID \(\downarrow\)	IS \(\uparrow\)	sFID \(\downarrow\)	PR \(\uparrow\)	SSIM \(\uparrow\)
			(M)
250	0%	Dense	675.13	5.33	273.62	17.84	0.84	–
	20%	Magnitude	541.35	5.17	272.27	17.35	0.84	0.788
		Wanda	541.55	5.57	276.77	18.38	0.84	0.785
		DiT-Pruning	541.43	5.46	276.26	17.83	0.84	0.825
	40%	Magnitude	407.58	9.39	163.42	19.81	0.70	0.545
		Wanda	407.94	13.09	152.66	41.74	0.68	0.537
		DiT-Pruning	407.73	6.53	224.27	23.39	0.79	0.617
	50%	Magnitude	340.69	42.06	51.27	34.91	0.41	0.474
		Wanda	340.69	45.97	49.20	97.17	0.62	0.461
		DiT-Pruning	340.69	13.86	134.94	36.18	0.66	0.534
100	0%	Dense	675.13	5.53	261.87	18.93	0.81	–
	20%	Magnitude	541.35	5.41	260.92	18.15	0.82	0.803
		Wanda	541.55	5.72	263.46	19.06	0.82	0.796
		DiT-Pruning	541.43	5.56	263.78	18.61	0.82	0.835
	40%	Magnitude	407.58	12.16	150.41	22.02	0.66	0.559
		Wanda	407.94	15.34	142.27	42.19	0.65	0.541
		DiT-Pruning	407.73	7.19	213.22	23.58	0.77	0.627
	50%	Magnitude	340.69	50.25	46.79	40.82	0.36	0.485
		Wanda	340.69	50.96	43.54	97.97	0.37	0.463
		DiT-Pruning	340.69	16.28	124.69	36.86	0.64	0.541
50	0%	Dense	675.13	6.68	236.27	21.62	0.78	–
	20%	Magnitude	541.35	6.38	237.56	20.43	0.79	0.823
		Wanda	541.55	6.57	238.33	20.78	0.79	0.813
		DiT-Pruning	541.43	6.41	242.23	20.68	0.79	0.853
	40%	Magnitude	407.58	17.29	131.89	26.50	0.62	0.584
		Wanda	407.94	19.20	122.57	43.60	0.62	0.555
		DiT-Pruning	407.73	8.78	194.21	24.50	0.75	0.646
	50%	Magnitude	340.69	62.40	37.98	50.50	0.32	0.508
		Wanda	340.69	57.32	38.84	101.30	0.34	0.473
		DiT-Pruning	340.69	20.29	111.95	38.95	0.60	0.559

\vskip

-0.1in \end{table}

Baselines. We compare our method with two representative approaches. Magnitude pruning [6], [20] is a classical pruning strategy and serves as a strong baseline in our experiments. Wanda [13] is a representative second-order pruning method originally designed for LLMs. To construct the calibration data, we collect 128 input samples from the model at randomly sampled diffusion timesteps [32].

Metrics. To comprehensively assess generated image quality, we employ the following metrics: We choose Fréchet Inception Distance (FID) [33] for fidelity evaluation, spatial FID (sFID) [34], [35] for spatial consistency, Inception Score (IS) [36] for visual quality and diversity, Clipscore (CLIP) [37] for text-image alignment, ImageReward (IR) [38] for human preference, Precision (PR) [39] for sample realism, SSIM [40] for structural coherence. All experiments are conducted under identical experiment settings on NVIDIA RTX A6000 GPUs.

4.2 Pruning Performance↩︎

We present a comprehensive assessment of our method against prevalent baseline methods in various settings.

DiT-XL/2 on ImageNet. Table ¿tbl:tab:dit95imagenet? reports the performance of pruned DiT-XL/2 model across various pruning settings and timesteps. At 20% sparsity, a mild pruning, all methods yield similar outcomes. As sparsity increase to 40%, our method consistently outperforms the baselines and remains close to the dense model, suggesting the existence of highly effective sparse sub-networks within DiTs. The advantage becomes more pronounced at 50% sparsity, where baseline methods suffer substantial degradation while our method maintains strong generative performance. For example, under 250 sampling steps and 50% sparsity, magnitude pruning and Wanda increase FID by 36.73 and 40.64, respectively, whereas our method incurs only an 8.53 increase, demonstrating the effectiveness under aggressive sparse.

PixArt-\(\Sigma\) on COCO. Table ¿tbl:tab:pixart95coco? presents the results of pruning PixArt-\(\Sigma\) on COCO. Our method consistently achieves competitive or superior performance in both image quality and text-image alignment. At 40% sparsity, it attains the best FID and CLIP scores, indicating strong preservation of visual quality and semantic fidelity. The advantage becomes more pronounced at 50% sparsity, where baseline methods suffer substantial degradation, while our approach maintains stable performance with lower FID and higher CLIP scores.

FLUX.1-dev on COCO and MJHQ. Table ¿tbl:tab:flux95256? and ¿tbl:tab:flux95512? report the results on FLUX.1-dev across COCO and MJHQ. Similar to the observations on PixArt-\(\Sigma\), all methods perform comparably at 20% and 40% sparsity, with only minor differences in image quality and text-image alignment metrics. As sparsity increases to 50%, the performance gap becomes substantially larger, where our method consistently achieves superior results and demonstrates stronger robustness.

\begin{table}[t]

4pt

Resolution	Sparsity	Method	Params	FID \(\downarrow\)	IS \(\uparrow\)	CLIP \(\uparrow\)	IR \(\uparrow\)	sFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
			(M)
256\(\times\)256	0%	Dense	612.08	70.31	33.33	0.265	0.807	249.59	0.49	–
	20%	Magnitude	493.42	69.44	32.92	0.266	0.821	249.16	0.49	0.543
		Wanda	493.64	68.60	30.62	0.264	0.727	248.03	0.50	0.631
		DiT-Pruning	493.46	69.01	31.84	0.265	0.793	249.52	0.51	0.688
	40%	Magnitude	374.48	69.87	31.36	0.258	0.498	239.25	0.44	0.407
		Wanda	374.87	69.32	28.86	0.256	0.244	235.79	0.47	0.456
		DiT-Pruning	374.55	67.50	31.06	0.260	0.635	244.21	0.51	0.484
	50%	Magnitude	315.02	92.70	21.03	0.231	-0.728	240.16	0.23	0.312
		Wanda	315.10	89.66	20.47	0.229	-0.923	234.46	0.27	0.381
		DiT-Pruning	315.10	69.79	29.76	0.251	0.183	237.60	0.42	0.385
512\(\times\)512	0%	Dense	612.08	62.36	36.12	0.256	0.929	274.47	0.59	–
	20%	Magnitude	494.29	61.25	37.76	0.257	0.968	268.88	0.59	0.484
		Wanda	494.51	61.22	36.92	0.255	0.879	269.71	0.60	0.674
		DiT-Pruning	494.34	61.71	35.78	0.257	0.942	272.10	0.60	0.693
	40%	Magnitude	375.35	61.12	34.87	0.252	0.587	257.66	0.53	0.378
		Wanda	375.75	61.26	33.33	0.247	0.360	256.87	0.57	0.475
		DiT-Pruning	375.43	56.74	35.98	0.253	0.803	264.47	0.62	0.514
	50%	Magnitude	315.90	108.86	19.66	0.214	-1.142	272.80	0.18	0.261
		Wanda	315.97	93.97	19.81	0.219	-1.143	263.23	0.29	0.373
		DiT-Pruning	315.97	61.78	33.27	0.244	0.285	254.82	0.51	0.439

\vskip

-0.1in \end{table}

\begin{table}[t]

4pt

Dataset	Sparsity	Method	Params	FID \(\downarrow\)	IS \(\uparrow\)	CLIP \(\uparrow\)	IR \(\uparrow\)	sFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
			(B)
MJHQ	0%	Dense	11.9	54.94	26.11	0.260	0.686	259.18	0.75	–
	20%	Magnitude	9.53	57.05	24.32	0.257	0.607	262.01	0.73	0.692
		Wanda	9.54	54.48	25.61	0.259	0.681	259.97	0.74	0.786
		DiT-Pruning	9.53	54.54	26.37	0.259	0.689	260.67	0.74	0.836
	40%	Magnitude	7.16	64.56	20.00	0.246	0.251	267.40	0.63	0.499
		Wanda	7.17	57.28	23.83	0.257	0.582	256.94	0.69	0.584
		DiT-Pruning	7.17	54.83	23.64	0.259	0.647	255.83	0.71	0.663
	50%	Magnitude	5.97	74.38	16.30	0.235	-0.384	273.47	0.49	0.422
		Wanda	5.98	62.43	19.00	0.256	0.281	255.13	0.58	0.477
		DiT-Pruning	5.98	56.78	23.51	0.257	0.517	256.29	0.69	0.567
COCO	0%	Dense	11.9	66.01	38.78	0.154	-0.964	295.08	0.73	–
	20%	Magnitude	9.53	66.54	37.26	0.154	-1.037	296.29	0.68	0.675
		Wanda	9.54	65.62	37.57	0.154	-0.991	293.87	0.72	0.784
		DiT-Pruning	9.53	65.31	37.69	0.154	-0.965	294.94	0.72	0.828
	40%	Magnitude	7.16	71.11	30.49	0.155	-1.227	308.19	0.57	0.493
		Wanda	7.17	68.27	36.15	0.155	-1.049	294.71	0.61	0.561
		DiT-Pruning	7.17	64.96	37.71	0.156	-0.992	292.77	0.69	0.645
	50%	Magnitude	5.97	78.46	24.14	0.159	-1.476	317.77	0.42	0.434
		Wanda	5.98	74.85	27.87	0.155	-1.179	287.41	0.44	0.452
		DiT-Pruning	5.98	67.56	35.91	0.161	-1.089	295.79	0.60	0.537

\vskip

-0.1in \end{table}

\begin{table}[h]

4pt

Dataset	Sparsity	Method	Params	FID \(\downarrow\)	IS \(\uparrow\)	CLIP \(\uparrow\)	IR \(\uparrow\)	sFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
			(B)
MJHQ	0%	Dense	11.9	48.64	29.56	0.272	0.899	266.47	0.78	–
	20%	Magnitude	9.53	50.76	26.74	0.271	0.867	270.15	0.76	0.731
		Wanda	9.54	48.82	26.98	0.272	0.921	265.42	0.77	0.801
		DiT-Pruning	9.53	48.93	28.75	0.273	0.913	265.72	0.78	0.847
	40%	Magnitude	7.16	57.36	22.32	0.262	0.635	269.75	0.71	0.561
		Wanda	7.17	51.98	23.66	0.271	0.895	267.40	0.75	0.624
		DiT-Pruning	7.17	49.56	27.07	0.272	0.899	265.03	0.77	0.692
	50%	Magnitude	5.97	66.03	19.77	0.249	0.133	269.02	0.57	0.493
		Wanda	5.98	61.44	18.73	0.265	0.684	270.97	0.62	0.504
		DiT-Pruning	5.98	52.42	24.10	0.271	0.828	264.81	0.76	0.603
COCO	0%	Dense	11.9	63.41	39.13	0.154	-0.801	298.68	0.75	–
	20%	Magnitude	9.53	63.82	36.45	0.155	-0.811	300.94	0.74	0.717
		Wanda	9.54	63.58	37.37	0.155	-0.812	299.05	0.72	0.793
		DiT-Pruning	9.54	63.12	37.73	0.155	-0.813	298.59	0.73	0.839
	40%	Magnitude	7.16	66.58	29.94	0.157	-0.916	306.87	0.71	0.546
		Wanda	7.17	65.73	34.38	0.156	-0.836	305.19	0.66	0.616
		DiT-Pruning	7.17	64.14	36.53	0.158	-0.824	299.66	0.71	0.684
	50%	Magnitude	5.97	69.18	26.77	0.156	-1.139	308.31	0.58	0.495
		Wanda	5.98	74.19	26.63	0.158	-0.915	310.66	0.49	0.495
		DiT-Pruning	5.98	66.15	33.29	0.158	-0.863	304.46	0.64	0.588

\vskip

-0.1in \end{table}

4.3 Ablation Study↩︎

To validate the effectiveness of our method, we conduct ablation studies on DiT-XL/2 over ImageNet at 256\(\times\)​256 resolution under different sparsity and sampling steps. We further evaluate FLUX.1-dev under structured sparsity on MJHQ at different resolutions to verify the effectiveness of STW. Starting from Wanda as the baseline, we incrementally incorporate each component, STW and CAG, and evaluate their individual contributions. As shown in Tables ¿tbl:tab:ablation1? and ¿tbl:tab:ablation2?, each component improves performance and their combination achieves the best. Appendix E shows more results.

\begin{table}[h]

4pt

Timesteps	Sparsity	Ablation	FID \(\downarrow\)	IS \(\uparrow\)	sFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
100	0%	Dense	5.53	261.87	18.93	0.81	/
	40%	Wanda	15.34	142.27	42.19	0.65	0.541
		+ STW	11.79	161.01	42.53	0.74	0.561
		+ STW + CAG	7.19	213.22	23.58	0.77	0.627
		Wanda	50.96	43.54	97.97	0.37	0.463
		+ STW	43.47	54.34	89.22	0.41	0.477
		+ STW + CAG	16.28	124.69	36.86	0.64	0.541
50	0%	Dense	6.68	236.27	21.62	0.78	/
	40%	Wanda	19.20	122.57	43.60	0.62	0.555
		+ STW	12.86	143.31	43.61	0.73	0.576
		+ STW + CAG	8.78	194.21	24.50	0.75	0.646
		Wanda	57.32	38.84	101.30	0.34	0.473
		+ STW	47.11	47.01	92.57	0.39	0.488
		+ STW + CAG	20.29	111.95	38.95	0.60	0.559

\vskip

-0.05in \end{table}

\begin{table}[h]

4pt

Resolution	Sparsity	Method	FID \(\downarrow\)	IS \(\uparrow\)	CLIP \(\uparrow\)	IR \(\uparrow\)	SFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
256\(\times\)256	0%	Dense	54.94	26.11	0.260	0.686	259.18	0.75	/
	4:8	Wanda	71.56	16.38	0.243	-0.206	260.76	0.49	0.441
		+ STW	63.95	20.51	0.246	0.189	257.96	0.61	0.489
	2:4	Wanda	79.35	14.63	0.235	-0.613	263.35	0.42	0.412
		+ STW	69.21	19.98	0.245	-0.078	254.42	0.55	0.444
512\(\times\)512	0%	Dense	48.64	29.56	0.272	0.899	266.47	0.78	/
	4:8	Wanda	61.31	20.09	0.259	0.414	256.52	0.64	0.522
		+ STW	57.58	23.27	0.261	0.575	254.47	0.71	0.548
	2:4	Wanda	63.07	18.36	0.251	0.054	242.66	0.56	0.466
		+ STW	59.89	23.32	0.255	0.340	241.22	0.65	0.497

\vskip

-0.1in \end{table}

Parameter Importance Evaluated by STW. Under 50% sparsity with 50 denoising steps, STW substantially improves performance over the baseline, reducing FID and sFID by 10.21 and 8.73, respectively, under the same pruning granularity, demonstrating its effectiveness in identifying important parameters. Similar gains are observed under 2:4 and 4:8 structured sparsity on FLUX.1-dev. For instance, under 2:4 sparsity at 256\(\times\)​256 resolution, STW reduces FID from 79.35 to 69.21 while improving CLIP, IR, and SSIM. These results demonstrate that the proposed importance criterion generalizes to both unstructured and structured pruning.

Sparse Allocation by CAG. Furthermore, we compare fixed pruning granularity with CAG under the same importance metric. By dynamically adapting pruning granularity to the underlying importance distribution, CAG consistently improves performance across sparsity levels and denoising steps. For example, as shown in Table ¿tbl:tab:ablation1?, at 40% sparsity with 100 denoising steps, FID decreases from 11.79 to 7.19, while at 50% sparsity with 50 denoising steps, it further drops from 47.11 to 20.29. These results suggest that adaptive granularity allocation enables more effective sparse allocation.

5 Conclusion↩︎

In this paper, we propose DiT-Pruning, a post-training pruning method for DiTs. Existing pruning methods suffer substantial performance degradation on DiTs due to their unique architecture, parameter distribution, and approximations that amplify weight contributions. To address this issue, we introduce a customized saliency criterion and pruning granularity to improve pruning performance. Specifically, we balance weight and activation contributions by squaring the weights, enabling more accurate identification of important elements. Furthermore, we observe that the proposed metric exhibits distinct clustering patterns in the two-dimensional weight space. Accordingly, we adopt a clustering-aware pruning granularity for effective sparse allocation. Extensive experiments show that our method significantly outperforms existing approaches, especially under high sparsity.

6 Weights and Activations↩︎

Weights. Due to space limitations, we present the weight visualizations of only two representative layers in the main text. Similar distribution patterns are consistently observed across all layers of the DiTs, indicating that the observed phenomenon is a general characteristic of the model rather than an isolated case. For completeness, the weight distributions of all remaining model layers are presented below. Figure 6 shows detailed information, weights are predominantly distributed around the \(10^{-1}\) scale, which is consistent with the distribution described before.

\begin{table}[h]

4pt

Layer	Metric	block5	block10	block15	block20	block25
atten.qkv	\(|X|\) mean	1.1012	1.6378	1.7370	1.8147	1.9152
	\(|X|\) max	22.3234	10.8295	9.2846	9.0306	9.4552
	\(|W|/|X|_{\rm mean}\)	0.1052	0.0979	0.0972	0.0905	0.0832
	\(|W|^2/|X|_{\rm mean}\)	0.0122	0.0157	0.0164	0.0149	0.0133
atten.proj	\(|X|\) mean	4.9624	8.5670	11.7076	9.9487	13.3197
	\(|X|\) max	16.7251	16.3386	36.9063	27.2713	36.1129
	\(|W|/|X|_{\rm mean}\)	0.0283	0.0213	0.0159	0.0193	0.0149
	\(|W|^2/|X|_{\rm mean}\)	0.0039	0.0039	0.0037	0.0037	0.0029
mlp.fc1	\(|X|\) mean	1.5179	1.5725	1.1350	1.1803	1.4208
	\(|X|\) max	17.1799	16.0704	8.7672	22.1027	41.0000
	\(|W|/|X|_{\rm mean}\)	0.0936	0.1047	0.1453	0.1373	0.1098
	\(|W|^2/|X|_{\rm mean}\)	0.0133	0.0172	0.0239	0.0223	0.0171
mlp.fc2	\(|X|\) mean	2.7627	3.4125	2.9495	2.8912	3.0111
	\(|X|\) max	16.0696	25.4478	22.8027	27.8377	36.2336
	\(|W|/|X|_{\rm mean}\)	0.0589	0.0532	0.0647	0.0609	0.0646
	\(|W|^2/|X|_{\rm mean}\)	0.0095	0.0097	0.0127	0.0129	0.0126
adaLN	\(|X|\) mean	0.0084	0.0084	0.0084	0.0084	0.0084
	\(|X|\) max	0.3409	0.3409	0.3409	0.3409	0.3409
	\(|W|/|X|_{\rm mean}\)	11.45	16.42	16.00	13.76	13.07
	\(|W|^2/|X|_{\rm mean}\)	1.1017	2.2638	2.1504	1.5908	1.4274

\vskip

-0.1in \end{table}

Activations. Table ¿tbl:tab:activation95stats1? and Table ¿tbl:tab:activation95stats2? report the activation statistics and weight-to-activation ratios of representative DiT and LLM layers, respectively. Owing to their much longer token sequences, LLMs exhibit significantly larger activation outliers. Despite this difference, the relative scales of weights and activations remain substantially more balanced in LLMs than in DiTs. STW reduces the dominant DiT weight scale from \(10^{-1}\) to \(10^{-2}\), yielding a quantitative relationship between weights and activations that is closer to the regime where Wanda has been shown to be effective.

\begin{table}[h]

4pt

Layer	Metric	block5	block10	block15	block20	block25
atten.q	\(|X|\) mean	11.4072	14.1838	15.8599	17.7827	19.9309
	\(|X|\) max	214.7673	229.1075	210.9702	185.4212	189.7775
	\(|W|/|X|_{\rm mean}\)	0.0019	0.0015	0.0013	0.0011	0.0009
atten.k	\(|X|\) mean	11.4072	14.1838	15.8599	17.7827	19.9309
	\(|X|\) max	214.7673	229.1075	210.9702	185.4212	189.7775
	\(|W|/|X|_{\rm mean}\)	0.0019	0.0009	0.0013	0.0011	0.0009
atten.v	\(|X|\) mean	11.4072	14.1838	15.8599	17.7827	19.9309
	\(|X|\) max	214.7673	229.1075	210.9702	185.4212	189.7775
	\(|W|/|X|_{\rm mean}\)	0.0012	0.0009	0.0013	0.0009	0.0009
atten.o	\(|X|\) mean	1.7841	5.1910	6.5438	6.4441	6.3077
	\(|X|\) max	14.9358	37.6035	26.8657	43.6976	36.2537
	\(|W|/|X|_{\rm mean}\)	0.0073	0.0027	0.0023	0.0027	0.0029
mlp.up	\(|X|\) mean	6.9415	8.6943	10.9269	14.2518	16.3169
	\(|X|\) max	95.0688	83.8602	53.8245	43.4382	58.0632
	\(|W|/|X|_{\rm mean}\)	0.0024	0.0020	0.0016	0.0012	0.0011
mlp.gate	\(|X|\) mean	6.9415	8.6943	10.9269	14.2518	16.3169
	\(|X|\) max	95.0688	83.8602	53.8245	43.4382	58.0632
	\(|W|/|X|_{\rm mean}\)	0.0027	0.0021	0.0016	0.0013	0.0011
mlp.down	\(|X|\) mean	1.7290	2.9034	4.5845	7.4560	8.5851
	\(|X|\) max	16.3258	18.8423	31.4849	65.0971	55.4523
	\(|W|/|X|_{\rm mean}\)	0.0096	0.0058	0.0038	0.0024	0.0021

\vskip

-0.1in \end{table}

7 Compression and Speedup Effects↩︎

Compression and Acceleration. DiT-Pruning introduces no additional parameters and therefore preserves the inherent benefits of sparsification. As shown in Table 1, both parameter count and computational cost decrease consistently with increasing sparsity. TFLOPs is used to measure the acceleration effect, Params is used to measure the storage reduction. Increasing compression sparsity results in significant reduction in inference cost, demonstrating its effectiveness.

4pt

Acceleration Analysis. Existing sparse inference frameworks [41], [42] have demonstrated substantial acceleration for unstructured sparse computation, while prior work reports approximately (1.6\(\times\)) runtime speedup for Transformer layers under 2:4 structured sparsity [13]. Since DiT-Pruning follows the same unstructured and 2:4 sparsity patterns without introducing additional modules, the achieved reductions in model size and FLOPs can be directly translated into practical efficiency gains under these sparse execution frameworks.

8 Experiment Settings↩︎

Parameter Settings. For DiT-XL/2, we follow the standard ImageNet generation setup with a classifier-free guidance (CFG) scale of 1.5, a fixed random seed of 42, and the MSE VAE for latent coding. Evaluations are conducted under different timesteps. For PixArt-\(\Sigma\), we use DPM-Solver with 20 denoising steps and a CFG scale of 4.5. For FLUX.1-dev, we adopt 50 inference steps with a guidance scale of 3.5. For fair comparison across resolutions, all reference images are resized to the target resolution when necessary.

Evaluations. We evaluate model performance using generation quality, text-image alignment, and structural consistency metrics. Generation quality is assessed by FID, sFID, Precision, and Inception Score (IS), following the official evaluation pipeline of OpenAI’s guided-diffusion repository with the corresponding reference batches. Text-image alignment is measured using CLIP with the official pretrained weights. For FLUX.1-dev, we additionally report Image Reward (IR), following the official evaluation protocol. Structural consistency is evaluated using SSIM. All evaluation settings follow the corresponding official implementations.

9 Calibration Samples↩︎

Number of Calibration Samples. To evaluate the sensitivity to calibration set size, we progressively increase the number of calibration samples while keeping all other settings fixed. Experiments are conducted on DiT-XL/2 at 256\(\times\)​256 resolution with 50 denoising steps and 5,000 generated samples. As shown in Table ¿tbl:tab:calibration95size?, the performance remains stable across a wide range of calibration sizes, demonstrating strong robustness to calibration data selection. Moreover, increasing the calibration set beyond 128 samples yields only marginal improvements at all sparsity levels, indicating that 128 samples are sufficient for reliable pruning.

\begin{table}[t]

4pt

Num. Calib.	Sparsity	FID \(\downarrow\)	sFID \(\downarrow\)	IS \(\uparrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
128	20%	10.32	250.38	37.52	0.79	0.853
	40%	12.75	194.92	40.46	0.74	0.649
	50%	24.14	111.01	53.50	0.59	0.560
1-7	20%	10.31	251.53	37.59	0.78	0.853
	40%	12.76	198.04	40.36	0.74	0.649
	50%	23.99	112.32	53.86	0.59	0.562
1-7	20%	10.47	250.96	37.45	0.77	0.854
	40%	12.76	197.51	40.33	0.73	0.651
	50%	23.98	111.94	53.82	0.59	0.561
1-7	20%	10.45	251.55	37.32	0.78	0.853
	40%	12.75	197.33	40.32	0.75	0.649
	50%	24.01	112.03	53.01	0.58	0.560

\vskip

0.1in \end{table}

Selection of Calibration Samples. We further investigate the impact of calibration data selection. Specifically, we fix the calibration set size to 128 and sample calibration data from different diffusion timesteps. As shown in Figure 7, pruning performance is highly sensitive to timestep coverage. Calibration data drawn exclusively from early or late stages consistently lead to inferior results, whereas samples covering a broad timestep range achieve substantially better performance. This observation suggests that different timesteps exhibit distinct activation distributions and parameter sensitivity patterns in DiTs. Therefore, effective calibration data should provide diversity not only in sample content but also across diffusion timesteps, ensuring representative activation statistics and more reliable pruning decisions.

10 More Ablation Materials↩︎

Due to space limitations, the main text reports ablation results only on DiT-XL/2 at 40% and 50% sparsity and N:M structured pruning results on FLUX. Table 2 extends the study to PixArt-\(\Sigma\) across different resolutions and sparsity levels. Similar to the observations on DiT-XL/2, STW consistently improves the baseline, and CAG provides additional gains when combined with STW. The resulting indicates that the two components are complementary and generalize well across different DiT-based architectures.

4pt

Effect of CFG Scale. We evaluate pruning performance under different classifier-free guidance (CFG) scales while keeping the sparsity ratio and denoising steps fixed. Following the official DiT setup, CFG=1.5 is used as the default configuration, while CFG=1 and CFG=2 are adopted to examine weaker and stronger guidance strengths, respectively. As shown in Table ¿tbl:tab:cfg95ablation?, varying CFG changes the trade-off between conditional guidance and generation diversity. Nevertheless, the proposed method consistently achieves the best performance across all CFG settings, maintaining its advantage under both weaker and stronger guidance. These results demonstrate the robustness of our method to variations in guidance strength and sampling configuration.

\begin{table}[ht]

4pt

Timesteps	Sparsity	CFG Scale	Method	FID \(\downarrow\)	IS \(\uparrow\)	sFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
50	50%	2.0	Magnitude	38.14	72.04	40.08	0.44	0.516
			Wanda	33.39	81.99	77.78	0.49	0.491
			DiT-Pruning	10.82	199.70	34.02	0.73	0.565
		1.5	Magnitude	62.40	37.98	50.50	0.32	0.508
			Wanda	57.32	38.84	101.30	0.34	0.473
			DiT-Pruning	20.29	111.95	38.95	0.60	0.559
		1.0	Magnitude	97.76	18.62	67.32	0.22	0.220
			Wanda	90.50	17.23	135.69	0.21	0.316
			DiT-Pruning	46.54	43.67	48.98	0.42	0.519

\vskip

-0.1in \end{table}

Effect of Optional \((m,n)\) Settings. To validate the choice of the proposed squared transformation, we perform an ablation study by varying the exponents \((m,n)\) in Eq. 6 , while keeping the CAG strategy fixed. As shown in Table ¿tbl:tab:mn95settings?, the proposed setting \((m=2,n=1)\) consistently achieves the best overall performance, yielding the lowest FID and highest IS, sFID, PRE, and SSIM. In contrast, both smaller and larger weight exponents lead to inferior results. These observations suggest that \((m=2,n=1)\) provides a more suitable balance between weight and activation contributions, supporting the effectiveness of the proposed STW.

\begin{table}[ht]

4pt

Timesteps	Sparsity	Settings	FID \(\downarrow\)	IS \(\uparrow\)	sFID \(\downarrow\)	PRE \(\uparrow\)	SSIM \(\uparrow\)
50	50%	\(m=1,\;n=1\)	22.23	100.91	41.71	0.57	0.557
		\(m=0.5,\;n=1\)	41.92	53.64	57.71	0.42	0.521
		\(m=3,\;n=1\)	22.76	104.36	42.27	0.57	0.547
		\(m=2,\;n=1\)	20.29	111.95	38.95	0.60	0.559

\vskip

-0.1in \end{table}

11 Additional Visualization Results↩︎

Figure 8, Figure 9-10, and Figure 11-14 provide qualitative comparisons between DiT-Pruning, Wanda, and the dense models under 50% sparsity level. Across all settings, our method generates images that remain visually closer to the dense model, preserving both structural details and semantic consistency. In contrast, Wanda exhibits noticeable degradation as sparsity increases, including distorted structures, blurred content, and weakened semantic alignment. These results demonstrate that DiT-Pruning effectively preserves generation quality under aggressive sparsification and generalizes consistently across different DiT-based architectures.

12 Limitations and Broader Impacts↩︎

Our method demonstrates effective pruning of DiTs, achieving favorable generation quality at different sparsity levels (e.g., 20%, 40%, 50%). However, several limitations remain. Under higher sparsity ratios (e.g., 70% and 80%), generated images exhibit noticeable degradation, including artifacts, reduced fidelity, and inaccurate semantics. This indicates that our pruning strategy, while efficient, cannot fully preserve model performance under extreme sparsity and suggests the trade-off between compression ratio and generation quality. Moreover, DiTs appear inherently more sensitive to pruning than LLMs, highlighting the need for strategies that better exploit their denoising dynamics and timestep-dependent characteristics. Finally, our experiments focus on standard image generation benchmarks, and the robustness of our approach in more diverse visual domains remains unexplored.

References↩︎