Diff-ES: Stage-wise Structural Diffusion Pruning via Evolutionary Search

Zongfang Liu; Shengkun Tang; Zongliang Wu; Xin Yuan; Zhiqiang Shen

cs.CVMar 5, 2026

Diff-ES: Stage-wise Structural Diffusion Pruning via Evolutionary Search

Authors: Zongfang Liu, Shengkun Tang, Zongliang Wu, Xin Yuan, Zhiqiang Shen

Organizations: Zhejiang University · Westlake University · Mohamed bin Zayed University of Artificial Intelligence

Abstract

Diffusion models have achieved remarkable success in high-fidelity image generation but remain computationally demanding due to their multi-step denoising process and large model sizes. Although prior work improves efficiency either by reducing sampling steps or by compressing model parameters, existing structured pruning approaches still struggle to balance real acceleration and image quality preservation. In particular, prior methods such as MosaicDiff rely on heuristic, manually tuned stage-wise sparsity schedules and stitch multiple independently pruned models during inference, which increases memory overhead. However, the importance of diffusion steps is highly non-uniform and model-dependent. As a result, schedules derived from simple heuristics or empirical observations often fail to generalize and may lead to suboptimal performance. To this end, we introduce Diff-ES, a stage-wise structural Diffusion pruning framework via Evolutionary Search, which optimizes the stage-wise sparsity schedule and executes it through memory-efficient weight routing without model duplication. Diff-ES divides the diffusion trajectory into multiple stages, automatically discovers an optimal stage-wise sparsity schedule via evolutionary search, and activates stage-conditioned weights dynamically without duplicating model parameters. Our framework naturally integrates with existing structured pruning methods for diffusion models including depth and width pruning. Extensive experiments on DiT and SDXL demonstrate that Diff-ES consistently achieves wall-clock speedups while incurring minimal degradation in generation quality, establishing state-of-the-art performance for structured diffusion model pruning.

Figures & tables

Figure 1: Comparison of sparsity ( $1-\text{density}$ ) schedules. An effective stage-wise sparsity schedule is crucial for maintaining image quality. Using the same structural second-order pruning method, our Diff-ES framework significantly outperforms MosaicDiff on SDXL by employing an optimized, adaptive stage-wise sparsity schedule. For fair comparison, the original 3-stage schedule of MosaicDiff is visualized in 10.

Figure 2: Visual comparison on SDXL-base-1.0 under 30% sparsity. Each row shows generations from the same prompts using 20 sampling steps (CFG 7.5): Dense, OBS-Diff, MosaicDiff with OBS, and Diff-ES with OBS. At the same sparsity level, Diff-ES remains closest to the dense model, preserving object identity, scene layout, and fine textures. OBS-Diff already shows clear structural differences with dense model (e.g., the teddy bear is rendered with three legs), while MosaicDiff exhibits much more severe semantic and perceptual degradation. This visual trend is consistent with FID.

Figure 3: Overview of Diff-ES . We divide the diffusion process into stages and evolve an optimal stage-wise sparsity schedule under a fixed global budget. A level-switch mutation redistributes sparsity across stages, while lightweight fitness evaluation (e.g., TOPIQ/CLIP-IQA/SSIM) guides survivor selection. For each stage, SNR-aware calibration supports pruning methods to obtain calibrated projection orders. For methods with weight updating such as OBS, all updated weights are stored once in a compact database and retrieved during evaluation through our efficient weight-routing mechanism, enabling rapid model assembly without recomputing second-order updates.

Pruning Strategy	Method	Sparsity (%)	MACs $\downarrow$	Latency (ms) $\downarrow$	Speedup $\uparrow$	IS $\uparrow$	FID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
Dense	Original	0%	2.37	83.60	1.00 $\times$	329.41	4.49	0.87	0.48
Diff-Pruning	Original	20%	1.90	81.72	1.03 $\times$	156.36	15.52	0.73	0.27
LayerDrop	MosaicDiff	17.86%	2.02	73.34	1.14 $\times$	1.50	366.96	0.00	0.00
LayerDrop	Diff-ES (Ours)	17.80%	2.02	71.22	1.17 $\times$	36.30	84.13	0.22	0.40
Wanda	MosaicDiff	12.50%	2.13	87.13	0.96 $\times$	66.57	43.89	0.39	0.46
	Diff-ES (Ours)	12.50%	2.13	78.87	1.06 $\times$	253.79	5.28	0.79	0.50

Table 1: Result comparison of baselines and Diff-ES (Ours) on DiT-XL/2. All results are reported using 20 sampling steps for a fair comparison. Diff-ES achieves the best performance compared with all baselines including MosaicDiff on different pruning strategies, demonstrating the effectiveness and generalizability of our method.

