Research map / Image generation / Text-to-image diffusion
Not All Prompts Are Made Equal: Prompt-based Pruning of Text-to-Image Diffusion Models
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.
Prunes Stable Diffusion V2.1 into a mixture of prompt-routed pruned experts (APTP) using a contrastive-trained prompt router with optimal-transport-balanced architecture codes, fine-tuned on CC3M and MS-COCO.
Problem
Static pruning of text-to-image diffusion models uses one sub-network for all prompts despite varying prompt complexity, while dynamic per-prompt pruning breaks GPU batch parallelism; the paper targets prompt-adaptive pruning of T2I models for a fixed compute budget on a target dataset.
Approach
A prompt router maps each prompt embedding (frozen Sentence Transformer + feed-forward predictor) to one of N learnable architecture codes, each a Gumbel-sigmoid width/depth mask defining an expert sub-network; contrastive loss makes similar prompts share codes and Sinkhorn optimal transport with equipartition prevents code collapse, so experts can be batched.
Key results
- On CC3M, APTP(0.85) reaches FID 36.77 and CMMD 0.675 vs. 141.04 and 1.646 for Norm pruning at @50k with similar MACs (1182.8G vs 1185.3G).
- On MS-COCO, APTP(0.78) reaches FID 22.60, CLIP 31.32, CMMD 0.569 at 1076.6G MACs, outperforming BKSDM (FID 26.31) at similar cost.
- Ablation on MS-COCO: optimal transport lowers FID from 48.78 to 38.56 and distillation further to 25.57 for the contrastive router.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use prompt-based APTP pruning at 0.85 budget instead of static SP pruning | Stable Diffusion 2.1 · CC3M validation | FID: 75.81 -> 36.77 | Helped |
| Use prompt-based APTP pruning at 0.85 budget instead of static SP pruning | Stable Diffusion 2.1 · CC3M validation | CLIP: 26.83 -> 30.84 | Helped |
| Use prompt-based APTP pruning at 0.85 budget instead of static SP pruning | Stable Diffusion 2.1 · CC3M validation | CMMD: 1.243 -> 0.675 | Helped |
| Use prompt-based APTP pruning at 0.64 budget instead of static pruning | Stable Diffusion 2.1 · MS-COCO validation | FID: 26.31 -> 39.12 | Mixed |
| Use prompt-based APTP pruning at 0.78 budget instead of static SP pruning | Stable Diffusion 2.1 · MS-COCO validation | CLIP: 28.98 -> 31.32 | Helped |
| Use prompt-based APTP pruning at 0.78 budget instead of static BKSDM pruning | Stable Diffusion 2.1 · MS-COCO validation | CMMD: 0.611 -> 0.569 | Helped |
| Use contrastive-only prompt routing rather than a single architecture (ablation) | Stable Diffusion 2.1 pruned to an expert mixture · MS-COCO validation | FID: 46.56 -> 48.78 | Hurt |
| Add optimal transport to the contrastive router (ablation) | Stable Diffusion 2.1 pruned to an expert mixture · MS-COCO validation | FID: 48.78 -> 38.56 | Helped |
| Add optimal transport to the contrastive router (ablation) | Stable Diffusion 2.1 pruned to an expert mixture · MS-COCO validation | CLIP: 28.90 -> 30.07 | Helped |
| Add optimal transport to the contrastive router (ablation) | Stable Diffusion 2.1 pruned to an expert mixture · MS-COCO validation | CMMD: 0.92 -> 0.74 | Helped |
| Add distillation to the contrastive router with optimal transport (ablation) | Stable Diffusion 2.1 pruned to an expert mixture · MS-COCO validation | FID: 38.56 -> 25.57 | Helped |
| Vary the number of experts (ablation) | APTP on Stable Diffusion 2.1 · MS-COCO | FID | Mixed |
Surprising findings
- Contrastive training alone fails to improve over a single pruned model, because routing collapses most prompts to one expert.
- The router automatically assigns text-image prompts and human-centric prompts to the highest-budget expert, matching known SD failure modes.
Limitations
- Limited experimental scope: only two target datasets and one base model (SD V2.1).
- Inherits SD's failures on text rendering, humans, and fingers.
- Optimal number of experts is dataset-dependent and must be tuned.
Experimental setup
Datasets: CC3M MS-COCO 2014 PartiPrompts
Benchmarks: CC3M validation MS-COCO 2014 validation PartiPrompts
Metrics: FID CLIP score CMMD PickScore MACs latency
Baselines: Norm pruning Structural Pruning (SP) BK-SDM Stable Diffusion V2.1
Builds on: Stable Diffusion BK-SDM Sinkhorn-Knopp Gumbel-sigmoid reparameterization Sentence Transformers
Closest papers by meaning
- On Discrete Prompt Optimization for Diffusion Models (ICML 2024 · similarity 0.91)
- AP-Adapter: Improving Generalization of Automatic Prompts on Unseen Text-to-Image Diffusion Models (NeurIPS 2024 · similarity 0.91)
- Prompting Hard or Hardly Prompting: Prompt Inversion for Text-to-Image Diffusion Models (CVPR 2024 · similarity 0.90)
- Prompt-Free Diffusion: Taking "Text" out of Text-to-Image Diffusion Models (CVPR 2024 · similarity 0.89)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.89)
- Learning to Sample Effective and Diverse Prompts for Text-to-Image Generation (CVPR 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.