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

## Abstract

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. Code is available here .

## 1 Introduction

In recent years, diffusion models [, , , ] have revolutionized generative modeling. For instance, modern Text-to-Image (T2I) diffusion models [, , ] like Stable Diffusion [, ] have achieved remarkable success in generating realistic images [, , ] and editing them [, , ]. Consequently, their integration into various applications is of great interest. However, the sampling process of T2I diffusion models is slow and computationally intensive, making them expensive to deploy on GPU clouds for a large number of users and preventing their utilization on edge devices. Thus, reducing the computational cost of T2I models is essential prior to serving them.

Two orthogonal factors underlie the sampling cost of T2I diffusion models: their large number of denoising steps and their parameter-heavy backbone architectures. Most acceleration methods for T2I models reduce the computation per sampling step or skip steps using techniques like distillation [, , ] and improved noise schedules [, ]. Other ideas address the second factor and propose efficient architectures for T2I models. Architecture modification methods [, ] modify the U-Net [] of Stable Diffusion (SD) [] to reduce its parameters and FLOPs while preserving performance. Still, they only provide a few architectural configurations, and generalizing their design choices to new compute budgets or datasets is highly non-trivial. Search-based methods [, ] search for an efficient architecture [] or an efficient mixture [] of pretrained T2I models in a model zoo. Yet, evaluating each action in the search process and gathering a model zoo of T2I models are both extremely costly, making search methods impractical in resource-constrained scenarios.

Figure: Figure 1: Overview: We prune a text-to-image diffusion model like Stable Diffusion (left) into a mixture of efficient experts (right) in a prompt-based manner. Our prompt router *routes* distinct types of prompts to different experts, allowing experts’ architectures to be separately specialized by removing layers or channels.
Refer to caption: x1.png

The efficient architecture design methods [, , ] have shown promise, but they aim to develop ‘one-size-fits-all’ architectures for *all* applications. We argue that this approach is misaligned with the common practice for two reasons: 1) Organizations typically fine-tune pretrained T2I models (*e.g.,* SD) on their proprietary *target* data and deploy the resulting model for their specific application. They prioritize performance on their *target* data distribution while meeting their computation budget, and the trade-off between efficiency and performance can vary depending on the complexity of the *target* dataset between organizations. 2) Verifying the validity of each design choice on large-scale datasets used in one-size-fits-all methods is costly and slow to iterate, making them impractical.

Model pruning [] can reduce a model’s computational burden to any desired budget with significantly less effort than designing [, ] or searching [] for efficient architectures. Still, T2I models have unique characteristics making existing pruning techniques unsuitable for them. Static pruning methods are input-agnostic and use the same pruned model for all inputs, but distinct prompts of T2I models may require different model capacities. Dynamic pruning employs a separate model for each input sample, but it cannot benefit from batch-parallelism in modern hardware like GPUs and TPUs.

In this paper, we introduce Adaptive Prompt-Tailored Pruning (APTP), a novel prompt-based pruning method for T2I diffusion models. APTP prunes a T2I model pretrained on a large-scale dataset (*e.g.,* SD) using a smaller *target* dataset given a desired compute budget.It tackles the challenges of static and dynamic pruning methods for T2I models by training a *prompt router* module along with a set of *architecture codes*. The prompt router learns to route an input prompt to an architecture code, determining the sub-architecture, called *expert*, of the T2I model to use. Each expert specializes in generating images for the prompts assigned to it by the prompt router (Fig. 1), and the number of experts is a hyperparameter. We train the prompt router and architecture codes using a contrastive learning objective that regularizes the prompt router to select similar architecture codes for similar prompts. In addition, we employ optimal transport to diversify the architecture codes and the resulting experts’ budgets, allowing deploying them on hardware with varying capabilities. We take CC3M [] and MS-COCO [] as the *target* datasets and prune Stable Diffusion V2.1 [] using APTP in our experiments. APTP outperforms the single-model pruning baseline, and we show that our prompt router learns to group the input prompts into semantic clusters. Further, our analysis demonstrates that APTP can automatically discover challenging prompts for SD, such as prompts for generating text images, found empirically by prior work [, ]. We summarize our contributions
as follows:

- •
We introduce APTP, a novel prompt-based pruning method for T2I diffusion models. It is more suitable than static pruning for T2I models as it is not input-agnostic. In addition, APTP enables batch parallelism on GPUs, which is not possible with dynamic pruning.
- •
APTP trains a prompt router and a set of architecture codes. The prompt router maps an input prompt to an architecture code. Each architecture code resembles a pruned sub-architecture expert of the T2I model, specialized in handling certain types of prompts assigned to it.
- •
We develop a framework to train the prompt router and architecture codes using contrastive learning and employ optimal transport to diversify the architecture codes and their corresponding sub-architectures’ computational requirements, given a desired compute budget.

## 2 Related Work

Several works have addressed improving the architectural efficiency of diffusion models, which is our paper’s primary focus.
Multi-expert [, , , ] methods employ several models each responsible for an interval of diffusion model’s denoising process. These methods design expert model architectures [, ], train several models with varying capacities from scratch [], or utilize existing pretrained models [, ]. However, pretrained experts may not be available, and training several models from scratch, as done by multi-expert methods, is prohibitively expensive.
Architecture Design [, , ] approaches redesign the U-Net architecture [] of diffusion models to enhance its efficiency. MobileDiffusion [] does so using empirical heuristics derived from performance metrics on MS-COCO []. BK-SDM [] removes some blocks from the SD’s U-Net and applies distillation [] to train the pruned model. Spectral Diffusion [] introduces a wavelet gating operation and performs frequency domain distillation. Yet, generalizing heuristics and design choices of the architecture design methods to other tasks and compute budgets is non-trivial.
Alternatively, SnapFusion [] searches for an efficient architecture for T2I models. Yet, evaluating each action in the search process requires 2.5 A100 GPU hours, making it costly in practice. Finally, SPDM [] estimates the importance of different weights using Taylor expansion and removes the low-scored ones. Despite their promising results, all these methods are ‘static’ in that they obtain an efficient model and utilize it for all inputs. This is sub-optimal for T2I models as input prompts may vary in complexity, demanding different model capacity levels.
Our method differs from existing approaches by introducing a prompt-based pruning technique for T2I models. It is the first method that allocates computational resources to prompts based on their individual complexities while ensuring optimal utilization and batch-parallelizability. We refer to Appendix A for a detailed review of related work.

## 3 Method

We introduce a framework for prompt-based pruning of T2I diffusion models termed Adaptive Prompt-Tailored Pruning (ATPT). APTP prunes a T2I model pretrained on large-scale datasets (*e.g.*, Stable Diffusion []) using a smaller *target* dataset. This approach mirrors the common practice where organizations fine-tune pretrained T2I models on their internal proprietary data and deploy the resulting model for their customers. The core component of APTP is a *prompt router* that learns to map an input prompt to an *architecture code* during the pruning process. Each architecture code corresponds to a specialized *expert* model, which is a sub-network of the T2I model. We train the prompt router and architecture codes in an end-to-end manner. After the pruning phase, APTP fine-tunes each specialized model using samples from the target dataset assigned to it by the prompt router. We elaborate on the components of APTP in the following subsections.

### 3.1 Background

Given an input text prompt, text-to-image (T2I) diffusion models generate a corresponding image by iteratively denoising a Gaussian noise [, , ]. They achieve this by training a denoising model, $\epsilon(\cdot;\theta)$, parameterized by $\theta$. For a given training image-text pair $(x_{0},p)\sim\mathcal{P}$, T2I models define a forward diffusion process, progressively adding Gaussian noise to the initial image $x_{0}$ over $T$ steps. This process is defined as $q(x_{t}|x_{0})=\mathcal{N}(x_{t};\sqrt{\bar{\alpha_{t}}}x_{0},(1-\bar{\alpha_{
t}})I)$, where $\bar{\alpha_{t}}$ is the forward noise schedule parameter, typically chosen such that $q(x_{t}|x_{0})\rightarrow\mathcal{N}(0,I)$ as $t\rightarrow T$. T2I models are trained using the variational evidence lower bound (ELBO) objective [].