Pruning Strategy	Method	Sparsity (%)	FID $\downarrow$	CLIP-Score $\uparrow$	SSIM $\uparrow$
Dense	Original	0%	25.46	0.3219	1.000
Deep-Cache	Original	—	26.32	0.3185	0.6709
Diff-Pruning	Original	20%	96.5	0.2702	0.5887
OBS-Diff	Original	30%	28.49	0.3215	0.6760
LayerDrop	MosaicDiff	10%	25.16	0.3198	0.6341
	Diff-ES (Ours)	10%	24.12	0.3214	0.767

Table 2: Result comparison of baselines and Diff-ES (Ours) on SDXL-Base-1.0 with 20 sampling steps. Our method consistently achieves better perceptual quality across pruning strategies and sparsity levels, with particularly strong gains at higher sparsity.

Model	Method	Latency (ms) $\downarrow$	Runtime (h) $\downarrow$	Speedup $\uparrow$	FID $\downarrow$
	Dense	83.60	0	1	4.49
DiT	MDiff	60.96	0.64	1.38	6.13
	Diff-ES	61.20	18.87	1.37	4.25
	Dense	1627.69	0	1	25.46
SDXL	MDiff	1379.40	6.62	1.18	59.09
	Diff-ES	1479.72	3.61	1.10	25.88

Table 3: Efficiency and quality comparison on DiT (

\sim

40% sparsity) and SDXL (30% sparsity). Diff-ES achieves strong acceleration while preserving high generative fidelity. For latency, we use a batch size of 8. Runtime denotes the one-time computational cost incurred by calibration and evolutionary search. For SDXL, we use SSIM as the fitness metric and perform a single calibration shared across stages. The number of generations is set to 30.

Panel of Figure 4 — Figure 4: Stage-wise sparsity schedules of Diff-ES and MosaicDiff across models. DiT: Diff-ES vs MosaicDiff (FID: 3.63 vs 4.01); SDXL: Diff-ES vs MosaicDiff (FID: 33.12 vs 98.56). On DiT, the sparsity schedules are similar between MosaicDiff and Diff-ES . On SDXL, the sparsity schedules diverge, indicating that MosaicDiff’s empirical schedule overfits to DiT. Stages 0–9 correspond to timesteps 0–999; sampling starts at timestep 999, and earlier stages correspond to earlier starting timesteps.

Fitness	FID $\downarrow$	CLIP $\uparrow$	SSIM $\uparrow$
CLIP-IQA	25.87	0.3214	0.7118
TOPIQ	26.04	0.3206	0.6777
SSIM	26.03	0.3223	0.8220

Table 4: Evaluation of Diff-ES + OBS under different fitness metrics on SDXL at 30% sparsity.

Setting	Memory (GB) $\downarrow$
Model Stitching	9.79
Weight Routing	5.61

Table 6: Loading memory comparison of Diff-ES with model stitching and proposed weight routing. We report results for DiT-XL/2 at

\sim

40% sparsity, the number of stages is 10. Weight routing saves

\sim

42.7% GPU memory compared to model stitching.

Figure 5: Out-of-distribution visual examples on AI-generated prompts (disjoint from the MS-COCO search set) across sparsity levels under the same fitness metric (CLIP-IQA). Diff-ES preserves semantic consistency and fine details up to moderate sparsity, with more visible degradation mainly at the highest sparsity level.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Steps	MACs (T)	FID $\downarrow$
Low-step comparison
Dense	20	67.64	25.46
Diff-ES	20	52.45	25.87
DeepCache + Diff-ES	20	36.05	27.10
DPM-Solver + Diff-ES	10	26.23	25.67

Appendix

Table 7: Compatibility of Diff-ES with different acceleration and sampling techniques on SDXL. The table demonstrates that Diff-ES seamlessly integrates with existing methods such as DeepCache and DPM-Solver, maintaining competitive FID while significantly reducing computation. We use Diff-ES -0.3 and DeepCache-N=2 here. The fitness metric for Diff-ES is CLIP-IQA.

Method	Fitness	Sparsity	FID $\downarrow$	CLIP-Score $\uparrow$	SSIM $\uparrow$
Dense	–	0%	25.46	0.3219	1.000
Diff-ES (Ours)	CLIP-IQA	10%	25.42	0.3211	0.7725
		20%	25.34	0.3215	0.7146
		30%	25.87	0.3214	0.7118
		40%	28.46	0.3209	0.6980
		50%	41.20	0.3158	0.6994

Appendix

Table 8: Comparison of OBS pruning under different fitness metrics on SDXL-Base-1.0.

Method	Sparsity	FID $\downarrow$	CLIP-Score $\uparrow$	SSIM $\uparrow$
Dense	0%	31.16	0.3187	1.0000
MosaicDiff	20%	30.16	0.3084	0.4602
Diff-ES (Ours)	20%	25.88	0.3150	0.6782

