Research map / Image generation / Fast diffusion sampling
DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DiP-GO, which casts diffusion pruning as a SubNet search over a SuperNet with backup connections, using a plugin pruner network trained with few-step gradient optimization (SSIM consistency + sparsity loss) to prune SD and DiT models without retraining.
Problem
Multi-step denoising makes diffusion model inference expensive; existing structural pruning needs retraining with large datasets, and training-free caching heuristics (e.g., DeepCache) use simple static strategies that leave redundancy unexploited.
Approach
Build a SuperNet by adding backup input connections from each block to its counterpart at the previous timestep (exploiting feature similarity across adjacent steps), then learn a differentiable plugin pruner with learnable queries and straight-through gates that selects a SubNet by optimizing SSIM consistency and a FLOPs-weighted sparsity loss, followed by bisection-based thresholding to hit a ta
Key results
- On SD-1.5 (PartiPrompts/COCO2017), DiP-GO at 0.80 pruning reaches 4.43x/4.40x speedup with CLIP Score 29.51/30.29, matching the 50-step PLMS baseline (30.30 on COCO).
- On DiT-XL/2 ImageNet, pruning 0.6 gives 2.43x speedup with FID-50K 3.01 vs 2.97 for the unpruned reproduced baseline; pruning 0.75 gives 3.60x with FID 3.14.
- On SD-1.5 at 80% sparsity, the SSIM consistency loss gives CLIP 30.29 vs 29.94 (L1), 29.71 (L2), 29.77 (L1+SSIM).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| prune 80% of SD-1.5 computation with DiP-GO | SD-1.5 · PartiPrompts | CLIP Score: 29.51 -> 29.51 | No clear effect |
| prune 80% of SD-1.5 computation with DiP-GO | SD-1.5 · COCO2017 | CLIP Score: 30.30 -> 30.29 | No clear effect |
| apply DiP-GO pruning at ratio 0.7 | SD-2.1 · COCO2017 validation | FID-5K: 27.29 -> 25.98 | Helped |
| prune DiT-XL/2* at ratio 0.6 | DiT-XL/2* · ImageNet validation | FID-50K: 2.97 -> 3.01 | Mixed |
| use SSIM instead of L1 consistency loss (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning 80% | CLIP Score: 29.94 -> 30.29 | Helped |
| use SSIM instead of L2 consistency loss (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning 80% | CLIP Score: 29.71 -> 30.29 | Helped |
| use SSIM instead of L1 + SSIM consistency loss (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning 80% | CLIP Score: 29.77 -> 30.29 | Helped |
| set α to 1.0 rather than 0.1 (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning ratio 0.80 | CLIP Score: 29.77 -> 30.29 | Helped |
| set α to 1.0 rather than 0.5 (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning ratio 0.80 | CLIP Score: 29.93 -> 30.29 | Helped |
| set α to 1.0 rather than 2.0 (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning ratio 0.80 | CLIP Score: 30.17 -> 30.29 | Helped |
| remove γ weighting from the sparsity loss (ablation) | DiP-GO on SD-1.5 · COCO2017 validation, pruning 80% | CLIP Score: 30.29 -> 29.50 | Hurt |
| use gradient-optimized DiP-GO instead of random search (ablation) | SD-1.5 · Pruning ratio 0.80 | CLIP Score: 28.73 -> 30.29 | Helped |
Surprising findings
- Pruner search (2.3 GPU hours) is far cheaper than random or genetic search, and those baselines reach much lower CLIP scores at 90% pruning (24.07/25.14 vs 28.72).
- Pruning gives low benefit on LCM few-step samplers (11.95T to 11.58T MACs) because adjacent-step feature redundancy is low.
Limitations
- Requires training a separate pruner network per pretrained diffusion model, adding setup time.
- Performance degrades at extremely high pruning ratios (e.g., 90%).
- Limited acceleration on few-step LCM samplers due to low feature redundancy across timesteps.
Experimental setup
Datasets: DiffusionDB PartiPrompts MS-COCO 2017 ImageNet
Benchmarks: PartiPrompts MS-COCO 2017 ImageNet
Metrics: MACs speedup CLIP Score FID
Baselines: PLMS BK-SDM DeepCache DDIM DPM-Solver LCM Random Search GA Search PixArt-alpha
Builds on: DeepCache DETR Lottery Ticket Hypothesis gradient checkpointing
Closest papers by meaning
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.92)
- Attention-Driven Training-Free Efficiency Enhancement of Diffusion Models (CVPR 2024 · similarity 0.90)
- Training-Free and Hardware-Friendly Acceleration for Diffusion Models via Similarity-based Token Pruning (AAAI 2025 · similarity 0.90)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.89)
- TinyFusion: Diffusion Transformers Learned Shallow (CVPR 2025 · similarity 0.89)
- Fast constrained sampling in pre-trained diffusion models (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.