$$ $\mathcal{L}_{\text{DDPM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}(x_{0},p)\sim \mathcal{P}\\ t\sim[1,T]\\ \epsilon\sim\mathcal{N}(0,I)\\ x_{t}\sim q(x_{t}|x_{0})\end{subarray}}||\epsilon(x_{t},p,t;\theta)-\epsilon|| ^{2}$ (1) $$

T2I models sample an image starting from a Gaussian noise $x_{T}$ $\sim$ $\mathcal{N}(0,I)$ and denoising it using the trained denoising model. We refer to appendix C for a thorough review of diffusion models.

Figure: Figure 2: Our pruning scheme. We train our prompt router and the set of architecture codes to prune a text-to-image diffusion model into a mixture of experts. The prompt router consists of three modules. We use a Sentence Transformer [] as our prompt encoder to encode the input prompt into a representation $z$. Then, the architecture predictor transforms $z$ into the architecture embedding $e$ that has the same dimensionality as architecture codes. Finally, the router routes the embedding $e$ into an architecture code $a^{(i)}.$ We use optimal transport to evenly assign the prompts in a training batch to the architecture codes. The architecture code $a^{(i)}=(u^{(i)},v^{(i)})$ determines pruning the model’s width and depth. We train the prompt router’s parameters and architecture codes in an end-to-end manner using the denoising objective of the pruned model $\mathcal{L}_{\text{DDPM}}$, distillation loss between the pruned and original models $\mathcal{L}_{\text{distill}}$, average resource usage for the samples in the batch $\mathcal{R}$, and contrastive objective $\mathcal{L}_{\text{cont}}$, encouraging embeddings $e$ preserving semantic similarity of the representations $z$.
Refer to caption: x2.png

### 3.2 Prompt Router and Architecture Codes

We employ our prompt router to determine the specialized sub-network of the T2I model to be used for a given input prompt. Specifically, we denote our prompt router with the function $f_{\text{PR}}(\cdot;\eta,\mathcal{A})$, parameterized by $\eta$, which maps an input prompt $p$ to an architecture code $a$:

$$ $a=f_{\text{PR}}(p;\eta,\mathcal{A})$ (2) $$

We define the set of learnable architecture codes as $\mathcal{A}=\{a^{(i)}\}_{i=1}^{N}$, where $N$ is a hyperparameter. Also, $a^{(i)}\in\mathbb{R}^{D}$, with $D$ being the number of prunable width and depth units in the T2I model. Given a desired constraint $T_{d}$ on the total compute budget (latency, MACs, etc.) for the prompts in the target dataset, we train the prompt router’s parameters $\eta$ and architecture codes in $\mathcal{A}$ to obtain a set of expert models. These experts are specialized, efficient, and performant sub-networks of the T2I model, as determined by the architecture codes. As illustrated in Fig. 2, the prompt router consists of a prompt encoder, an architecture predictor, and a router module. We provide details of these components in the following subsections and present our pruning objective in Eq. 15.

#### 3.2.1 Prompt Encoder and Architecture Predictor

Our primary intuition in designing the prompt router is that it should route semantically similar prompts to similar sub-networks of a T2I model. Accordingly, we use a pretrained frozen Sentence Transformer model [] as our prompt encoder module. It can effectively encode input prompts $p$ into semantically meaningful embeddings $z$:

$$ $z=f_{\text{PE}}(p)$ (3) $$

$f_{\text{PE}}(\cdot)$ is the prompt encoder. We do not explicitly show the prompt encoder’s parameters as we do not train them in our pruning method. The Sentence Transformer can encode semantically similar prompts to nearby embeddings and distant from dissimilar ones. We leverage this property in our framework (Eq. 13) to ensure the prompt router maps similar prompts to similar architecture codes.

The architecture predictor module transforms prompt embeddings $z$ into architecture embeddings $e$:

$$ $e=f_{\text{AP}}(z;\eta)$ (4) $$

$\eta$ denotes the parameters of the architecture predictor $f_{\text{AP}}(\cdot)$, implemented with a single feed-forward layer (Appendix D.1). The embeddings $e$ have the same dimensionality as the architecture codes $a$.

#### 3.2.2 Router

The router module takes an architecture embedding $e$ and routes it to an architecture code $a\in\mathcal{A}$. A
straightforward way to implement the router is to map an input embedding $e$ to its nearest architecture code in $\mathcal{A}$. However, we found that this approach may lead to the collapse of architecture codes, as they could converge to a single code that meets the desired compute budget $T_{d}$ with a relatively decent performance, causing the prompt router to route all input prompts to it.

To tackle this challenge, we employ optimal transport in our router module during the pruning phase. Formally, let $E=[e_{1},\cdots,e_{B}]$ represent $B$ architecture embeddings in a training batch, and $A=[a^{(1)},\cdots,a^{(N)}]$ represent the $N$ codes. The goal is to find an assignment matrix $Q=[q_{1},\cdots,q_{B}]$ that maximizes the similarity between architecture embeddings and their assigned architecture codes:

$$ $~{}\max_{Q\in\mathcal{Q}}\text{Tr}(Q^{T}A^{T}E)+\epsilon\text{H}(Q)$ (5) $$

H is an entropy term $\text{H}(Q)=-\sum_{i,j}q_{ij}\log(q_{ij})$ and $\epsilon$ is the regularization strength. It has been shown [, ] that high values of $\epsilon$ lead to a uniform assignment matrix $Q$, causing all codes in $A$ to collapse to a single code. Thus, we set $\epsilon$ to a small value. Further, we impose an equipartition [] constraint on $Q$ so that architecture embeddings in a batch are assigned equally to architecture codes:

$$ $\mathcal{Q}=\{Q\in\mathbb{R}_{+}^{N\times B}~{}|~{}Q1_{B}=\frac{1}{N}1_{N},~{} Q^{T}1_{N}=\frac{1}{B}1_{B}\}$ (6) $$

Here, $1_{\{B,N\}}$ are vectors of ones with length $B$ and $N$, respectively. These constraints enforce that, on average, $\frac{B}{N}$ embeddings are assigned to an architecture code per training batch, thereby ensuring that each architecture code gets enough samples for training. The optimal transport problem in Eq. 5 with the constraints in Eq. 6 can be solved using the fast version [] of the Sinkhorn-Knopp algorithm and the solution has the normalized exponential matrix form []:

$$ $Q^{*}=\text{diag}(m)\text{exp}(\frac{A^{T}E}{\epsilon})\text{diag}(n)$ (7) $$

Here, $m$ and $n$ are renormalization vectors that can be calculated using a few iterations of the Sinkhorn-Knopp algorithm. With $\epsilon$ set to a small value, the matrix $BQ^{*}$ will have columns that are close to one-hot vectors, which we use to assign the architecture embeddings to the architecture codes. In summary, the router module’s function during the pruning process is:

$$ $a=f_{\text{R}}(e,\mathcal{A};Q^{*})$ (8) $$

After the pruning stage, the trained router simply routes an input architecture embedding to an architecture code with the highest cosine similarity. We provide the definition of $f_{\text{R}}$ in Appendix D.2.

### 3.3 Pruning

We divide each architecture code $a^{(i)}\in\mathcal{A}$ into two sub-vectors $a^{(i)}=(u^{(i)},v^{(i)})$. We utilize the vectors $u^{(i)}$ and $v^{(i)}$ to prune the depth layers and determine widths of the layers, respectively.