Appendix

Table 9: Comparison of MosaicDiff and our Diff-ES with Layer Dropping on FLUX.1-schnell. The CFG scale is set to 3.5, which is the default value of FLUX.1-schnell.

Batch Size	FULL (ms) $\downarrow$	PRUNED (ms) $\downarrow$	Speedup $\uparrow$
1	158.7	188.5	0.842 $\times$
2	201.5	200.9	1.003 $\times$
8	670.3	489.6	1.369 $\times$
16	1320.9	1021.6	1.293 $\times$
32	2877.4	2366.3	1.216 $\times$

Appendix

Table 10: Overall latency and speedup on DiT with 43.75% (

\sim

40%) sparsity under different batch sizes. Latency is reported in milliseconds (ms), and speedup is computed as FULL/PRUNED.

Search Method	IS $\uparrow$	FID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
Greedy Search	83.66	33.75	0.5028	0.5758
Evolutionary Search ( Diff-ES )	156.23	12.86	0.66	0.59

Appendix

Table 11: Comparison between evolutionary search and greedy search on DiT-XL/2 with OBS pruning at 50% sparsity with fitness TOPIQ. The number of stages is 10.

Diffusion Model Pruning

See all 12 papers

Jul 7, 2026

Bridging Diffusion Pruning and Step Distillation with Teacher-Aligned Repair

Diffusion models generate high-quality images, but their inference cost comes from two sources: large denoising networks and repeated denoising steps. Existing compression pipelines usually attack these costs separately. Pruning reduces the network, but most pruning methods still rely on a long post-pruning retraining stage to recover a many-step sampler. Step distillation reduces the number of denoising steps, but it usually assumes a student that can already follow the teacher well enough to receive useful distillation gradients. This paper asks whether post-pruning retraining can be replaced by step distillation. We find that the direct replacement fails: after pruning an EDM2-XS teacher, starting SiDA from the pruned checkpoint produces unusable samples. We introduce a short teacher-alignment repair stage as a bridge between pruning and step distillation. The bridge matches the pruned generator to the teacher on noisy real-image latents, then hands the repaired checkpoint to one-step distillation. On ImageNet-512, the original EDM2-XS baseline uses 124.713M parameters and 63 network evaluations, reaching an FID of 3.53. With a suitable distillation objective, our 20% pruned one-step generator uses 98.826M parameters and one network evaluation, reaching an FID of 3.12. With 30% pruning, the model uses 88.029M parameters and one network evaluation, with an FID of 4.26.

Jincheng Ying, Li Wenlin, Minghui Xu +1

School of Big Data and Artificial Intelligence Guangdong University of Finance and Economics Guangzhou, China · School of Computer Science Shandong University Qingdao, China · GDUFE-UCM Joint Institute Guangdong University of Finance and Economics Guangzhou, China

Diffusion Model Distillation Few-Step Diffusion Sampling Diffusion Model Pruning Diffusion Model Inference Acceleration Diffusion Model Compression

Ask PDF

Jul 22, 2026

Importance-Aware OBS Pruning for Diffusion Models

We propose importance-aware pruning for diffusion models, a training-free framework that prioritizes preserving parameters critical to semantically salient image regions. To do so, we incorporate spatial importance maps -- derived from conditioning signals or model attention -- into the pruning objective. This produces parameter rankings aligned with perceptual relevance rather than uniform reconstruction error. On MS-COCO dataset, our proposed approach consistently retains subject fidelity and structural correctness at high compression ratios where conventional pruning causes visible degradation. These results demonstrate that content-aware objectives are key to perceptually faithful compression of generative models.

Ba-Thinh Lam, Srijan Das, Hieu Le

Department of Computer Science UNC Charlotte NC, USA

Diffusion Model Pruning Neural Network Compression Neural Network Pruning Diffusion Model Compression

Ask PDF

Apr 4, 2026

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Previous works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inefficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations in several steps in these token cache methods. To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise sparsity allocation in an end-to-end manner through a learnable network combined with a dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency. We conducted extensive experiments on a range of diffusion-transformer models, including DiT-XL/2, PixArt-

α

, FLUX, and Wan2.1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt-

α

with 20 sampling steps, we reduce computational cost by

54\%

while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.

Haowei Zhu, Ji Liu, Ziqiong Liu +4

Advanced Micro Devices, Inc. · Tsinghua University · BNRist

Diffusion Model Inference Acceleration Diffusion Models Diffusion Transformer Diffusion Model Caching

Ask PDF

Abstract

Diffusion Model Pruning

Explore similar work

Bridging Diffusion Pruning and Step Distillation with Teacher-Aligned Repair

Importance-Aware OBS Pruning for Diffusion Models

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity