Not All Prompts Are Made Equal: Prompt-based Pruning of Text-to-Image Diffusion Models

Authors: Alireza Ganjdanesh, Reza Shirkavand, Shangqian Gao, Heng Huang
Primary area: generative models
Keywords: Model Pruning, Diffusion Models, Inference Efficiency

Abstract:
Text-to-image (T2I) diffusion models have demonstrated impressive image generation capabilities. Still, their computational intensity prohibits resource-constrained organizations from deploying T2I models after fine-tuning them on their internal *target* data. While pruning techniques offer a potential solution to reduce the computational burden of T2I models, static pruning methods use the same pruned model for all input prompts, overlooking the varying capacity requirements of different prompts. Dynamic pruning addresses this issue by utilizing a separate sub-network for each prompt, but it prevents batch parallelism on GPUs. To overcome these limitations, we introduce Adaptive Prompt-Tailored Pruning (APTP), a novel prompt-based pruning method designed for T2I diffusion models. Central to our approach is a *prompt router* model, which learns to determine the required capacity for an input text prompt and routes it to an architecture code, given a total desired compute budget for prompts. Each architecture code represents a specialized model tailored to the prompts assigned to it, and the number of codes is a hyperparameter. We train the prompt router and architecture codes using contrastive learning, ensuring that similar prompts are mapped to nearby codes. Further, we employ optimal transport to prevent the codes from collapsing into a single one. We demonstrate APTP's effectiveness by pruning Stable Diffusion (SD) V2.1 using CC3M and COCO as *target* datasets. APTP outperforms the single-model pruning baselines in terms of FID, CLIP, and CMMD scores. Our analysis of the clusters learned by APTP reveals they are semantically meaningful. We also show that APTP can automatically discover previously empirically found challenging prompts for SD, *e.g.,* prompts for generating text images, assigning them to higher capacity codes.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730468665592

thread_type:
official_review

<thread_object>
creator: reviewer_xwxf

<content>
summary:
The authors propose a mixture-of-expert-esque strategy for efficiently, which they coin Adaptive Prompt Tailored Pruning (APTP). The methods combine the benefits of dynamic and static pruning methods and archives good generative metrics while decreasing the computational cost.

strengths:
* Good Illustrations and concise explanation
* The core idea is an interesting combination of known concepts in a new, but highly relevant setup.
* A good amount of comparison to other methods using multiple datasets with different image-types (COCO being primarily real-world images and CC3M being more diverse)

weaknesses:
* While FID, CLIP-Score and CMMD alongside the visual examples provide a good overview, I, personally, would have preferred some human-user study (which I see is a lot of effort and for this reason, would not ask from the authors). As an alternative substitute, I propose compute the pick-scores [Kirstain et al. 2023] against the pruning metrics similar to [Pernias et al 2024] on a set of diverse prompts like Partiprompts could provide additional, easily interpretable evidence of this methods' superiority in terms of generation quality.

Kirstain et al. 2023 https://proceedings.neurips.cc/paper_files/paper/2023/hash/73aacd8b3b05b4b503d58310b523553c-Abstract-Conference.html
[Pernias et al. 2024] https://proceedings.neurips.cc/paper_files/paper/2023/hash/73aacd8b3b05b4b503d58310b523553c-Abstract-Conference.html

questions:
* I would be curious to know how big the router-model itself is (in terms of parameters and memory footprint) and, by extension, how does the size of the router model affect the maximum batch size on an A100 GPU?
* Did you do any additional experiments concerning varying resolutions and aspect ratios and their impact on the pruned-image quality?
* Did you try applying this technique to transformer-based image-generation models like Pix-Art-Alpha / Sigma, do you see any major hurdles regarding the switch from ConvNets to Transformers?
* How specific is the APTP-model to its training data? How do out-of-distribution prompts impact the quality of the generations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730565265861

thread_type:
official_review

<thread_object>
creator: reviewer_vgxh

<content>
summary:
This paper presents a new approach for accelerating the sampling from diffusion models using adaptive pruning denoted as APTP. The method is tailored for a specific scenario where the inference is on a specific known data distribution and the data for it is given (e.g. a company’s private dataset). Using this data, APTP trains a prompt router module that selects which parts of the architecture should be pruned for each given prompt. The selection is from a fixed set of reduced architecture states (denoted as architecture codes). Both the prompt router and arch. codes are trained end-to-end.

strengths:
- The idea of pruning different parts of the network for each prompt is non-trivial and interesting.
- Visual results show APTP seems to use various pruning modes, and does not collapse to a single one.

weaknesses:
**(1) Missing Baselines of Less Diffusion Steps:** My main concern is that the paper does not compare to approaches that perform less diffusion steps. More specifically the following should be considered as baselines:

- *Step Skipping Distillation of Diffusion Models [1,2,3]:* As mentioned in the paper, several methods tackled faster inference of diffusion models using this idea, mostly by knowledge distillation of the model to inference with fewer (between 1-4) steps. These methods should cut the latency by 8-25 times from the original model, while the proposed APTP was shown to cut latency by much less (not even 50% of the original model timing). These approaches also require training as APTP does, and their weights are readily available for many SoTA architectures, e.g. SDXL, SD 1.5, etc. 

- *Caching Based Acceleration [4]:* These approaches cache features of the U-Net architecture and reuse them in later steps. Such approaches hold a distinct advantage of being training-free.

- *Less Denoising Steps at Inference Time:* A trivial training-free baseline is to simply sample images with less denoising steps. I wonder how that would compare to the proposed APTP. Can APTP further accelerate such a setting?

As step skipping approaches became more popular recently, I believe including some of these as baselines is a must for such an approach.

**(2) Quality of Writing (Sec. 3):** Sec. 3 is complicated and difficult to understand. Specifically, I think Sec. 3.2,3.3 are overloaded and could be shortened to a much clearer version. These subsections are the only place in the paper that describe the actual approach of APTP (and what does arch. codes mean), therefore are especially important for the proposed approach.

**(3) Visual Comparisons:** The paper offers a very limited selection of visual comparisons - having only 8 qualitative comparisons to the original baseline, and no such comparisons to previous approaches. Could the authors please supply a larger set of comparisons to the original model and baselines (including a step skipping distillation [1,2,3], caching [4] and less denoising steps). 

**(4) Clustering of Pruning Modes:** While this qualitative analysis was promised in the abstract and introduction, it only exists in figures at the last 2 pages of the appendix without any textual description. Given it is mentioned in the abstract I think it should be included as a part of the main manuscript.

**(5) Limited Setting and Empirical Results:** Unlike other approaches, the proposed method is limited to a specific data distribution. Although the task is much more confined, the reduction in latency is not substantial: To keep performance comparable to the original model in terms of FID or CLIP score, APTP can only reduce 20% of the latency (Tab.1). 

[1] Liu, Xingchao, et al. "Instaflow: One step is enough for high-quality diffusion-based text-to-image generation." The Twelfth International Conference on Learning Representations. 2023.

[2] Sauer, Axel, et al. "Adversarial diffusion distillation." European Conference on Computer Vision. Springer, Cham, 2025.

[3] Meng, Chenlin, et al. "On distillation of guided diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[4] Ma, Xinyin, Gongfan Fang, and Xinchao Wang. "Deepcache: Accelerating diffusion models for free." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

questions:
- Did the authors try their approach on other architectures? even other backbones of Stable Diffusion?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730659705881

thread_type:
official_review

<thread_object>
creator: reviewer_uiaa

<content>
summary:
The authors propose prompt-based tuning of text-to-image diffusion models, in which different sub-networks within pre-trained models are trained for different prompts/concepts. They authors performe experiments on multiple datasets to show that for a given latency, their model performs comparably to higher latency pretrained models.

strengths:
1. The idea behind the paper is technically sound and novel.
2. The paper is well written. 
3. The authors present some interesting interepretability experiments on expert assignment that aid the proposed concepts.

weaknesses:
1. Limited experiments - Although I find the proposed ideas novel, I believe that the paper lacks extensive experimentation on 
 - different types of architecture - It is currently unknown if the proposed methods is generalizable across architectures (DiT, MMDiT etc). 
 - Small datasets - How does the method perform when data is limited?
 - Fine grained concepts - How does their method handle expert assignment when concepts are fine-grained (breeds of different animals)