A simple approach to prune the width of a layer (a depth layer) is to use binary vectors $v^{(i)}$ ($u^{(i)}$), indicating whether the channels (layers) should be pruned. Yet, doing so is not differentiable, and one needs to solve a discrete optimization problem to find the optimal binary vectors. Instead, we employ soft vectors $\textbf{v}^{(i)}$ ($\textbf{u}^{(i)}$) that are continuous and differentiable for pruning. We calculate them as:

$$ $\textbf{v}^{(i)}=\text{sigmoid}(\frac{v^{(i)}+g_{v}}{\gamma}),~{}~{}~{}~{}~{}~ {}\textbf{u}^{(i)}=\text{sigmoid}(\frac{u^{(i)}+g_{u}}{\gamma})$ (9) $$

$g_{\{u,v\}}\sim\text{Gumbel}(0,1)$ represents a noise vector sampled from the Gumbel distribution [], and $\gamma$ is the temperature. This formulation, known as the Gumbel-sigmoid reparameterization [, ], is a differentiable approximation of sampling from Bernoulli distributions with parameters $\texttt{sigmoid(}v^{(i)}\texttt{)}$ and $\texttt{sigmoid(}u^{(i)}\texttt{)}$. When the temperature $\gamma$ is set appropriately, the vectors $\textbf{v}^{(i)}$ and $\textbf{u}^{(i)}$ will be close to binary vectors, and we use them for pruningthe layers’ width and depth layers. Assuming $\textbf{v}^{(i)}=[\textbf{v}_{l}^{(i)}]_{l=1}^{L}$ where $L$ is the number of the model’s layers, we prune the width of the $l$-th layer as:

$$ $\mathcal{\widehat{F}}_{l}=\mathcal{F}_{l}\odot\textbf{v}_{l}^{(i)}$ (10) $$

In this equation, $\mathcal{F}_{l}$ represents feature maps of the $l$-th layer, and $\odot$ is the element-wise multiplication. We prune the channels of the convolution layers in the ResBlocks [], attention heads in the Transformer layers [], and the channels of the feed-forward layers in the Transformer layers.

In a similar manner, we apply the vectors $\textbf{u}^{(i)}=[\textbf{u}_{j}^{(i)}]_{j=1}^{M}$ (where M is the number of depth layers that we prune) for pruning the model’s depth. Specifically, we prune the $j$-*th* depth layer $f_{j}$ as:

$$ $\mathcal{\widehat{F}}_{j}=\mathbf{u}_{j}^{(i)}f_{j}(\mathcal{F}_{j-1})+(1- \mathbf{u}_{j}^{(i)})\mathcal{F}_{j-1}$ (11) $$

$\mathcal{F}_{j-1}$ denotes the previous layer’s feature maps. The granularity of our depth pruning is a ResBlock or a Transformer layer in the U-Net of the T2I model. We provide more details in Appendix D.3. In summary, we employ the architecture vectors $\textbf{a}^{(i)}=(\textbf{v}^{(i)},\textbf{u}^{(i)})$ to prune the T2I model.

#### 3.3.1 Training the Prompt Router and Architecture Codes

We jointly train the prompt router and architecture codes in an end-to-end manner, guiding the prompt router to map similar prompts to similar architecture codes. In addition, we regularize the architecture codes to correspond to performant sub-networks of the T2I model, be diverse, and adhere to the desired compute budget $T_{d}$ on aggregate.

##### Contrastive Training

Given a training batch with $B$ prompts, we compute their prompt embeddings $z$ (Eq. 3) and architecture embeddings $e$ (Eq. 4). Then, we calculate architecture vectors $\textbf{e}^{\prime}$:

$$ $\textbf{e}^{\prime}=\text{sigmoid}(\frac{e+g}{\gamma})$ (12) $$

where $\gamma$ and $g$ have the same definitions as Eq. 9. We define the cosine similarity of two vectors as
$\text{sim}(\textbf{m},\textbf{n})=\textbf{m}^{T}\textbf{n}/||\textbf{m}||.||
\textbf{n}||$, and we use the following objective to train the architecture predictor:

$$ $\mathcal{L}_{\text{cont}}(\eta)=\frac{1}{B^{2}}\sum_{i=1}^{B}\sum_{j=1}^{B}r_{ i,j}\text{log}(s_{i,j})+(1-r_{i,j})\text{log}(1-s_{i,j})$ (13) $$

$$ $r_{i,j}=\frac{\text{exp}(\text{sim}(z_{i},z_{j})/\tau)}{\sum_{k=1}^{B}\text{ exp}(\text{sim}(z_{i},z_{k})/\tau)},~{}~{}~{}~{}s_{i,j}=\frac{\text{exp}(\text {sim}(\textbf{e}_{i}^{\prime},\textbf{e}_{j}^{\prime})/\tau)}{\sum_{k=1}^{B} \text{exp}(\text{sim}(\textbf{e}_{i}^{\prime},\textbf{e}_{k}^{\prime})/\tau)}$ (14) $$

$\tau$ is a temperature parameter. Eq. 13 regularizes the architecture predictor to map representations $z$ to the regions of the space of the architecture embeddings $e$ such that their corresponding architecture vectors $\textbf{e}^{\prime}$ maintain the similarity between the prompts. Note that we do not actually use the architecture vectors $\textbf{e}^{\prime}$ to prune the
model (Sec. 3.2.2, 3.3). Yet, we apply our contrastive regularization to the vectors $\textbf{e}^{\prime}$ instead of embeddings $e$. This design choice is a result of our observation that applying $\mathcal{L}_{\text{cont}}$ to embeddings $e$ may not lead to diverse architecture vectors $\textbf{a}^{(i)}=(\textbf{v}^{(i)},\textbf{u}^{(i)})$ that we use for
pruning (Sec. 3.3). For instance, embeddings $e$ (and their nearby architecture codes $a$) can be distributed in the embedding space (before Gumbel-Sigmoid estimation), but all be in saturation regions of the sigmoid function (Eqs. 9, 12), resulting in similar architecture vectors $\textbf{e}^{\prime}$ and $\textbf{a}^{(i)}$.

In contrast, Eq. 13 implicitly diversifies the architecture vectors $\textbf{a}^{(i)}$. The reason is that the router routes embeddings $e$ to codes $a$ (Eq. 8) that have high similarity with each other (Eq. 7). Also, Eq. 13 distributes embeddings $e$ in the space of the architecture embeddings such that the architecture vectors $\textbf{e}^{\prime}$ become similar (different) for similar (dissimilar) prompts. As the vectors $\textbf{e}^{\prime}$ and $\textbf{a}^{(i)}$ are calculated in a similar manner (Eqs. 9, 12), the diversity of vectors $\textbf{e}^{\prime}$ implies the same for vectors $\textbf{a}^{(i)}$.

Assuming $B^{(i)}$ samples get routed to the architecture code $a^{(i)}$ in a training batch ($\sum_{i}B^{(i)}=B$), we train the prompt router and the architecture codes using the following objective:

$$ $\displaystyle\min_{\eta,\mathcal{A}}\mathcal{L}=$ $\displaystyle[\frac{1}{N}\sum_{i=1}^{N}[\frac{1}{B^{(i)}}\sum_{j=1}^{B^{(i)}}[ \mathcal{L}_{\text{DDPM}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})+\lambda_{\text{ distill}}\mathcal{L}_{\text{distill}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})]]]$ (15) $\displaystyle+\lambda_{\text{res}}\mathcal{R}(\widehat{T}(\mathcal{A}),T_{d})+ \lambda_{\text{cont}}\mathcal{L}_{\text{cont}}(\eta)$ $$

$\mathcal{L}_{\text{DDPM}}((x^{(i)}_{j},p^{(i)}_{j});a^{(i)})$ denotes the denoising objective for the sample $(x^{(i)}_{j},p^{(i)}_{j})$ routed to sub-network chosen by the architecture code $a^{(i)}$. $\mathcal{R}(\widehat{T}(\mathcal{A}),T_{d})$ regularizes the weighted average of the MACs used by architecture codes ($\widehat{T}(\mathcal{A})=\sum_{i}\frac{B^{(i)}}{B}[\widehat{T}(a^{(i)})]$) to be close to $T_{d}$. We define $\mathcal{R}(x,y)=\text{log}(\text{max}(x,y)/\text{min}(x,y))$, and
$\{\lambda_{\text{distill}},\lambda_{\text{res}},\lambda_{\text{cont}}\}$
are hyperparameters. Finally, $\mathcal{L}_{\text{distill}}$ is the distillation objective [] regularizing the pruned model having similar outputs to the original one. We refer to Appendix D.4 for more details about our distillation objective.

### 3.4 Fine-tuning the Pruned Expert Models

After the pruning stage, we use the learned architecture codes to prune the T2I model into our experts. Then, we fine-tune experts using samples routed to them. We use the same training objective as the pretraining stage while adding distillation to fine-tune experts and refer to Appendix D.5 for more details. At test time, our trained prompt router routes an input prompt to one of the experts, and we use the expert to
generate an image for it.

Figure: Figure 3: Samples of the APTP-Base experts after pruning the Stable Diffusion V2.1 using CC3M [] and COCO [] as the target datasets. Expert IDs are shown on the top right of images. (See Table 6 for prompts)
Refer to caption: extracted/5669950/figs/grids/grid.png

## 4 Experiments

We use Conceptual Captions 3M (CC3M) [] and MS-COCO [] as our target datasets to prune the Stable Diffusion (SD) V2.1 model to demonstrate APTP’s effectiveness. On CC3M, we prune the model with Base: ($0.85$ MACs, $16$ experts) and Small: ($0.66$ MACs, $8$ experts) configurations. On MS-COCO, we perform Base: ($0.78$ MACs, $8$ experts) and Small: ($0.64$ MACs, $8$ experts) pruning settings. We set $(\lambda_{\text{distill}},\lambda_{\text{res}},\lambda_{\text{cont}})$=$(0.2,2.0,100)$ (Eq. 15) and the temperature $\tau$ (Eq. 14) to $0.03$. We evaluate all models with FID [], CLIP [], and CMMD [] scores using 14k samples in the validation set of CC3M and 30k samples from the MS-COCO’s validation split. For all models, we sample the images at the resolution of $768$ and resize them to $256$ for calculating the metrics, using the 25-steps PNDM [] sampler following BK-SDM []. We refer to appendix E for more details.

### 4.1 Comparison Results

As APTP is the first pruning method to prune a pretrained T2I model on a target dataset, we compare its performance with SD V2.1 and the weight norm pruning [] in this section, and Table 1(b) shows the results. We found that the APTP and Norm pruning baseline reach their final performance after about 30k and 50k iterations fine-tuning after pruning.

##### CC3M:

Table 1(a) summarizes the results on CC3M. With a similar MACs budget, APTP ($0.85$) significantly outperforms the Norm pruning [] baseline in terms of FID, CLIP, and CMMD scores. It also achieves $15\%$ less latency while showing close performance scores to SD V2.1. Notably, APTP (0.66) has $23\%$ less latency and $21\%$ lower MACs budget than Norm [], and at the same time, shows a large gap with the Norm pruning.

##### MS-COCO:

We present the results for MS-COCO in Table 1(b). APTP (0.78) reduces the latency of SD by $22.5\%$ while preserving its CLIP score and achieving a close CMMD score. Further, APTP (0.78) significantly outperforms the model pruned by Norm [] with a similar latency. Similar to the results on CC3M, APTP (0.64) achieves $19.3\%$ less latency than the Norm pruning while outperforming it on FID, CLIP, and CMMD scores. In summary, our quantitative evaluations show the clear advantage of prompt-based compared to static pruning for T2I models. We provide samples of APTP on the validation sets of CC3M and MS-COCO in Fig. 3 (See Appendix. E.4.6 for more samples).

**Table 1: Results on CC3M and MS-COCO. We report performance metrics using samples generated at the resolution of $768$ then downsampled to $256$ []. We measure models’ MACs/Latency with the input resolution of $768$ on an A100 GPU. @30/50k shows fine-tuning iterations after pruning.**
| CC3M |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Method | Complexity | Performance |  |  |  |
| MACs<br>(@768) | Latency ($\downarrow$)<br>(Sec/Sample)<br>(@768) | FID ($\downarrow$) | CLIP ($\uparrow$) | CMMD ($\downarrow$) |  |
| Norm []<br>@50k | 1185.3G | 3.4 | 141.04 | 26.51 | 1.646 |
| APTP(0.66)<br>@30k | 916.3G | 2.6 | 60.04 | 28.64 | 1.094 |
| APTP(0.85)<br>@30k | 1182.8G | 3.4 | 36.77 | 30.84 | 0.675 |
| SD 2.1 | 1384.2G | 4.0 | 32.08 | 31.12 | 0.567 |

**Table 2: The most frequent words in prompts assigned to each expert of APTP-Base pruned on CC3M. The resource utilization of each expert is indicated in parentheses.**
| Expert 1 (0.72) | Expert 2 (0.73) | Expert 3 (0.75) | Expert 4 (0.76) |
| --- | --- | --- | --- |
| View - Sunset - City - Building - Sky | View - Boat - Sea | Artist - Actor | Actor - Dress - Portrait |
| Expert 5 (0.77) | Expert 6 (0.78) | Expert 7 (0.79) | Expert 8 (0.79) |
| Illustration - Portrait - Photo | Player - Ball - Game - Team | Background - Water - River - Tree | Biological Species - Dog - Cat |
| Expert 9 (0.79) | Expert 10 (0.80) | Expert 11 (0.81) | Expert 12 (0.81) |
| Illustration - Vector | People | Car - City - Road | Person - Player - Team - Couple |
| Expert 13 (0.86) | Expert 14 (0.90) | Expert 15 (0.95) | Expert 16 (0.98) |
| Room - House | Art - Artist - Digital | Food - Water | Person - Man - Woman - Text |

### 4.2 Analysis of the Prompt Router

We analyze our prompt router’s behavior by examining the prompts it assigns to experts in the APTP-Base ($0.85$) model for the CC3M experiment. Table 2 displays the most frequent words in the prompts routed to each expert, along with the experts’ MACs budgets. Our prompt router effectively "specializes" experts by assigning distinct topics to experts with varying budgets. For instance, Expert 1 focuses on cityscapes, Expert 8 on animals, and Expert 13 on house interiors and exteriors. Notably, the prompt router assigns images of textual content and human beings to Expert 16, which has the highest budget. These categories have been empirically found to be challenging for SD 2.1 [, ], and our prompt router can automatically discover and route them to higher-capacity experts. In contrast, paintings and illustrations are assigned to lower-budget experts, as they seem to be easier for the model to generate. We provide examples of high and low resource prompts for APTP-Base on both CC3M and COCO in Fig. 4.

Figure: Figure 4: Comparison of samples generated by low and high budget experts of APTP-Base vs. SD V2.1 on CC3M and MS-COCO validation sets.
Refer to caption: x3.png

### 4.3 Ablation Study

We conduct two ablation experiments to study the impact of APTP’s components on its performance.
First, we implement a naive baseline that uses parameterizations $\textbf{v}=\text{sigmoid}((\eta_{v}+g_{v})/\tau)$ and $\textbf{u}=\text{sigmoid}((\eta_{u}+g_{u})/\tau)$ (Eq. 9) to prune a single model. It directly trains two vectors $(\eta_{v},\eta_{u})$ to do so (Uni-Arch Baseline’ in Table 3). Then, we start with a router trained only using the contrastive objective (Eq. 13) and add optimal transport and distillation incrementally to prune SD V2.1 into $8$ experts with $80\%$ MACs budget on MS-COCO []. We fine-tune all pruned models with 10k iterations, and Table 3 presents the results. We observe that the contrastive training (Eq. 13) alone fails to improve results from pruning a single model to a mixture of experts. This is because although the contrastive objective makes the architecture codes diverse, it does not enforce the prompt router to distribute the prompts between architecture codes. Thus, it routes most of the input prompts to a single expert. As a result, all experts except one receive insufficient training samples and generate low-quality samples after fine-tuning, leading to the mixture showing worse metrics than the baseline. Employing optimal transport (Eq. 5) in the prompt router significantly improves FID (10.22), CLIP (1.17), and CMMD (0.18) scores of the mixture. In addition, distillation can further improve the architecture search process of the experts, resulting in a more performant mixture. In summary, these results validate the effectiveness of our design choices for APTP.

**Table 3: Ablation results of APTP’s components on 30k samples from MS-COCO [] validation set. We fine-tune all models for 10k iterations after pruning.**
| Method | MACs(@768) | Latency(@768) | FID ($\downarrow$) | Clip Score ($\uparrow$) | CMMD ($\downarrow$) |
| --- | --- | --- | --- | --- | --- |
| Uni-Arch Baseline | 1088.8G | 3.1 | 46.56 | 29.11 | 0.91 |
| Contrastive Router | 1079.5G | 3.1 | 48.78 | 28.90 | 0.92 |
| + Optimal Transport | 1076.6G | 3.1 | 38.56 | 30.07 | 0.74 |
| + Distillation (APTP) | 1076.6G | 3.1 | 25.57 | 31.13 | 0.58 |

Figure: Figure 5: Ablation Results for the number of experts of APTP on MS-COCO.
Refer to caption: x4.png

In our second ablation experiment, we explore the impact of the number of experts on APTP. We prune SD V2.1 to $80\%$ MACs budget using APTP with $4$, $8$, and $12$ experts on MS-COCO. Fig. 5 shows the results. Interestingly, the results demonstrate that the relationship between FID and CLIP scores with the number of experts is nonlinear, and the optimal number of experts is dataset-dependent. This observation illustrates that prompt-based pruning is more suitable than static pruning for T2I models.

## 5 Conclusion

In this paper, we develop Adaptive Prompt-Tailored Pruning (APTP), the first prompt-based pruning method for text-to-image (T2I) diffusion models. APTP takes a T2I model, pretrained on large-scale data, and prunes it using a target dataset, resembling the common practice that organizations fine-tune T2I models on their internal data before deployment. The core element of APTP is a prompt router module that learns to decide the model capacity required to generate a sample for an input prompt, routing it to an architecture code given a desired compute budget. Each architecture code corresponds to a sub-network of the T2I model, specializing in generating images for the prompts that prompt router routes to it. APTP trains the prompt router and architecture codes in an end-to-end manner, encouraging the prompt router to route similar prompts to similar architecture codes. Further, we utilize optimal transport in the prompt router of APTP during pruning to diversify the architecture codes. Our experiments in which we prune Stable Diffusion (SD) V2.1 using CC3M and MS-COCO as target datasets demonstrate the benefit of prompt-based pruning compared to conventional static pruning methods for T2I models. Further, our analysis on APTP’s prompt router reveals that it can automatically discover challenging prompt types for SD, like generating text, humans, and fingers, routing them to experts with high compute budgets.

## Appendix A Related Work

##### Efficient Diffusion Models:

Methods for accelerating and improving the efficiency of diffusion models fall into two categories. The first group focuses on reducing the complexity or the number of sampling steps of diffusion models. They use techniques like distillation [, , , ], learning optimal denoising time-steps [, ],
designing improved noise schedules [, , ],
caching intermediate computations [, , ],
and proposing faster solvers [, , ]. The second group is aimed at improving the architectural efficiency of diffusion models, which is the main focus of our paper. Multi-expert [, , , ] methods employ multiple expert models each responsible for a separate part of the denoising process of diffusion models. These methods either design expert model architectures [, ], train several models with varying capacities from scratch [], or utilize existing pretained experts [, ] with different capacities. However, pretrained experts are not necessarily available, and training several models from scratch, as required by multi-expert methods, is prohibitively expensive in practice. Architecture Design approaches [, , ] redesign the architecture of diffusion models to enhance efficiency. MobileDiffusion [] modifies the U-Net [] architecture of Stable Diffusion (SD) based on empirical heuristics derived from the model’s performance on the MS-COCO [] dataset. BK-SDM [] removes some blocks from the U-Net model of SD and applies knowledge distillation from the original SD model to the pruned model. Spectral Diffusion [] introduces a wavelet gating operation and performs frequency domain distillation from a pretrained teacher model into a small student model. Yet, generalization of the heuristics and design choices in the architecture design methods to other tasks and compute budgets is non-trivial.
Quantization methods [, , , , , , , ] reduce the precision of model weights and activations during the forward pass to accelerate the sampling process.  Different from these methods, SnapFusion [] searches for an efficient architecture for T2I models. However, evaluating each action in the search process consumes about 2.5 A100 GPU hours, rendering it impractical for resource-constrained scenarios. Finally, SPDM [] estimates the importance of different weights in the model using Taylor expansion and removes the low-scored architectures. Despite their promising results, all these methods are ‘static’ in that they obtain an efficient model and utilize it for all inputs. This is suboptimal for T2I models as input prompts may vary in complexity, demanding different model capacity levels.

##### Pruning and Neural Architecture Search (NAS):

Our paper also intersects with model pruning [] and NAS [, , , , , ]
methods. These methods prune pretrained models and search for suitable architectures given a specific task and computational budget.
Existing pruning techniques can be categorized into Static and Dynamic methods. Static pruning approaches [, , ]
prune a pretrained model and use the pruned model for all inputs. Conversely, dynamic pruning methods [, , , ] employ a separate sub-network of the model for each input. Although static and dynamic pruning methods have been successful for image classification, they are not suitable to be directly applied to T2I models. Static pruning methods neglect prompt complexity and employ the same model for all prompts while dynamic pruning ideas cannot utilize batch-parallelism in GPUs. We refer to recent surveys [, , ] for a comprehensive review of pruning and NAS methods.

Our method distinguishes itself from existing approaches by introducing a prompt-based pruning technique for T2I models. This work marks the first instance where computational resources are allocated to prompts based on their individual complexities while ensuring optimal utilization and batch-parallelizability.

## Appendix B Limitations and Broader Impacts

We believe a limitation of our work may be that our experimental scope is relatively limited. We chose CC3M [] and MS-COCO [] as they contain at least 100k image-text pairs and are diverse to enable us to simulate the real-world case of fine-tuning a pretrained T2I model like SD v2.1 on a target dataset before deploying it. Yet, exploring APTP’s performance on other tasks is an interesting future work. In addition, APTP, like Stable Diffusion, has challenges generating images of text, humans, and fingers (Sec. 4.2, Fig. 4), and one can employ techniques in recent works [, , ] to alleviate these issues. Finally, we note that our method can enable organizations to serve their models with a lower cost for GenAI applications, but it may also facilitate harmful content generation, necessitating responsible usage of generative models.

## Appendix C Overview of Diffusion Models

Given a random variable $\mathbf{x_{0}}\sim\mathcal{P}$, the goal of diffusion models [, ] is to model the underlying distribution $\mathcal{P}$ using a training set $\mathcal{D}=\{x_{0}\}$ of samples. To do so, first, diffusion models define a forward process parameterized by $t$ in which they gradually perturb each sample $x_{0}$ with Gaussian noise with the variance schedule of $\beta_{t}$:

$$ $q(x_{t}|x_{t-1})=\mathcal{N}(x_{t};\sqrt{1-\beta_{t}}x_{t-1},\beta_{t}I)$ (16) $$