2. Comparison to Mixture-of-Experts (MoE) models - How does the proposed method compare to other prompt-conditinal architectures like MoE text-to-image diffusion models? Currently the competitors in Table 1 (a and b) are static structural pruning baselines, but I believe the paper's contribution is prompt-conditional pruning, which demands comparison to prompt-conditional efficient architectures like MoEs like [1].
3. I am concerned about the 4 and 7 point drop in FID of the proposed method in Table 1. The authors have not presented any trade-off between latency and performance, which would help understand how limiting computational budget affects performance


[1] RAPHAEL: Text-to-Image Generation via Large Mixture of Diffusion Paths, Xue et al

questions:
See Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730689874046

thread_type:
official_review

<thread_object>
creator: reviewer_ff7o

<content>
summary:
The paper introduces adaptive prompt-based pruning strategy to reduce the computation cost of diffusion model. The proposed approach involves encoding input prompts into architecture embeddings, which are mapped to specialized architecture codes. These codes determine the routing of each prompt to a pruned sub-network. By training a prompt router using a combination of contrastive learning and optimal transport, the proposed method ensures that prompts are dynamically assigned to appropriate sub-networks. The results of the paper demonstrate the reduction in computational cost while maintaining FID and CLIP scores.

strengths:
1. The paper introduces a novel approach by proposing adaptive prompt-based pruning that routes input prompts to specialized pruned sub-networks based on their characteristics. This represents a difference from conventional static and dynamic pruning methods,
2. The empirical results training on datasets like CC3M and MS-COCO demonstrate the method’s effectiveness compared to other pruning methods. The results show that the proposed method outperforms other baselines by significantly reducing computational cost while maintaining or improving output quality as measured by metrics like FID, CLIP score, and CMMD score.

weaknesses:
The major concern is the empirical evaluation of the proposed method: 

1. as stated in the paper, most organizations typically fine-tune pre-trained diffusion models on their target data but evaluate these models on broader benchmarks to demonstrate generalizability. In this study, however, the authors only fine-tune their model on CC3M and MS-COCO and limit their evaluation to the corresponding validation sets. Expanding the evaluation to a common benchmark would better showcase the model’s generalization capabilities. Specifically, demonstrating that the prompt router can handle prompts outside the training distribution would be more convincing.

2. The paper also references other model pruning methods, such as MobileDiffusion[1], SnapFusion[2], and LD-Pruner[3]. However, it does not include quantitative comparisons with these approaches. It would be helpful for the authors to explain why these comparisons were omitted.

3. In efficient inference for stable diffusion, recent papers show that one-step or few-step generation can speed up the generation. This paper does not include comparisons with methods like INSTAFLOW[4], which would have provided valuable insights into how APTP compares with state-of-the-art approaches in rapid generation.



[1] Zhao, Yang, et al. "Mobilediffusion: Subsecond text-to-image generation on mobile devices." arXiv preprint arXiv:2311.16567 (2023).

[2] Li, Yanyu, et al. "Snapfusion: Text-to-image diffusion model on mobile devices within two seconds." Advances in Neural Information Processing Systems 36 (2024).

[3] Castells, Thibault, et al. "LD-Pruner: Efficient Pruning of Latent Diffusion Models using Task-Agnostic Insights." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[4] Liu, Xingchao, et al. "Instaflow: One step is enough for high-quality diffusion-based text-to-image generation." The Twelfth International Conference on Learning Representations. 2023.

questions:
1. The authors utilize a pre-trained sentence transformer as the prompt encoder in their training process. Do the authors have any insights into how the size of the prompt encoder influences the overall performance, as the size of the prompt encoder will affect the models' ability to understand the input prompt? 

2. Training diffusion models often incorporate classifier-free guidance. Is the proposed method compatible with training under this manner?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732215193523

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response to All Reviewers (1/3)

comment:
We thank the reviewers for their efforts and valuable feedback. We address the common concerns among them here. 

## 1. Comparison with related works [1, 2, 3, 4, 5, 6]) (R-Ff7o, R-uiaa, R-vGxh): 
We describe that the scope of our method is different from the works that the reviewers mentioned in the following:

### **1-1. Comparison with MobileDiffusion [1] (Architecture Design) and  SnapFusion[2] (Architecture Search) methods (R-Ff7o) :**
We did not compare with MobileDiffusion and SnapFusion as the scope as well as the experimental setup in our paper is significantly different from these methods.

As we describe in the Related Work section, architecture design methods like MobileDiffusion [1] develop heuristics to design an efficient architecture guided by some proxy metrics like FID and CLIP scores on MSCOCO. Similarly, SnapFusion [2] as an architecture search approach trains an architecture with elastic dimensions. Then, it searches for a performant sub-network of it while optimizing for CLIP score on MSCOCO and latency. Yet, these methods have two key drawbacks: 
- They require extremely large training data size and compute budget to train their models and validate their design choices. For instance, MobileDiffusion uses a training dataset of 150 million text-image pairs from the public web and consumes approximately 512 TPUs spanning 15 days to complete the network search. In addition, SnapFusion uses internal proprietary data of unknown size to train its model with 256 NVIDIA A100 GPUs. 
- Their heuristics and design choices are either non-trivial (MobileDiffusion) or costly to generalize (SnapFusion) to new compute budgets and datasets. 

In contrast to these methods, APTP addresses the inference efficiency of a computationally intensive model *after* its pretraining phase by pruning it in a prompt-based manner while optimizing its performance on a *target* data distribution. We prune Stable Diffusion V2.1 using MSCOCO and CC3M, and we show the advantage of APTP compared to recent static pruning baselines for diffusion models. 

In summary, our work and these methods differ in the following aspects:

   1. **Required Training Resources:** MobileDiffusion and SnapFusion train models from scratch, and their training process consumes tens of thousands of GPU/TPU hours. For instance, MobileDiffusion uses 15 Days $\times$ 24 hours $\times$ 512 TPUs = 184320 TPU hours. In contrast, our pruning phase requires approximately 40 GPU hours, and its fine-tuning takes around 120 GPU hours. 
    
   2. **Training dataset size:** While MobileDiffusion and SnapFusion train on tens of millions of data points, we demonstrate that our method works effectively with merely 80k images in the MSCOCO training set. 
    
   3. **Difference in Objectives:** Most importantly, our work aims to show that not all prompts necessitate the full capacity of a pre-trained diffusion model and that one-arch-for-all static pruning architecture approaches are suboptimal for T2I models. We propose a fast and efficient framework to prune an existing off-the-shelf pre-trained T2I diffusion model into a Mixture of Experts, enabling dynamic budget adjustment based on the prompt's complexity. Unlike designing a new model like MobileDiffusion or searching for it from scratch like SnapFusion, our framework is applicable to *any* pretrained diffusion model. 

### **1-2. Comparison with Few/One-Step Generation using distillation [3,4,5] and Caching [6] methods(R-Ff7o, R-uiaa, R-vGxh)**
We discussed in the introduction (L47-51) and related work (L925-930 in the revised pdf) sections that step-reduction methods, like few-step generation ones [3, 4, 5] and caching [6], are **orthogonal** to architectural efficiency methods like ours. In more details, these methods address the "high number of sampling steps" aspect of the broader challenge of the "slow sampling process of diffusion models." They are complementary rather than mutually exclusive with "architectural efficiency of diffusion models." In fact, one can combine step-reduction methods with our approach to further improve inference efficiency, highlighting that our method is not an alternative but a complementary solution to these techniques. Accordingly, the scope of our paper is the architectural efficiency of diffusion models, and we benchmarked APTP against the recent pruning techniques to validate the effectiveness of APTP.

We emphasize that the latency reduction of our method is a result of **model size and memory usage reduction**, which is the **goal of architectural efficiency approaches**. The benefit of architectural efficiency methods is not limited to improved latency. **For instance, one cannot deploy a diffusion model that requires 40GB GPU memory on a GPU with 24GB of memory, regardless of the number of sampling steps and sampling speed-up techniques like [3, 4, 5, 6] that they employ.** In contrast, they can prune it using APTP to smaller experts and deploy an expert on it.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In contrast to the traditional approach where only one pruned T2I model is called by all prompts, this paper proposes to apply different prompts to different pruned models, determined by a router model. This router model is learned by contrastive learning and optimal transport.

The evaluation of this approach is limited to relatively small datasets, and it is still unknown whether the router can be generalized to outside of training distribution.

Overall, it's an interesting idea with some good results. I agree with one reviewer that "The core idea is an interesting combination of known concepts in a new, but highly relevant setup."