where $t\in[1,T].$ Thus, $q(x_{t}|x_{0})$ has a Gaussian form:

$$ $q(x_{t}|x_{0})=\mathcal{N}(x_{t};\sqrt{\bar{\alpha_{t}}}x_{0},(1-\bar{\alpha_{ t}})I)$ (17) $$

where $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$. The noise schedule $\beta_{t}$ is usually selected [] such that $q(x_{T})\rightarrow\mathcal{N}(0,I)$. Assuming $\beta_{t}$ is small, diffusion models approximate the denoising distribution $q(x_{t-1}|x_{t})$ by a parameterized Gaussian distribution $p_{\theta}(x_{t-1}|x_{t})=\mathcal{N}(x_{t-1};\frac{1}{\sqrt{\alpha_{t}}}(x_{t
}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha_{t}}}}\epsilon_{\theta}(x_{t},t)),
\sigma^{2}_{t}I)$, and $\sigma^{2}_{t}$ is often set to $\beta_{t}$. Diffusion models implement $\epsilon_{\theta}(.)$ with a neural network called the denoising model and train it with the variational evidence lower bound (ELBO) objective []:

$$ $\mathcal{L}_{\text{DDPM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}t\sim[1,T]\\ \epsilon\sim\mathcal{N}(0,I)\\ x_{t}\sim q(x_{t}|x_{0})\end{subarray}}||\epsilon(x_{t},t;\theta)-\epsilon||^{2}$ (18) $$

Similarly, T2I diffusion models train a denoising model using pairs of image and text prompts $(x_{0},p)\sim\mathcal{P}$ to model the distribution $\mathcal{P}(\textbf{x}_{0}|\textbf{p})$ of images given an input text prompt:

$$ $\mathcal{L}_{\text{DDPM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}(x_{0},p)\sim \mathcal{P}\\ t\sim[1,T]\\ \epsilon\sim\mathcal{N}(0,I)\\ x_{t}\sim q(x_{t}|x_{0})\end{subarray}}||\epsilon(x_{t},p,t;\theta)-\epsilon|| ^{2}$ (19) $$

T2I Diffusion models generate a new sample by sampling an initial noise from $x_{T}\sim p(x_{T})=\mathcal{N}(0,I)$ and iteratively denoising it using the denoising model by sampling from $p_{\theta}(x_{t-1}|x_{t},p)$. Thus, the sampling process requires $T$ sequential forward calculation of the denoising model, making it a slow and costly process.

## Appendix D More Details of APTP

In this section we provide more details of the architecture predictor, the prompt router module, and our pruning method.

### D.1 Architecture Predictor

$f_{\text{AP}}(\cdot)$ in Eq. 4 is the architecture predictor. The architecture predictor changes the
dimensionality of $z$ to the dimensionality of $e$ and codes $a$, which is the number of prunable width and depth units in the T2I model. We implement the $f_{\text{AP}}$ function with a single feed-forward layer. Our prompt encoder, Sentence Transformer [], has an output dimension of $\mathbf{768}$. The total number of prunable units in SD 2.1 is $\mathbf{1620}$, so the architecture predictor has an input dimension of $\mathbf{768}$ and an output dimension of $\mathbf{1620}$ in all of our experiments.

### D.2 Router Module Definition

During the pruning process, we use the assignment matrix $Q^{*}$ calculated by optimal transport (Eq. 7) to route the architecture embeddings $e$ to the architecture codes $a$. Thus, during the pruning process, the definition of function $f_{\text{R}}$ is as follows:

$$ $I=\texttt{argmax}(BQ^{*},\texttt{dim}=0)$ (20) $$

$$ $f_{\text{R}}(e,~{}\mathcal{A},~{}Q^{*})=A[:,I]$ (21) $$

where $e$ represents the architecture embeddings in a pruning batch, $\mathcal{A}$ is the set of architecture codes, $Q^{*}$ is the calculated optimal assignment matrix, and $A$ is the matrix of architecture codes.

At test time, the router function routes an input architecture embedding $e$ to the most similar trained architecture code:

$$ $f_{\text{R}}(e,\mathcal{A})=\text{argmax}_{a\in\mathcal{A}}\frac{e^{T}a}{||e|| ||a||}$ (22) $$

### D.3 Pruning Depth

We apply the vectors $\textbf{u}^{(i)}=[\textbf{u}_{j}^{(i)}]_{j=1}^{M}$ (Eq. 9) to prune the model’s depth with M being the total number of layers we prune. As the U-Net [] architecture of T2I models
has skip connections, we prune the depth layers in the encoder and decoder sides separately.

##### Pruning Depth in Encoder.

We prune the $j$-*th* depth layer $f_{j}$ in the encoder by applying the vector $u_{j}^{(i)}$ with the following
formulation:

$$ $\mathcal{\widehat{F}}_{j}=\mathbf{u}_{j}^{(i)}f_{j}(\mathcal{F}_{j-1})+(1- \mathbf{u}_{j}^{(i)})\mathcal{F}_{j-1}$ (23) $$

where $\mathcal{F}_{j-1}$ is the feature maps of the previous layer. When $u_{j}^{(i)}$ is close to one/zero, the $j$-th layer will be kept/pruned.

##### Pruning Depth in Decoder.

In the decoder, the input of each layer is a concatenation of feature maps of the previous layer
and feature maps of its corresponding encoder layer from the skip connection $\mathcal{F}_{j,skip}$. Thus, we prune the $j$-th depth layer $f_{j}$ in the decoder as:

$$ $\mathcal{\widehat{F}}_{j}=\mathbf{u}_{j}^{(i)}f_{j}(\mathcal{F}_{j-1}|| \mathcal{F}_{j,skip})+(1-\mathbf{u}_{j}^{(i)})\mathcal{F}_{j-1}$ (24) $$

### D.4 Distillation Objective

Assuming $B^{(i)}$ samples get routed to the architecture code $a^{(i)}$ in a training batch ($\sum_{i}B^{(i)}=B$), we train the prompt router and the architecture codes using the following objective:

$$ $\displaystyle\min_{\eta,\mathcal{A}}\mathcal{L}=$ $\displaystyle[\frac{1}{N}\sum_{i=1}^{N}[\frac{1}{B^{(i)}}\sum_{j=1}^{B^{(i)}}[ \mathcal{L}_{\text{DDPM}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})+\lambda_{\text{ distill}}\mathcal{L}_{\text{distill}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})]]]$ (25) $\displaystyle+\lambda_{\text{res}}\mathcal{R}(\widehat{T}(\mathcal{A}),T_{d})+ \lambda_{\text{cont}}\mathcal{L}_{\text{cont}}(\eta)$ $$

$\mathcal{L}_{\text{DDPM}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})$ denotes the denoising objective for the sample $(x^{(i)}_{j},p^{(i)}_{j})$ routed to the sub-network chosen by the architecture code $a^{(i)}$ (Eq. 19). $\mathcal{R}(\widehat{T}(\mathcal{A}),T_{d})$ regularizes the weighted average of the MACs used by architecture codes ($\widehat{T}(\mathcal{A})=\sum_{i}\frac{B^{(i)}}{B}\widehat{T}(a^{(i)})$) to be close to $T_{d}$. We define

$$ $\mathcal{R}(x,y)=\text{log}(\text{max}(x,y)/\text{min}(x,y))$ (26) $$

as it can keep the resource usage close to the target value. $\mathcal{L}_{\text{cont}}$, given by Eq. 13, is the contrastive loss that guides the architecture predictor to map prompt representations to the regions of the space of the architecture embeddings such that their corresponding architecture vectors maintain the similarity between the prompts. Finally,
$\mathcal{L}_{\text{distill}}$ is the distillation objective [], regularizing the pruned model to have similar outputs to the original one. We do distillation at two levels, output level and block level. The output level distillation objective is:

$$ $\mathcal{L}_{\text{output-distill}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})=\mathbb{E }_{\begin{subarray}{c}t\sim[1,T]\\ \epsilon\sim\mathcal{N}(0,I)\\ x^{(i)}_{j,t}\sim q(x_{t}|x^{(i)}_{j})\end{subarray}}[||\epsilon_{\text{ Teacher}}(x^{(i)}_{j,t},p^{(i)}_{j},t;\theta)-\epsilon_{\text{Sub-Net}}(x^{(i) }_{j,t},p^{(i)}_{j},t;\theta,a^{(i)})||^{2}]$ (27) $$

Here, $\epsilon_{\text{Teacher}}(x^{(i)}_{j,t},p^{(i)}_{j},t;\theta)$ denotes the original model’s output and $\epsilon_{\text{Sub-Net}}(x^{(i)}_{j,t},p^{(i)}_{j},t;\theta,a^{(i)})$ denotes the output of the sub-network chosen by the architecture code $a^{(i)}$. The way we prune the U-Net preserves the output shape of each block. Doing so enables us to do distillation at block level as well and regularize the sub-network to match the output of the original model at each block. The block-level distillation objective is:

$$ $\mathcal{L}_{\text{block-distill}}(x^{(i)}_{j},p^{(i)}_{j};a^{(i)})=\mathbb{E} _{\begin{subarray}{c}t\sim[1,T]\\ \epsilon\sim\mathcal{N}(0,I)\\ x^{(i)}_{j,t}\sim q(x_{t}|x^{(i)}_{j})\end{subarray}}[\sum_{b}||\epsilon_{ \text{Teacher}}^{b}(x^{(i)}_{j,t},p^{(i)}_{j},t;\theta)-\epsilon_{\text{Sub- Net}}^{b}(x^{(i)}_{j,t},p^{(i)}_{j},t;\theta,a^{(i)})||^{2}]$ (28) $$

Here, $\epsilon_{\text{Teacher}}^{b}$ and $\epsilon_{\text{Sub-Net}}^{b}$ denote the outputs of block $b$ of the original model and the chosen sub-network of it, respectively. The total distillation loss $\mathcal{L}_{\text{distill}}$ is simply the sum of the output-level loss and the block-level distillation loss:

$$ $\mathcal{L}_{\text{distill}}=\mathcal{L}_{\text{output-distill}}+\mathcal{L}_{ \text{block-distill}}$ (29) $$

### D.5 Fine-tuning

The finetuning loss is a weighted average of the original DDPM objective (Eq. 19) and the distillation loss term (Eq. 29):

$$ $\mathcal{L}_{\text{finetuning}}=\alpha_{\text{DDPM}}\mathcal{L}_{\text{DDPM}}+ \alpha_{\text{distill}}\mathcal{L}_{\text{distill}}$ (30) $$

## Appendix E Experiments

### E.1 Models and Datasets

We use two datasets as our target datasets: Conceptual Captions 3M (CC3M) [] and MS-COCO Captions 2014 [] with approximately 2.5M and 400K image-caption pairs, respectively. We apply APTP to the Stable Diffusion 2.1 (SD 2.1) [] model. On CC3M, we prune SD2.1 with two settings: Base ($0.85$ budget, $16$ experts) and Small ($0.66$ budget, $8$ experts). Similarly, for COCO, we have two settings: Base ($0.78$ budget, $8$ experts) and Small ($0.64$ budget, $8$ experts). We use a pruned SD 2.1 using weight norm pruning [] as a baseline for our main experiments.

### E.2 Experimental Settings

We train at a fixed resolution of 256$\times$256 across all settings. During pruning, we first train the architecture predictor for $500$ iterations as a warm-up phase. During this warm-up phase, we directly use its predicted architectures for pruning. Then, we start architecture codes and train the architecture predictor jointly with the codes for an additional $2500$ iterations. We use the AdamW [] optimizer and a constant learning rate of $0.0002$ for both modules, with a $100$-iteration linear warm-up. The effective pruning batch size is 1024, achieved by training on 16 NVIDIA A100 GPUs with a local batch size of 64. The temperature of the Gumbel-Sigmoid reparametrization (Eq. 9) is set to $\gamma=0.4$. We set the regularization strength of the optimal transport objective (Eq. 5) to $\epsilon=0.05$. We use 3 iterations of the Sinkhorn-Knopp algorithm [] to solve the optimal transport problem []. We set the contrastive loss temperature $\tau$ to $0.03$. The total pruning loss is the weighted average of DDPM loss, distillation loss, resource loss, and contrastive loss (see Eq. 15) with weights $\lambda_{\text{distill}}=0.2$, $\lambda_{\text{res}}=2.0$, and $\lambda_{\text{cont}}=100.0$. After the pruning phase, we fine-tune the experts with the prompts assigned to them for 30,000 iterations using the AdamW optimizer, a fixed learning rate of $0.00001$, and a batch size of $128$. Upon experiments, we observed that higher weights of the DDPM loss result in unstable fine-tuning and slow convergence. As a result, we set the DDPM loss weight in the fine-tuning loss (Eq. 30) $\alpha_{\text{DDPM}}$ to $0.0001$. We set $\alpha_{\text{distill}}=1.0$. For sample generation, we use the classifier-free guidance [] technique with the scale of $7.5$ and 25 steps of the PNDM sampler [].

### E.3 Evaluation

For quantitative evaluation of models pruned on CC3M, we use its validation dataset of approximately 14k samples. For COCO, we sample 30k captions of unique images from its 2014 validation dataset. We report Fréchet inception distance (FID) [], CLIP score [], and Maximum Mean Discrepancy with CLIP Embeddings (CMMD) [] for APTP, the baselines, and SD 2.1 itself.

### E.4 Results

#### E.4.1 Training Loss

Introduced in section 3.3.1, the resource loss regularizes the weighted average of the MACs used by architecture codes ($\widehat{T}(\mathcal{A})=\sum_{i}\frac{B^{(i)}}{B}[\widehat{T}(a^{(i)})]$) to be close to $T_{d}$. We define resource loss as:

$$ $\mathcal{R}(x,y)=\text{log}(\text{max}(x,y)/\text{min}(x,y))$ (31) $$

Fig. 6(a) illustrates the resource loss when applying APTP-Base to prune Stable Diffusion 2.1 using the COCO dataset as the target. This is shown under two conditions: with and without optimal transport following the initial warm-up phase (refer to E for details). APTP effectively regularizes the model so average MACs of the architecture codes is very close to the target budget. Fig. 6(b) presents the contrastive loss (Eq. 13) under the same conditions. APTP maps the prompt embeddings to regions of architecture embeddings such that their corresponding architectures maintain the similarity between prompts.

Figure: (a) Resource loss.
Refer to caption: extracted/5669950/figs/losses/resource.png

#### E.4.2 Visualization of the Impact of Optimal Transport

Fig. 7 displays the assignment of prompts to experts with and without optimal transport. By adding optimal transport, the assignment becomes more diverse, ensuring all experts get enough samples. This results in a significant improvement of performance metrics (See Table 3). Fig. 8 shows the distribution of the number of training CC3M samples mapped to each expert of APTP-Base.

Figure: (a) Without Optimal Transport
Refer to caption: extracted/5669950/figs/ablations/without-ot-assignment.png

Figure: Figure 8: Distribution of CC3M Samples Mapped to Each Expert of APTP-Base, Including Resource Utilization Ratios
Refer to caption: extracted/5669950/figs/ablations/expert_pie_chart.png

#### E.4.3 Analysis of Prompt Router on COCO

Table 4 demonstrates the most frequent words in the prompts assigned to each expert of APTP-Base pruned on COCO, revealing distinct topics and effective specialization. For example, Expert 1 specializes in indoor scences, Expert 5 in wildlife, and Expert 6 in urban scences. Expert 8 which has the highest resource utilization, focuses on images of human beings and hands, an observation consistent with our prompt analysis on CC3M (Table 2). Hands are another category that have been found to be challenging for SD 2.1 [].

**Table 4: The most frequent words in prompts assigned to each expert of APTP-Base pruned on COCO. The resource utilization of each expert is indicated in parentheses.**
| Expert 1 (0.65, Indoor Scenes and Dining) | Expert 2 (0.77, Food and Small Groups) |
| --- | --- |
| table - plate - kitchen - sitting | food - pizza - sandwich |
| Expert 3 (0.78, People and Objects) | Expert 4 (0.79, Sports and Activities) |
| skateboard - surfboard - laptop - tie - phone | tennis - baseball - racquet - skateboard - skis |
| Expert 5 (0.79, Wildlife and Nature) | Expert 6 (0.80, Urban Scenes and Transportation) |
| giraffe - herd - sheep - zebra - elephants | street - train - bus - park - building |
| Expert 7 (0.81, Outdoor Activities and Nature) | Expert 8 (0.83, Domestic Life and Pets) |
| beach - ocean - surfboard - kite - wave | man - woman - girl - hand - bed - cat |

#### E.4.4 Quantitative Results

The quantitative results on the CC3M and MS-COCO datasets (Table 5(b)) indicate that APTP significantly outperforms the baseline Norm method in terms of FID and CLIP scores while also reducing complexity. For both datasets, APTP at 30k and 50k fine-tuning iterations shows lower FID and CMMD values and higher CLIP scores compared to the baseline, demonstrating the advantage of prompt based pruning compared to static pruning ideas for T2I models.

**Table 5: Quantitative results on CC3M and MS-COCO. We report the performance metrics using samples generated at the resolution of $768$ and downsampled to $256$ []. We measure models’ MACs and Latency with the input resolution of $768$ on an A100 GPU. @30/50k shows the model’s fine-tuning iterations after pruning.**
| CC3M |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Method | Complexity | Performance |  |  |  |
| MACs<br>(@768) | Latency ($\downarrow$)<br>(Sec/Sample)<br>(@768) | FID ($\downarrow$) | CLIP ($\uparrow$) | CMMD ($\downarrow$) |  |
| Norm []<br>@30k | 1185.3G | 3.4 | 157.51 | 26.23 | 1.778 |
| Norm []<br>@50k | 1185.3G | 3.4 | 141.04 | 26.51 | 1.646 |
| APTP(0.66)<br>@30k | 916.3G | 2.6 | 60.04 | 28.64 | 1.094 |
| APTP(0.66)<br>@50k | 916.3G | 2.6 | 54.95 | 29.08 | 1.017 |
| APTP(0.85)<br>@30k | 1182.8G | 3.4 | 36.77 | 30.84 | 0.675 |
| APTP(0.85)<br>@50k | 1182.8G | 3.4 | 36.09 | 30.90 | 0.669 |
| SD 2.1 | 1384.2G | 4.0 | 32.08 | 31.12 | 0.567 |

#### E.4.5 Expert Architectures

Figures 9 and 10 display the block-level U-Net architecture of all experts of APTP-Base, with CC3M and COCO as the target datasets, respectively. The retained MAC patterns are markedly different, further corroborating that different datasets require distinct architectures and that a ’one-size-fits-all’ approach is not well-suited for T2I models. For CC3M, down-sampling blocks generally are retained with higher ratios. Overall, APTP tends to prune Resnet Blocks more frequently and aggressively compared to Attention Blocks.

Figure: Figure 9: The block-level retained MAC ratio of the UNet architecture of all experts of APTP-Base applied to Stable Diffusion 2.1 with CC3M as the target dataset. The groups of ResBlocks and the heads of Attention Blocks are pruned based on the outputs of the architecture predictor. The intensity of the color of each block represents the resource utilization of it. The number in each block indicates the precise ratio of retained MACs of the block. Conv_in, Conv_out, and skip connections between corresponding down and up blocks are omitted for brevity.
Refer to caption: extracted/5669950/figs/arch/cc3m-main.png

Figure: Figure 10: The block-level retained MAC ratio of the UNet architecture of all experts of APTP-Base applied to Stable Diffusion 2.1 with COCO as the target dataset. The groups of ResBlocks and the heads of Attention Blocks are pruned based on the outputs of the architecture predictor. The intensity of the color of each block represents the resource utilization of it. The number in each block indicates the precise ratio of retained MACs of the block. Conv_in, Conv_out, and skip connections between corresponding down and up blocks are omitted for brevity.
Refer to caption: extracted/5669950/figs/arch/coco-main.png

#### E.4.6 Samples of APTP

Table 6 displays the prompts for the images in Fig. 3 from CC3M and COCO validation sets.

**Table 6: Prompts for Fig. 3**
| CC3M Prompts |
| --- |
| Saw this beautiful sky on my way home. |
| A silhouetted palm tree with boat tied to it rests on a beach that a man walks across during golden hour. |
| Sketch of a retro photo camera drawn by hand on a white background. |
| From left: person person, the dress on display. |
| Never saw a doll like this before but she sure is sweet looking. |
| The team on the summit. |
| A water drop falls towards a splash already made by another water drop. |
| Husky dog in a new year’s interior. |
| Old paper with a picture of flowers, ranked in a moist environment. |
| People on new year’s eve! |
| I drive over stuff - a pretty cool jeep, 4x4, or truck t-shirt. |
| The crowds arrive for day of festival. |
| A scary abandoned house under a starry sky. |
| Freehand fashion illustration with a lady with decorative hair. |
| Introduce some new flavors to your favorite finger food with these inspired chicken wings. |
| Gloomy face of a sad woman looking down, zoom in, gray background. |
| COCO Prompts |
| A white plate topped with a piece of chocolate covered cake. |
| Decorated coffee cup and knife sitting on a patterned surface. |
| A desk topped with a laptop computer and speakers. |
| A man sitting on the beach behind his surfboard. |
| A pizza type dish with vegetables and a lemon wedge. |
| The browned cracked crust of a baked berry pie. |
| A tennis player in an orange skirt walks off the court. |
| The skier in the helmet moves through thick snow. |
| A giraffe walks leisurely through the tall grass. |
| Several brown horses are standing in a field. |
| A red fire hydrant on a concrete block. |
| A bus driving in a city area with traffic signs. |
| There is a man on a surf board in the ocean. |
| A boat parked on top of a beach in crystal blue water. |
| A small kitchen with low a ceiling. |
| A calico cat curls up inside a bowl to sleep. |

We also provide more samples from APTP-Base pruned on CC3M and COCO. Fig. 11 presents samples from the validation set of CC3M generated by each 16 experts of APTP-Base at $0.85$ MACs budget. Fig. 12 shows samples from the validation set of COCO from each of the 8 experts of APTP-Base pruned to $0.77$ MACs budget.

Figure: Figure 11: Samples of the APTP-Base experts after pruning the Stable Diffusion V2.1 using CC3M [] as the target dataset. Each row corresponds to a unique expert. Please refer to Table 2 for the groups of prompts assigned to each expert.
Refer to caption: extracted/5669950/figs/grids/cc3m-grid-supp.png

Figure: Figure 12: Samples of the APTP-Base experts after pruning the Stable Diffusion V2.1 using MS-COCO [] as the target dataset. Each row corresponds to a unique expert. Please refer to Table 4 for the groups of prompts assigned to each expert.
Refer to caption: extracted/5669950/figs/grids/coco-grid-supp.png

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 2
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?
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: See Weaknesses
soundness: 2
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.
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.
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

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: - Did the authors try their approach on other architectures? even other backbones of Stable Diffusion?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 3
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)
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.
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

Meta-review: 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."

Decision: Accept (Poster)