Not All Prompts Are Made Equal: Prompt-based Pruning of Text-to-Image Diffusion Models
Authors: Alireza Ganjdanesh
Reza Shirkavand
Shangqian Gao
Heng Huang
Venue: ICLR 2025 Poster (2025)
Keywords: Model Pruning
Diffusion Models
Inference Efficiency
Primary area: generative models
TL;DR: We propose a prompt-based pruning method for Text-to-Image diffusion Models, which prunes a pretrained model for a target task into a set of specialized efficient models for different categories of input prompts, given a target compute budget.

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.

Published as a conference paper at ICLR 2025
NOT ALL PROMPTS ARE MADE EQUAL: PROMPT- BASED PRUNING OF TEXT-TO-IMAGE DIFFUSION MOD- ELS
Alireza Ganjdanesh" Reza Shirkavand Shanggiian Gao 2, Heng Huang 1f
Depparmeent of Computer Science. University of Maryland, College Park 2 Department of Computer Science, Florida State University
ABSTRACT
Text-to-image (T21) diffusion models have demonstrated impressive image genera- tion capabilities. Still, their computational intensity prohibits resource-constrained organizations from deploying T2I models after fine-tuning them their inter- nal 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 different prompts. Dynamic pruning addresses this issue by utilizing separate sub- network for each prompt, but it prevents batch parallelism on GPUs. To overcome these limitations, we introduce Adaptive Prompt-Tailoree Pruning (APTP), novel prompt-based pruning method designed for T2I diffusion models. Central to our approach a prompt router model, which learns to determine the required capacity for an input text prompt and routes it to an architecture code, given 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 hyperparameter. We train the prompt router and architecture codes using contrastive learning, ensur- ing that similar prompts mapped 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 (SI) 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. Our code is available here
INTRODUCTION
In recent years, diffusion models (Sohi-Dickstein et al. 2015 Ho et al. 2020; Song & Ermon 2019
Song et al. 2021b) have revolutionized generative modeling. For instance, modern Text-to- Image
(T2I) diffusion models (Adobe FireFly) 2023] Midjourney 2023 Betker al. 2023) like Stable
Diffusion (Rombach et al. 2022 Podell et al. 2024) have achieved remarkable success in generating
realistic images Niicoll et al. 2222 Saharia et al.) 2022, Ramesh et al., 2022) and editing them (Kim) ct al. 2022: Hertz et al. [022) Zhang et al.) 2023b). 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 deploy GPU clouds for large number of users and preventing their utilization on edge devices. Thus, reducing the computational cost of T21 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 T21 models reduce the computation per sampling step or skip steps using the techniques such as distillation (Meng et al. 2023 Salimans & Ho 2022 Habibian al. 2023) and improved noise
Authors Contributed Equally. Conrespondence to (aliganj, zezashkv sanaahkveccs (as und.edu I Partially supported by NSF IIS 2347592,, 2348169, DBI 2405416, CCF 2348306, CNS 2347617.

---

Published as a conference paper at ICLR 2025
Py: bescball gome where the
pyi baseball gamee
bot is worting for the pitch
batter woiting for the pitch
py: beseball player on baseball field
Pyy a bosebell player
Egsebell field
Prompt Router
Pur A skutebearder is doing jumping trick.
Pe Skafchoarder
doing jumping trick
Pe Lags a skoteboord going)
py Logs on skateboard going
down small
down smoll romp
Figure Overview: We prune a text-to-image diffusion model like Stable Diffusion (left) into mixture of efficient experts (right) in a pomppt-bseet manner. Our prompt router routes distinct types of prompts to different experts, allowing experts' architectures be separately specialized by removing layers or channels.
schedules [Song et al. 2021a Nichol& & Dhariwal 202I). Other ideas address the second factor and
propose efficient architectures for T21 models. Architecture modification methods (Zhao et al. 2023 Kim et al. 2023) modify the U-Net (Ronneberger et 2013) of Stable Diffusion (SD) (Rombach al. 2022) 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 (Li et al., 2024] Liu et al., (023a) search for an efficient architecture (Li et al. (024)) or an efficient mixture (Liu et al. (023a)) of pretrained T2I
models in a model zoo. Yet, evaluating each action in the search process and gathering model zoo of T21 models are both extremely costly, making search methods impractical in resource-constrained scenarios.
The efficient architecture design methods (Zhao et al. 2023 Kim et al.. 2023 Li et al., 2024 have shown promise, but they aim to develop 'one-size-fits-ar 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 proprictary target data and deploy the resulting their target data distribution
model for their specific application. They prioritize performance while meeting their computation budget, and the trade-off between efficiency and performance can vary depending
the complexity of the target dataset between organizations. 2) Verifying the validity of each design choice on large-scale datasets used in nne-size-fits-all methods is costly and slow iterate, making them impractical.
Model pruning (Cheng et al. (023) can reduce a model's computational burden to any desired budget with significantly less effort than designing (Zhao al., 2023 Kim al. 2023) searching (Li et al. 2024) 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 cannot benefit from batch-parallelism in modern hardware like GPUs and TPUs.
In this paper, we introduce Adaptive Prompt-Tailoree Pruning (APTP), a novel prompt-based pruning method for T21 diffusion models. APTP prunes 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 set of architecture codes. The prompt router learns to route an input prompt to an architecture code, determining the sub-archa aub-architecture, called expert, of the T2I model use. Each expert specializes in generating images for the prompts assigned to it by the prompt router (Fig. []. and the number of experts is hyperparameter. We train the prompt router and architecture codes using 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 (Sharma et al.) 2018) and MS-COCO (Lin et al. 2014) as the target datasets and prune Stable Diffusion V (Rombach et al. 2022) using APTP in our experiments. APTP outperforms the single-mode pruning baselines, 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 (Chen et al. 2024 Yang et al. 2024j. We summarize our contributions as follows:
We introduce APTP, novel prompt- based pruning method for T2I diffusion models. It is more suitable than static pruning for T2I models as is not imput-agnostic In addition, APTP enables batch parallelism GPUs, which is not possible with dynamic pruning.
2

---

Published as a conference paper at ICLR 2025
APTP trains a prompt router and a set of architecture codes. The prompt router router maps imput 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 framework to train the prompt router and architecture codes using contrastive learning and employ optimal transport to diversify the architecture codes and their corre- sponding 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 s primary focus. Multi-expert (Ganjdanesh et al.) 2024c Zhang et al. 2023a Liu
al., 2023a| Pan al., 2024, Xue et al. 2023) methods employ several models each responsible
for an interval of diffusion model's 's denoising process. These methods design expert model archi-
tectures (Lee et al. 2024, /Zhang et al. [2023]] Xue et al. 2023), train several models with varying
utilize existing pretrained models (Liu et al., 2203aa Pan et al., 2024). However, pretrained experts may not be available, and training several models from scratch, as done by multi-expert methods, is prohibitively expensive. Architecture Design (Zhao iet al. [2023] Yang et al., 2023a Kim et al. (203) approaches redesign the U-Net architecture (Ron- nneberger et al. 2015) of diffusion models enhance its efficiency. MobileDiffusion (Zhao al.
capacities from scratch (Liu et al. (203aa),
2023) does so using empirical heuristics derived from performance metrics on MS-COCO (Lin
al. 2014). BK-SDM (Kim al., 2023) removes some blocks from the SD's U-Net and applies distillation (Hinton et al. 2015) to train the pruned model. Spectral Diffusion (Yang et al. 2023a) introduces 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 non-trivial. Alternatively, SnapFusion (Li et al.) (2024) 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 (Fang et al. 2023) 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 input prompts may vary in complexity, demanding different model capacity levels. Our method differs from existing approaches by introducing a prompt- based pruning tech- nique for T21 models. It is the first method that allocates computational resources prompts based on their individual complexities while ensuring optimal utilization and batch-paralle mulch-parallelizabiiyy See Appendix B for detailed review of related work, with more focus on sampling efficiency studies that are orthogonal to our approach.
3 METHOD
We introduce a framework for prompt-based pruning of T21 diffusion models termed Adaptive Prompt-Tailored Pruning (ATPT). APTP prunes T2I model pretrained on large-scale scale datasets (e.g., Stable Diffuusion (Rombach et al. 2022)) using a smaller target dataset. This approach mirrors the common practice where organizations fine-tune pretrained T21 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 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 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 Gaussian noise (Ho et al., 2020; [Song & Ermon 2019 Sich--icksteinn et al.) 2015). They achieve this by training a denoising model, e(;; 0), parameterized by For a given training image-text pair (IO, p) ~ P. T2I models define forward diffusson process, progressively adding Gaussian noise to the initial image 26 over steps. This process is defined as q(z, (2o) = ((1;;; Vāzzo, (1 I). where di is the forward noise schedule parameter, typically chosen such that (( (n)) + I0, as > T. T2I models are trained using the variational evidence lower bound (ELBO) objective (Ho et al. 2020
3

---

Published as a conference paper at ICLR 2025
Prompt Router
Architecture Codes
skateboarder doing jjumping trick.
Prompt Encoder
Architecture Predictor
Optimal
3
e
Router
Transport
cont
R
utt, 1(0
Remove j- th block
Channel Pruning given
Given 14;
DDPM
Lasttill
1-th layer
Pruning Width
Pruning Depth
Our pruning scheme. We train our prompt router and the set of architecture codes to prune text-to-image diffusion model into mixture of experts. The prompt router consists of three modules. We use Sentence Transformer (Reimers & Gurevych, 2019) our prompt encoder to encode the input prompt into representation z. Then, architecture predictor transforms into the architecture embedding that has the same dimensionality architecture codes. Finally, the router routes the embedding into an architecture code (*), We use optimal transport to evenly assign the prompts in a training batch to the architecture codes. The architecture code a = (u((), v() 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 LDDPM, distillation loss between the pruned and original models Lassill average resource usage for the samples in the batch R, and contrastive objective Leomi encouraging embeddings e preserving semantic similarity of the representations z.
Figure 2:
(DDPM E (Zo [l((z,,,,t;0)) =||2 [.2]
(0,1)
T2I models sample image starting from Gaussian noise FT (N and denoising it using the trained denoising model. We refer to appendix for thorough review of diffusion models.
3.2 PROMPT ROUTER AND ARCHITECTURE CODES
Our approach prunes a diffusion model into a Mixture of Experts, where each expert specializes in handling a distinct group of prompts with different complexities. Each expert corresponds unique sub-network (*) 1). where D represents the total number of prunable units in the T2I model. These sub-networks are optimized to efficiently process their assigned prompt groups while maintaining high performance. To assign prompts the appropriate expert sub-network, design prompt router model comprising three key components:
1. Prompt Encoder:
This module encodes input prompts into semantically meaningful em- beddings. Prompts with similar semantics are mapped to embeddings that are close in the embedding space, ensuring that similar prompts are routed similar sub-networks.
2. Architecture Predictor: The encoded prompt embeddings
further transformed into architecture embeddings e. This step bridges the gap between the high-level semantics of prompts and the architectural configurations needed to process them efficiently.
3. Router Module: The architecture embeddings are finally mapped to specific architecture codes a, which define the structure of the sub-network (expert) that will handle the prompt.
4

---

Published as a conference paper at ICLR 2025
The remainder of this section provides detailed descriptions of how each component functions and contributes to the overall pruning and routing process.
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 fPR A), parameterized by v/. which maps input prompt p to architecture code at
JPR 7), A)
We define the set of learnable architecture codes as A = |a where N is a hyperparame- ter. Also, a E IR, with D being the number of prunable width and depth units in the T2I model. Given a desired constraint Ta on the total compute budget (latency, MACs, etc.) for the prompts in the target dataset, we train the prompt router's ' parameters lj and architecture codes in A to obtain a set of expert models. These experts are specialized, efficient, and performant sub-networks of the T2I model, 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 T2I model. Accordingly, we use a pretrained frozen Sentence Transformer model (Reimers & Gurevych] 2019) our prompt encoder module. It can effectively encode input prompts p into semantically meaningful embeddings z:
2 (PE(P)
fPE(.) is the prompt encoder. We do not explicitly show the prompt encoder's parameters do not train them in our pruning method. The Sentence Transformer can encode semantically similar prompts nearby embeddings and distant from dissimilar ones. We leverage this property in our framework (Eq. [3 to ensure the prompt router maps similar prompts to similar architecture codes.
The architecture predictor module transforms prompt embeddings into architecture embeddings
fAp(2;
7) denotes the parameters of the architecture predictor fAp(-), implemented with single feed- forward layer (more details Appendix D). The embeddings e have the same dimensionality as the architecture codes
3.2.2 ROUTER
The router module takes an architecture embedding e and routes it to an architecture code a E A. A straightforward way to implement the router is to map an input embedding to its nearest architecture code in 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 Ta with a relatively decent performance, causing the prompt router to route all input prompts to it.
To tackle this challenge, we employ optimal transport our router module during the pruning phase. Formally, let E = (ei.... enn represent B architecture embeddings in training batch, and A (a(1), a (N)) represent the N codes. The goal is to find an assignment matrix Q =[0,.-. 9g) that maximizes the similarity between architecture embeddings and their assigned architecture codes:
max T ATE) cH(Q) QEQ
H is entropy term H(Q) - Es, 4zj log(g;z) and is the regularization strength. It has been shown (YM. al. 2020 Caron al. 2020) that high values of lead to uniform assignment matrix Q. causing all codes in to collapse to a single code. Thus, we set c to a small value. Further, we impose an equipartition (YM. et al. (2020) constraint on Q that architecture embeddings in a batch are assigned equally to architecture codes:
1 IN, QTIN= = B 1 N
Q = {Q E R.R.N NxB QQ1 =
(6)
5

---

Published as a conference paper at ICLR 2025
Here, [(B,N) are vectors of ones with length and N. respectively. These constraints enforce that, on average, 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. with the constraints in Eq. cann be solved using the fast version (Cuturi) 2013) of the Sinkhorn-Knopp algorithm and the solution has the normalized exponential matrix form (Cuturi 2013):
A F diag(())
Q' (exxpl
Here, and are renormalization vectors that can be calculated using few iterations of the Sinkhorn- Knoop algorithm. With e set to a small value, the matrix BQ will have columns that are close one-hot vectors, which we use assign the architecture embeddings the architecture codes. In summary, the router module's function during the pruning process is:
a = Ax
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 FR in Appendix D.2
3.3 PRUNING
We divide each architecture code ( E A into two sub-vectors ( = (x())..))) We utilize the vectors (t) and (()) to prune the depth layers and determine widths of the layers, respectively.
())).
simple approach to prune the width of a layer (a depth layer) is to use binary vectors ( (u(0), 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 (f) (u (t) that are continuous and differentiable for pruning. We calculate them as:
(()) sigmoid (+() 9u
( + 9u
sigmoid
0m ~ Gumbel (0, 1) represents a noise vector sampled from the Gumbel distribution (Gumbel 1954). and y is the temperature. This formulation, known as the Gumbel- sigmoid reparameteriza- tion (Jang et al. 2017] Maddison et al. (017)) is a differentiable approximation of sampling from Bernoulli distributions with parameters s gmoid cl (w(") and sigmoid (u()) ). When the tempera- ture y is set appropriately, the vectors (() and ( will be close to binary vectors, and we use them
for pruning the layers' width and depth layers. Assuming (() = [v] 1 the model's layers, we prune the width of the (-th layer as:
R=1 where L is the number of
F = Fi
(10)
In this equation, Fi represents feature maps of the I-th layer, and O is the element-wise multiplica- tion. We prune the channels of the convolution layers in the ResBlocks (He et al.) 2016)), attention heads in the Transformer layers Vaaswani et al. 2017). and the channels the forward layers in the Transformer layers.
In a similar manner, we apply the vectors ( that we prune) for pruning the model's depth. Specifically, we prune the j-th depth layer fy as:
[u 1= (where M the number of depth layers
M
= ( + (1- F-11
Jy-1 denotes the previous layer's feature maps. The granularity of our depth pruning is ResBlock or Transformer layer in the U-Net of the T2I model. We provide more details in Appendix [D3] In summary, we employ the architecture vectors (6) (v(() 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 Ta on aggregate.

---

Published as a conference paper at ICLR 2025
Contrastive Training. Given training batch with B prompts, compute their prompt embed- dings 2 (Eq. 3) and architecture embeddings e (Eq. 4). We calculate architecture vectors e':
sigmoid c+
(12)
where y and g have the same definitions Eq. 9 We define the cosine similarity of two vectors as (m, n) m² /// m ||||| and we use the following objective train the architecture predictor:
1
Leom (1)
(oglog()s,)) + (1 - Ta.j Voig))oo(11 - 81.))
(13)
B
exp(sim( 24, TT)
exp(sim(e,, ej)/T)
TJ
SiJ
(14)
Σ 1= exp (ap(sim( 24, 2x)/T)
Σ k= exp(sim(e/, e()/T)
T is temperature parameter. Eq. 3 regularizes the architecture predictor to map representations to the regions of the space of the architecture embeddings e such that their corresponding architecture vectors e' maintain the similarity between the prompts. Note that we do not actually use the architec- ture vectors e' to prune the model (Sec. 3.2.2 3.3 Yet, we apply our contrastive regularization
vectors e' instead of embeddings This design choice is result of our observation that applying feom to embeddings may not lead to diverse architecture vectors (6) (v(6), u()) 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 e' and a (*),
In contrast, Eq. 13 implicitly diversifies the architecture vectors a (f) The reason is that the router routes embeddings codes (Eq. { that have high similarity with each other (Eq. 7 Also, Eq. [3 distributes embeddings e in the space of the architecture embeddings such that the architecture vectors become similar (different) for similar (dissimilar) prompts. As the vectors e' and a are calculated in similar manner (Eqs. 12). the diversity of vectors e' implies the same for vectors a (1).
Assuming (+) samples get routed the architecture code (*)) in training batch (E, C,BO) =B), we train the prompt router and the architecture codes using the following objective:
N
73(0)
min L = n.A
(1), :a .P (+))+ + Adistin Crisin ( (x .P (0) ;
B
(15)
Ares R (AA), Ta) + Acone
(6) CDDEM((Z) Pj ); ) denotes the denoising objective for the sample (x)) p) routed to sub- network chosen by the architecture code ( (A), Ta) regularizes the weighted average of be close to Ta. We define R(x, log (max(x, v) w)/min(2z, (z, and (Adisill Ares Acont) are hyperparameters. Finally, Cdisnll is distillation objective (Kim et al. 2023) regularizing the pruned model having similar outputs to the original one. We refer to Appendix D.4 for more details about our distillation objective.
the MACs used by architecture codes (T(AA) Σ. B(0) (T()())) [T())((
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 DD3 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.
EXPERIMENTS
We Conceptual Captions 3M (CC3M) (Shamma al. 2018 and MS-COCO (Lin et al. 2014) as our target datasets to prune the Stable Diffusion (SD) V2.1 model to demonstrate APTP's ef- fectiveness. On CC3M, we prune the model with Base: (0.85 MACs, 16 experts) and Small: (0.66
7

---

Published as a conference paper at ICLR 2025
CC3M
COCO
Figure 3:
Samples of the APTP-Base experts after pruning the Stable Diffusion V2.1 using CC3M (Sharma al. 2018) and COCO (Lin et al. 2014) the target datasets. Expert IDs shown on the top right of images. (See Table for prompts)
MACs, 8 experts) configurations. On MS-COCO, we perform Base: (0.78 MACs, 8 experts) and
Small: (0.64 MACs, experts) pruning settings. We set (Adistlll Aess Acont )= 2.0, 100) (Eq. 15 and the temperature r (Eq. 14) to 0.03. We evaluate all models with FID (Heusel et al. 2017]. CLIP (Hessel et al.) 2021), 1) and CMMD (Jayasumana et al. (023) 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 256 for calculating the metrics, using
the 25-steps PNDM (Liu et al. 2022) sampler following BK-SDM (Kim al. 2023). We refer to Appendix B for more details.
4.1 COMPARISON RESULTS
As APTP is the first pruning method specifically designed to prune a pretrained T2I model on target dataset, we compare its performance with SD V2.1, weight norm pruning (Li et al.] 2017), and two recently proposed static pruning baselines, namely Structural Pruning (SP) (Fang et al. 2023) and BKSDM (Kim al. 2023). Table[[]shows the results. We fine-tune APTP, SP and BKSDM for 30k iterations after pruning and give Norm-pruning 50k fine-tuning iterations to ensure they all reach their final performance level.
CC3M: Table Ta summarizes the results on CC3M. With a similar MACs budget and latency values, APTP (0.85) significantly outperforms the Norm pruning (Li et al. 2017), SP (Fang al., (023). and BKSDM (Kim et al. (023)) baselines with a significant margin in terms of FID, CLIP, and CMMD scores. It also achieves 15% less latency while showing close performance scores to SD 12. Notably, APTP (0.66) has approximately 23% less latency and 21% lower MACs budget than the baselines, but still outperforms them on all metries. These results illustrate the advantages of prompt-based compared to static pruning.
We present the results for MS-COCO in Table Tb APTP (0.78) reduces the latency of SD by 22.5% while preserving its CLIP score and achieving a close CMMD score. Further, with a similar latency, APTP (0.78) significantly outperforms the static pruning baselines with at least 3.71 FID (BKSDM), 2.34 CLIP (SP), and 0.042 CMMD (BKSDM) scores. Similar to the results on CC3M, APTP (0.64) achieves approximately 19.3% less latency than the Norm pruning and SP while outperforming them all 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. and Appendix E.4.6
MS-COCO:
8

---

Published as a conference paper at ICLR 2025
Table 1: Results on CC3M and COCO. We report performance metries using samples generated at resolution of 768 then downsampled to 256 (Kim et all 2023). We measure models' MACs/Latency with the input resolution of 768 on an A100 GPU, @ 30/50k shows fine-tuning iterations after pruning.
MS-COCO
Complesty
Pedformance
Complesity
Performnce
Mebod
Latreey (1)
Mellod
Lanmry (I) (Stu/Saippee)
(87760)
(SecSutgle)
(1) (LPP())
(MMP
FED (1)
CUIPT
CMMD
Nonn
Li
2017
1155.3G
344
141.04
26.51
1.646
Numa
Li ctal. 2017
1677.4G
3.1
47.35
26.51
1.136
SP Fask ct 2023
1192.G
3.5
75.81
26.83
1243
SP
Fag al. 2025
1071.G
3,3
53.00
26:06
0.966
BXSDM
et
3.3
97.27
26.56
1083.4G
3.1
26331
25,50
0611
APTINO 66)
916.3G
266
60.04
28
1,004
APTINO
890.00
2.5
30.12
20.56
0867
AFTPILRS)
1182.8G
3.4
36.77
30.84
0.673
APTPN6,75)
1078.60
3.1
22.60
31.32
0.589
(a)
(h)
Table 2: The most frequent words in prompts assigned to each expert of APTP- pruned on CC3M. The resource utilization of each expert indicated parentheses.
Expert (0.72)
Expert (0.73)
Expert 3 (0.75)
Expert (0.76)
View Sunset City Building Sky
View Boal Sea
Artist Actor
Actor Dress
Poetrait
Expert 5 (0.77)
Expert 6 (0,78)
Expert 7 (0.79)
Expert (0.79)
Illustration Portrait Photo
Player Rall
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
Person Player Team Couple
Expert (0.86)
Expert (0,90)
Expert (0.95)
Expert 16 (0.98)
Room House
Art Artist Digital
Fond Water
Person Man Woman Text
CC3M low budget
COCO low budget
CC3M high budget
CoOO high budget
APTP (Ours)
MAtiA
50 V2.1
Patterns sand beach.
White paper Christmas tree red matte Background.
slice of pizza that is sitting on paper plate.
Pointing of ship in grass with
wamon background looks down morual. and away from camera with forlorn look on her face.
Tiny
infant holding baby toothbrush his hand looking of it.
dag a little girl stands near.
lighthouse in
the background.
Figure 4: Comparison of samples generated by low and high budget experts APTP-Base vs. SD V2.1 on CC3M and MS-COCO validation sets.
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. dispeyjssplays 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 (Chen) et al. 2024| Yang et al.] 2024h, and our prompt router can automatically discover and route them to nigher-capacity experts. In contrast, paintings and illustrations are assigned to lower-budget experts, as they seem to be casier for the model to generate. We provide examples of high and low resource prompts for APTP-Bass on both CC3M and COCO in Fig.

---

Published as a conference paper at ICLR 2025
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
sigmoid( +g)//T) and u = (igmoid((1: + 9u)/T) (Egq. 9 to prune a single model. It directly trains two vectors (7v /uu to do so ("Uni-Arch Baseline' in Table 3. Then, we start with 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 (Lin et al. 2014). We fine-tune all pruned models with 10k iterations, and [able(]present the results. We observe that the contrastive training (Eq. [3] alone fails to improve results from pruning a single model 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 single expert. As result, all experts except receive insufficient training samples and generate low-quality samples after fine-tur tinettniig, 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. addition, distillation can further improve the architecture search process the experts, resulting in more performant mixture. In summary, these results validate the effectiveness of our design choices for APTP.
Table Ablation results of APTP's s components on 30k samples from MS-COCO (in et al. validation set. We fine-tune all models for 10k iterations after pruning.
2014
Method
MACs( @ 768)
Latency( @ 768)
FID (J)
Clip Score (t)
CMMD (4)
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
In our second ablation experiment, we explore the impact of the number of experts APTP. We prune SD V to 80% MACs budget using APTP with 4, 8, and 12 experts on MS-COCO. Fig 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 experts dataset-dependent. This observation illustrates that prompt-basee based pruning is more suitable than static pruning for T2I models.
5 CONCLUSION
In this paper, we develop Adaptive Prompt-Tailored Prun- ing (APTP), the first prompt-based pruning method for text-to-image (T2I) diffusion models. APTP takes T2I model, pretrained on large-scale data, and prunes using a targel dataset, resembling the common practice that orga- nizations fine-tune T2I models on their internal data before deployment. The core element of APTP is prompt router module that learns to decide the model capacity required to generate sample for an input prompt, routing it to architecture code given a desired compute budget. Each architecture code corresponds model, specializing in generating images for the prompts that prompt router routes to 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 experi- ments in which prune Stable Diffusion (SD) V 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, analysis on APTP's prompt router reveals that it automatically discover challenging prompt types for SD, like generating text, humans, and fingers, routing them to experts with high compute budgets.
35
31.4
31.3
33
31
31.1
29
30.9
30.8
sub-network of the T2I
25p
30.7
Number of Experts
Figure 5: Ablation Results for the number experts of APTP on MS-COCO.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Adobe FireFly. Firefly, 2023. Firefly: y.mefl
Shubham Agarwal, Subrata Mitra, Sarthak Chakraborty, Srikrishna Karanam, Koyel Mukherjee, and Shiv Saini. Approximate caching for efficiently serving diffusion models. arXiv preprint arXiv:2312.04429 2023.
James Betker, Gabriel Goh, Li Jing, Tim Brooks, Jianfeng Wang, Linjie Li, Long Ouyang, Juntang Zhuang, Joyce Lee, Yufei Guo, et al. Improving image generation with better captions. Computer Science. https://cdn. openai. com/papers/dall---. pdf. 2(3):8, 2023.
Han Cai, Chuang Gan, Tianzhe Wang. Zhekai Zhang, and Song Han. Once-for-all: Train one network and specialize for efficient deployment. In International Conference on Learning Representations, 2020. httLttttt////perrreet
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems, 13:9912-999 9924924 2020.
Thibault Castells, Hyoung-Ki Kyu Song, Bo-Kyeong Kim, and Shinkook Choi. Ld Ld-pruner: Efficient pruning of latent diffusion models using agnosttc insights. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2024 Workshops, Seattle, WA, USA, June 17-18, 2024, pp. 221- IEFE, 2024. doi: 10.1 10090//RR663.22222.000000 URL https //doi org 0. 11 www.010PPw6332222222 2024. 00087
Haoxing Chen, Zhuoer Xu, Zhangxuan Gu, Yaohui Li, Changhua Meng, Huijia Zhu, Weiqiang Wang, et al. Diffute: Universal text editing diffusion model. Advances in Neural Information Processing Systems, 36, 2024.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, al. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.00426, 2023.
Hongrong Cheng, Miao Zhang, and Javen Qinfeng Shi. A survey deep neural network pruning- taxonomy, comparison, analysis, and recommendations. arXiv preprint arXiv:2308.0676) 2023.
Huanpeng Chu, Wei Wu, Chengjie Zang, and Kun Yuan. Qned: Quantization noise correction for diffusion models. arXiv preprint 240:::2400.9914000 2024.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Ming Ding, Wendi Zheng, Wenyi Hong, and Jie Tang. Cogview2: Faster and better text-to-image generation via hierarchical transformers. Advances in Neural Information Processing Systems, 35: 16890-16902, 2022.
Sara Elkerdawy, Mostafa Elhoushi, Hong Zhang, and Nilanjan Ray. Fire together wire together: A dynamic pruning approach with self-supervised mask prediction. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 12454-12463, 2022.
Gongfan Fang, Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. In Advances in Neural Information Processing Systems, 2023.
Rohit Gandikota, Joanna Materzynska, Tingrui Zhou, Antonio Torralba, and David Bau. Concept sliders: Lora adaptors for precise control in diffusion models. arXiv preprint arXhv:2311.12092. 2023.
Alireza Ganjdanesh, Shangqian Gao, and Heng Huang. Interpretations steered network pruning via amortized inferred saliency maps. In European Conference on Computer Vision, pp. 278-296 Springer, 2022.
11

---

Published as a conference paper at ICLR 2025
Alireza Ganjdanesh, Shangqian Gao, and Heng Huang. Effconv: efficient learning of kernel sizes for convolution layers of cnns. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7604-7612, 2023.
Alireza Ganjdanesh, Shangqian Gao, Hirad Alipanah, and Heng Huang. Compressing image-to- image translation gans using local density structures on their learned manifold. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 12118 22118-12126 2024a.
Alireza Ganjdanesh, Shangqian Gao, and Heng Huang. Jointly training and pruning cnns via learnable agent guidance and alignment. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16058-16069, 2024b.
Alireza Ganjdanesh, Yan Kang, Yuchen Liu, Richard Zhang, Zhe Lin, and Heng Huang. Mixture of efficient diffusion experts through automatic interval and sub-network selection. European Conference Computer Vision, pp. 54-71. Springer, 2024c.
Emil Julius Gumbel, Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Government Printing Office, 1954.
Amirhossein Habibian, Amir Ghodrati, Noor Fathima, Guillaume Sautiere, Risheek Garrepalli, Fatih Porikli, and Jens Petersen. Clockwork diffusion: Efficient generation with model-step distillation. arXiv preprint arXiv:2312.08122 2023.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Yang He and Lingao Xiao. Structured pruning for deep convolutional neural networks: A survey. IEEE Transactions Pattern Analysis and Machine Intelligence, pp. 1-20, 2023. doi: 10.1109/
TPAML.2023.333461
Yefei He, Luping Liu, Jing Liu, Weijia Wu, Hong Zhou, and Bohan Zhuang. Ptqd: Accurate post- training quantization for diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Yihui He, Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. Ame: Automl for model compression and acceleration mobile devices. In Proceedings of the European conference computer vision (ECCV), pp. 784-800, 2018
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt- to-prompt image editing with cross attention control. arXiv preprint arXiv:2208. 200::2228.016666 2022.
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference- free evaluation metric for image captioning. arXiv preprint arXiv:2104.07711 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge neural network. arXiv preprint arXiv:1503.0253: 2015.
Jonathan Ho and Tim Salimans. Classifier-frre diffusion guidance. CoRR, abs/2207.12598 2022. doi: 1044550//RRXXX220711155 12598
URL htttp:::/aoi..or/11.4455555
2207
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6855 6851, 2020.
Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. Advances in Neural Information Processing Systems, 33: 9782-9793, 2020.
12

---

Published as a conference paper at ICLR 2025
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with somtex- In International Conference on Learning Representations, 2017. URL URL/atpps://ppeeeeiier
Sadeep Jayasumana, Srikumar Ramalingam, Andreas Veit, Daniel Glasner, Ayan Chakrabarti, and Sanjiy Kumar. Rethinking fid: Towards a better evaluation metric for image generation. arXiv preprint urXiv:2401.09000 2023.
Minguk Kang, JanJJunnaa Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10124-10134, 2023.
Bo-Kyeong Kim, Hyoung-Kyu Song, Thibault Castells, and Shinkook Choi. On architectural compression of text-to-image diffusion models. arXiv preprint arXiv:2305. 15798,, 2023.
Gwanghyun Kim, Taesung Kwon, and Jong Chul Ye. Diffusionclip: Text-guided diffusion models for robust image manipulation. In Proceedings of the IEEE/CVE Conference on Computer Vision and Puttern Recognition, pp. 2426-2435, 2022.
Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, [baccka00644003333002333333-3333
Ashish Kumar, Daneul Kim, Jaesik Park, and Laxmidhar Behera. Pick-or-mix: Dynamic channel sam- pling for convnets, 2024.
wowoooooommmooorroooidooo :dd= FXB4V
Yunsung Lee, Jin Young Kim, Hyojun Go, Myeongho Jeong, Shinhyeok Oh, and Seungtaek Choi. Multi-architecture multi-expert diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 13427-13436, 2024.
Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. In International Conference on Learning Representations, 2017. https/
//openreview.oo wwwp:evwwwwweeeeeotreemmemamm rid
Xiuyu Li, Yijiang Liu, Long Lian, Huanrui Yang. Zhen Dong, Daniel Kang, Shanghang Zhang, and Kurt Keutzer. Q-diffusion: Quantizing diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.
17535-17545, 2023.
Yanyu Li, Huan Wang, Qing Jin, Ju Hu, Pavlo Chemerys, Yun Fu, Yanzhi Wang, Sergey Tulyakov, and Jian Ren. Snapfusion: Text-to-image diffusion model on mobile devices within two seconds. Advances in Neural Information Processing Systems, 36, 2024.
Tao Lin, Sebastian U. Stich, Luis Barba, Daniil Dmitriev, and Martin Jaggi. Dynamic model pruming with feedback. In International Conference on Learning Representations, 2020. URL
net/ www.foooouuddeeeoo Seee 1SFWB
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision- ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings. Part V 13, pp. 740-755. Springer, 2014.
Enshu Liu, Xuefei Ning, Zinan Lin, Huazhong Yang, and Yu Wang. Oms-dpm: Optimizing the model schedule for diffusion probabilistic models. arXiv preprint arXiv:2306.08866 2023a.
Enshu Liu, Xuefei Ning, Huazhong Yang, and Yu Wang. A unified sampling framework for solver searching of diffusion probabilistic models. arXiv preprint 13X0v:2312.07243, 2023b.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In International Conference on Learning Representations, 2019. URL /atts::::: h/openree//eeeeeieeer
net/
13

---

Published as a conference paper at ICLR 2025
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL /tRLtttss////eeeeeeeee forum? 1 d=P1 KWVd2yBkY
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. In The Twelfih International Conference on Learning Representations, ICIR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024a. UtI.tttss:// /ttp::////p/eeeeeeeeeeeme net/forurut d=1k4y2bb40pp
Xuewen Liu, Zhikai Li, Junrui Xiao, and Qingyi Gu. Enhanced distribution alignment for post-training quantization of diffusion models. arXiv preprint arXiv:2401. 04588, 2024b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, JCLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https:://openerr/////
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM--olver++: Fast solver for guided sampling of diffusion probabilistic models, 2023. URL https:/7
Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepcache: Accelerating diffusion models for free. arXiv preprint arXiv:2312.00856 2023.
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous
relaxation of discrete random variables. In International Conference on Learning Representations, 2017. /ttpsttpss//ppennee.......eee Stt JE5L5g]
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffuusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Midjoumey. Midjourney, 2023. URL https: / www mi wwt.:mimjjounneego wtm. cwm.home
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint I 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffussion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171, PMLR, 2021.
Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, and Anima Anandkumar. T-stitch: Accelerating sampling in pre-trained diffusion models with trajectory stitching. arXiv preprint urXiv:2402. 14167, 2024.
Nilesh Prasad Pandey, Marios Fournarakis, Chirag Patel, and Markus Nagel. Softmax bias correction for quantized generative models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1453-1458, 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. SDXL: Improving latent diffusion models for high-resolution image synthesis. In The Twelfth International Conference Learning Representations, 2024. URL
/tttps:/copenevvee www.eeiiewweeeeeeeeeeeeee nnt/ fow.mmidddii222Re50
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint 200:22044.0125,, 1(2):3, 2022.
Nils Reimers and Iryna Gurevych. Centence-BERT: Sentence embeddings using Siamese BERT- networks. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-JCCNPP), pp. 3982-3992 Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v 1/D19- 1410.
URL www.http://eec //acss://acl antho cgy.orrg/D19-1411
14

---

Published as a conference paper at ICLR 2025
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- on computer vision and pattern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp
Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCA. 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 241. Springer, 2015.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35:36479-36494, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations, 2022. URL
pest/oorumiid-hiee TiIIIIIIIIIIIIIIII
Piyush Sharma, Nan Ding, Sebastian Goodman, and Radu Soricut. Conceptual captions: A cleaned, hypernymed, image alt-text dataset for automatic image captioning. In Proceedings of ACL, 2018.
Junhyuk So, Jungwon Lee, Daehyun Ahn, Hyungjun Kim, and Eunhyeok Park. Temporal dynamic quantization for diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2266-2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- fional Conference Learning Representations, 2021a. URL httttttt///ppeeee....ee n
[oormm?d==StliarrCe
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
net formm?
iid-PXTI12RRHH3
Yuda Song, Zehao Sun, and Xuanwu Yin. Sdxs: Real-time one-step latent diffusion models with image conditions. arXiv preprint arXiv:2403.16627 2024.
Siao Tang, Xin Wang, Hong Chen, Chaoyu Guan, Zewen Wu, Yansong Tang, and Wenwu Zhu. Post-training quantization with progressive calibration and activation relaxing for text-to-image diffusion models. arXiv preprint arXiv:2311.0663 2023.
Yehui Tang, Yunhe Wang, Yixing Xu, Yiping Deng, Chao Xu, Dacheng Tao, and Chang Xu. Manifold regularized dynamic network pruming. In Proceedings of the IEEE/CVF conference computer vision and patlern recognition, pp. 5018-5028, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers for dif- fusion models by differentiating through sample quality. In International Conference Learning Representations, 2022a. URL https:: /ttp::::://ppeeeee ew net f wtt.rrummmiivva3aueeeee
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning to efficiently sample from diffusion probabilistic models, 2022b. URL https://///......tt..tt http://ppppeeee..eee/ee net htt/fooummfiii
T.Oz OxOx0x
15

---

Published as a conference paper at ICLR 2025
Colin White, Mahmoud Safari, Rhea Sukthanker, Binxin Ru, Thomas Elsken, Arber Zela, Debadeepta Dey, and Frank Hutter. Neural architecture search: Insights from 1000 papers. arXiv preprins arXiv:2301.08777, 2023.
Felix Wimbauer, Bichen Wu, Edgar Schoenfeld, Xiaoliang Dai, Ji Hou, Zijian He, Artsiom Sanakoyeu, Peizhao Zhang, Sam Tsai, Jonas Kohler, et al. Cache me if you can: Accelerating diffusion models through block caching. arXiv preprint WrXii:2312.032099 2023.
Yilun Xu, Mingyang Deng. Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. arXiv preprint arXiv:2306. 18006.1488 14878, 2023.
Zeyue Xue, Guanglu Song, Qiushan Guo, Boxiao Liu, Zhuofan Zong, Yu Liu, and Ping Luo. RAPHAEL: text-to-image generation via large mixture of diffusion paths. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023, URL ht 621655c7dc483688 38c88888c888555 oot.oo/ooooooaattttttototteeeeee
httpt::pppppppp..eeeep papersppppspaaa nips. cc/paper htmml 7 -bstefffA-bbtttcc
Xingyi Yang, Daquan Zhou, Jiashi Feng, and Xinchao Wang. Diffusion probabilistic model made slim. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 225622 2023a.
Yuewei Yang, Xiaoliang Dai, Jialiang Wang, Peizhao Zhang, and Hongbo Zhang. Efficient quantiza- tion strategies for latent diffusion models. arXiv preprint arXiv:2312.054311 2023b.
Yukang Yang, Dongnan Gui, Yuhui Yuan, Weicong Liang, Haisong Ding, Han Hu, and Kai Chen. Glyphcontrol: Glyph conditional control for visual text generation. Advances in Neural Information Processing Systems, 36, 2024.
Lewei Yao, Renjie Pi, Hang Xu, Wei Zhang, Zhenguo Li, and Tong Zhang. Joint-detnas: Upgrade your detector with nas, pruning and dynamic distillation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10175-10184, 2021.
Asano YM., Rupprecht C., and Vedaldi A. Self-labelling via simultaneous clustering and rep- resentation learning. In International Conference on Learning Representations, 2020. URL
htttss//oppeneevier .ne httt:/oorumiiid-yyx.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Bureu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu. Scaling autoregressive models for content-rich text-to-image generation. Trans. Mach. Learn. Res., 2022, 2022. URL
ner/forum? id JKHND
Huijie Zhang, Yifu Lu, Ismail Alkhouri, Saiprasad Ravishankar, Dogyoon Song, and Qing Qu Improving efficiency of diffusion models via multi-stage framework and tailored multi-decoder architectures. arXiv preprint arXiv:2312.0918) 2023a.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control text-to-image diffusion models. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pp. 3836-3847, 2023b.
Qinsheng Zhang, Molei Tao, and Yongxin Chen. gDDIM: Generalized denoising diffusion implicit models. In The Eleventh International Conference Learning Representations, 2023c. URL htttps://openeeee
net. 1
Yang Zhao, Yanwu Xu, Zhisheng Xiao, and Tingbo Hou. Mobilediffusion: Subsecond text-to-image generation on mobile devices. arXiv preprint arXiv:2311.16567, 2023.
Barret Zoph and Quoc Le. Neural architecture search with reinforcement learning. In International Conference Learning Representations, 2017. URL http::/:s////peeeeee.ee... net fformm? id=r1Ue8Hcxq
16

---

Published as a conference paper at ICLR 2025
A LIMITATIONS AND BROADER IMPACTS
We believe limitation of our work may be that our experimental scope relatively limited. We chose CC3M Sharma et al. (2018) and MS-COCO et al. (2014) as they contain at least 100k image-text pairs and are diverse to enable us to simulate the woal-world case of fine-tuning pretrained T2I model like SD v2.1 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 Chen et al. )024); Yang et al. (2024): Gandikota et al. (2023) to alleviate these issues. Finally, note that our method applications, but it may also facilitate harmful content generation, necessitating responsible usage of generative models.
targer dataset before deploying it. Yet. exploring APTP's performance
other
enable organizations to serve their models with lower cost for GenAl
B 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 Habibian et al. (2023);
Meng et al. )003)); Salimans & Ho (2022); Song et al. (2024); Liu et al. (2024a), learning optimal denoising time-steps Watson al. (2022b a), designing improved noise schedules Nichol & Dhariwal
(20TT); Song et al. (2021a) Zhang et al. (202cc), caching intermediate computations Ma et al. (2023):
Agarwal et al. (2023); Wimbauer al. (2023), and proposing faster solvers Lu al. (2023); Xu
et al. (2023); Liu et al. (2023b). The second group is aimed at improving the architectural efficiency of diffusion models, which is the main focus of our paper. Multi-expert [ee et al. (2024); /Zhang et al. (2023a) Liu et al. (2023a); |ann et al. (2024) Xue et al. (2023) 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 Lee et al. (2024); Zhang et al. (2023a); Xue et al. (2023), train several models with varying capacities from scratch et al. (2023a) or utilize existing pretained experts Liu et al. 2023a);; Pan et al. (2024) 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 ZZao et al. (2023): Yang et al. (2023)) Kim et al. (2023) redesign the architecture of diffusion models to enhance efficiency. MobileDiffusion Zhao et al. (2023) modifies the U-Net Ronneberger al.) (2015) architecture of Stable Diffusion (SD) based on empirical heuristics derived from the model's performance on the MS-COCO Lin et al. (2014) dataset. BK-SDM Kim et al. (2023) 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 Yang et al. (2023a) introduces 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 the architecture design methods to other tasks and compute budgets is non trivial..... Quantization methods So et al. (2024); He et al. (2024); Pandey et al. (2023) ) Yang et al. (2023b); [iu et al. (2024b); Tang et al. (2023); Chu et al. 2024 Li et al. (2023) reduce the precision of model weights and activations during the forward pass to accelerate the sampling process. Different from these methods, SnapFusion Li et al. (2024) 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 Fang et al. (023) 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 prun-
ing Cheng et al. (2023); Ganjdanesh et al. 2022 2024b) and NAS Zoph & Le 2017); Liu et al. (2019); Cai et al. (2020); Yao et al. (2021 Hou et al. (200); White et al. (2023); Ganjdanesh et al. (202)) 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 pruming approaches He et al. (2018); Li et al. (2017); Han et al.
(2015): Castells al. (2024); (Gajjdanesh al. (002a)) prune pretrained model and use the pruned
17

---

Published as a conference paper at ICLR 2025
model for all imputs. Conversely, dynamic pruning methods Elkerdawy et al. (2022 Kumar et al. (2024): Lin et al. (2020); Tang et al. (2021) employ a separate sub-network the model for each input. Although statie 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 White et al. (2023); He & Xiao (2023); Cheng al. 2023) 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-pares autch-parallelizbbiity
C OVERVIEW OF DIFFUSION MODELS
Given a random variable Xo ~ P, the goal of diffusion models models/Sohl-Dickseeee et al. (2015); Ho et al. (2020) is to model the underlying distribution P using a training set 'D { } of samples. To do so, first, diffusion models define a forward process parameterized by in which they gradually perturb each sample zo with Gaussian noise with the variance schedule of B5:
q( (z V1 11111-Bac 1, 3,1)
(16)
where E [1, Thus, (221/20) has a Gaussian form:
(<x1(20)= = Nass Varzo, -1-60 1I)
(17)
where 0₂ = - Bt and ā₂ I ag The noise schedule 3, is usually selected Ho et al. 2020 such that ((rr) + (O, Assuming 3e is small, diffusion models approximate denoising distribution (xx+-1(222) (2x)) by parameterized Gaussian distribution po Po(x1-1 (2x) = ((x-111 /o. (ax
9. co f)), of and of is often set to 3t. Diffusson models implement ca(.) with neural VI-ai network called the denoising model and train it with the variational evidence lower bound (ELBO)
objective Ho et al. (2020):
(DDPM = E t11t]] ||e(2,,t;0) - -el1 (o.1).
(18)
Similarly, T2I diffusion models train a denoising model using pairs of image and text prompts (zo.P) model the distribution P(xo p) of images given an input text prompt:
(f) E
(ZoP))>P [-[1T]
[lc( c/12
(19)
c=N().T)
T21 Diffusion models generate new sample by sampling initial noise from øт p(xT)
(0 /) and iteratively denoising using the denoising model by sampling from po po(x,-11 (rr, p). Thus, the sampling process requires T sequential forward calculation of the denoising model, making it slow and costly process.
D MORE DETAILS OF APTP
In this section we provide more details of the architecture predictor, the prompt router module, and pruning method.
18

---

Published as a conference paper at ICLR 2025
D.1 ARCHITECTURE PREDICTOR
(AP(-) in Eq. 4 is the architecture predictor. The architecture predictor changes the dimensionality of the dimensionality of e and codes a, which is the number of prunable width and depth units in the T2I model. We implement the fAP function with single feed-forward layer. Our prompt encoder, Sentence Transformer Reimers Gurevych| (2019), has an output dimension of 768. The total number of prunable units in SI> 2.1 is 1620, so the architecture predictor has an input dimension of 768 and an output dimension of 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 the architecture codes Thus, during the pruning process, the definition of function fR as follows:
= argaxx( diim = 0)
(20)
fR(e, A, Q*) -{;=][[1]] 4]
(21)
where represents the architecture embeddings in pruning batch, 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:
eT |||||al||
fR(e, A) = argmax, aEA
(22)
D.3 PRUNING DEPTH
(Eq. 9 to prune the model's depth with M being the total number layers we prune. As the U-Net Ronneberger al. (2015) architecture of T2I models has skip connections, we prune the depth layers in the encoder and decoder sides separately.
We apply the vectors u u M 1=1
Pruning Depth in Encoder. We prune the j-th depth layer fj in the encoder by applying the vector (4) with the following formulation: Mj
(23)
where Fj-1 is the feature maps of the previous layer. When uj 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 Fj.skip- Thus, we prune the j-th depth layer fú in the decoder as:
8, 3 1 + (1-1, (
(24)
D.4 DISTILLATION OBJECTIVE
Assuming (+) samples get routed to the architecture code (+) in training batch (E, B() B). we train the prompt router and the architecture codes using the following objective:
N
B(
min L n.d
=
(DDeM(
-D () + Adisti Caisun (3 -D a (1)
Bit)
(25)
1=1
+ Ares R (TAA), Ta) + Acone Coone
19

---

Published as a conference paper at ICLR 2025
(1) ( denotes the denoising objective for the sample (2) P) routed the sub- CDDPM (1j .Pj network chosen by the architecture code (+)) (Eq. 19.. (A), Ta) regularizes the weighted average
(1)
( (6)
of the MACs used by architecture codes (A) Σ. B ((a (()))) to close to Ta. We define
R(x, y) log(max(x, //)//mmn(xx y))
(26)
as it can keep the resource usage close to the target value. Lcont. given by Eq. 13 is the contrastive loss that guides the architecture predictor to map prompt representations the regions of the space of the architecture embeddings such that their corresponding architecture vectors maintain the similarity between the prompts. Finally, Laisin the distillation objective Hinton al. (2015). 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:
(0) Couput distill (1))) PP (0)
E
[-1,7] (0,10.
[Teachrr (3 (6) (8) Pj ,t; ,t; Sub Nec (x, (0)
t: 0t;0, )
(27)
Here, TTeccer (anucher()),)) (T).. (2)
1; denotes the original model's output and 6Suh 6Suh (x)
(6) 0, a,a()) ) de-
notes the output of the sub-network chosen by the architecture code (*)) The way we prune the U-Net preserves the output shape of each block. Doing so enables us to distillation 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:
Cotock-disnn (x)) PD) ) W|1T c~N(0.1)
e feachher (2), 11
(6) .t:,),)))) ( (6) ( Sab-NNa
Pj t: 0, al) (1
(28)
Here, exeaceer and Lubb-Net denote the outputs of block 6 of the original model and the chosen sub- network of respectively. The total distillation loss Ldistill is simply the sum of the output-level loss and the block-level distillation loss:
Edistill = Coutput-distill + Lplock-disill
(29)
D.5 FINE-TUNING
The finetuning loss is weighted average of the original DDPM objective (Eq. 19 and the distillation loss term (Eq.
Liinetning CDDPPM + Eedistin
(30)
E EXPERIMENTS
E. MODELS AND DATASETS
We use two datasets our target datasets: Conceptual Captions 3M (CC3M) Sharma et al. (2018)
and MS-COCO Captions 2014 Lin et al. (2014) with approximately 2.5M and 400K image-caption
pairs, respectively. We apply APTP to the Stable Diffusion 2.1 (SD 2.1) Rombach et al. 2022 model.
On CC3M, we prume SD2.1 with two settings: Base (0.85 budget, 16 experts) and Small (0.66 budget, 8 experts). Similarly, for COCO, have two settings: Base (0.78 budget, 8 experts) and Small (0.64 budget, experts). We use a pruned SD 2.1 using weight norm pruming Li et al. (201)) as a baseline for our main experiments.
20

---

Published as a conference paper at ICLR 2025
E.2 EXPERIMENTAL SETTINGS
We train at a fixed resolution of 256 across all settings. During pruning, we first train the architecture predictor for 500 iterations 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 jointy with the codes for an additional 2500 iterations. We use the Adam W Loshchilov & Hutter (2019) optimizer and constant learning rate of 0.0002 for both modules, with T00-iteration linear warm-up. The effective pruning batch size 1024, achieved by training on 16 NVIDIA A100 GPUs with a local batch size of 64. The temperature of the Gumbel-Sigmoid reparametrizationn (Eq. is set to = 0.4. We set the regularization strength of the optimal transport objective (Eq. 5) to 0.05. We use 3 iterations of the Sinkhorn- Linkhorn-Knopp algorithm Cuturi (2013) solve the optimal transport problem Caron et al. (2020). We set the contrastive loss temperature T 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 ai = 0.2, Aves 2.0, and Acom = 100.0. 0. After the pruning phase, fine-tune the experts with the prompts assigned them for 30,000 iterations using the AdamW optimizer, a fixed learning rate of 0. 0.00001, and a batch size of 128. Upon experiments, we observed that higher weights of the DDPM loss result unstable fine-tuning and slow convergence. As result, we set the DDPM loss weight in the fine-tuning loss (Eq. 33) aDDPM to 0.0001. We set Oddssill = 1.0. For sample generation, we use the classifier-fre free guidance Ho & Salimans (2022) technique with the scale of 7.5 and 25 steps of the PNDM sampler Liu al. 2022).
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 unique images from its 2014 validation dataset. We report Fréchet inception distance (FID) Heusel et al. (2017), CLIP score Hessel et al. (021). and
Maximum Mean Discrepancy with CLIP Embeddings (CMMD) Jayasumana et al. (2023 for APTP, the baselines, and SD 2.1 itself.
E.4 RESULTS
E.4.1 TRAINING LOSS
Introduced in section 3.3.
the resource loss regularizes the weighted average of the MACs used by 8
architecture codes (T E:
B [T (a)) to be close to Ta. We define resource loss as:
R(x, y) = log(max(<r, hmipm//inx.//// y))
(31)
Fig. 6a illustrates the resource loss when applying APTP- 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 foor details). APTP effectively regularizes the model average MACs of the architecture codes is very close the target budget. Fig. (6b presents the contrastive loss (Eq. [3) under the same conditions. APTP maps the prompt embeddings to regions of architecture embeddings such that their corresponding architectures maintain the similarity between prompts.
E.4.2 VISUALIZATION OF THE IMPACT OF OPTIMAL TRANSPORT
Fig. 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 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.
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-Bari pruned 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 which has the highest resource utilization, focuses images of human beings and hands, an observation
21

---

Published as a conference paper at ICLR 2025
Resource Loss
Contrastive Loss
-Wiha Saagpn Op Torqpet
(a) Resource loss.
(b) Contrastive loss.
Figure 6: Resource and Contrastive loss observed when applying APTP-Bas -9Te-Base with a MAC budget of
0.77 to prune Stable Diffusion 2.1 using the COCO dataset. The comparison is made between two settings: with and without optimal transport. APTP both adheres to the target MAC budget and finds architecture vectors that maintain the similarity between the prompts.
045
375
280
575
365
670
360
3.55
235
Without Optimal Transport
(b) With Optimal Transport
Figure 7: Comparison of sample assignments in a batch to experts with and without optimal transport. The incorporation of optimal transport results more diverse assignment pattern. In the figure, each square represents prompt within the batch, and the color signifies the budget level of the expert assigned to the prompt. Higher-resoe resource experts are indicated by darker blue.
22

---

Published as a conference paper at ICLR 2025
Samples Mapped Expert Including Rasource Utilization
(Uil: 0.23!
Expert 16 tuna: 0.72)
8.75
10.8%
Expert
0.215
Export 11 [unl:: 0.75)
2.0%
1.6%
(Ubli 0.801
Expert
21.5%
0.855
(uil:: b.act
Expert [Uhii 6.801
11.5%
2.8%
7.7%
Pxpert
Expert (upel 0.81)
Expert
experfysl::
Expert 3 (UBl: n.so)
Figure 8: Distribution of CC3M Samples Mapped Each Expert of APTP-Base, Including Resource Utilization Ratios
23

---

Published as a conference paper at ICLR 2025
consistent with our prompt analysis CC3M (Table 2). Hands are another category that have been found to be challenging for SD 2.1 Gandikota et al. (2023).
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 (0.65, Indoor Scenes and Dining)
Expert 2 (0.77, Food and Small Groups)
table plate kitchen sitting
food pizza sandwich
Expert (0.78, People and Objects)
Expert 4 (0.79, Sports and Activities)
skateboard surfboard laptop- tie phone
tennis baseball raoquet- skateboard skis
Expert (0.79, Wildlife and Nature)
Expert (0.80, Urban Scenes and Transportation)
giraffe- herd- sheep- zebra- elephants
street train bus park building
Expert (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 indicate that APTP signifi- cantly 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.
We also benchmark our method and the baselins on the PartiPrompt Yu et al. [022) dataset using PickScore/Krsttin
PickScore Kirstain et al. (202)). The prompts in this benchmark can be considered diitt--diittibttion for model as many of them are significantly longer and semantically different from the ones in MSCOCO. We report the PickScore Kirstain et al. (2023) as a proxy for human preference and present the results in Table 2 below. We can observe that the Pickscore of the pruned model is only 1% below the original Stable Diffusion 2.1, indicating that the pruned model can preserve the general knowledge and generation capabilities of the Stable Diffusion model. Table 6 demonstrates the results.
Conducting systematic, controlled evaluations of large-scale text-to-image models remains challeng- ing, as most existing models, datasets, or implementations are not publicly available. Training a new model from scratch is prohibitively expensive, even when training code is accessible. Extending the design decision of architecture design methods to other compute budgets or settings is highly non-trivial. In contrast, our approach can be applied in plug-and-plass manner to any target model capacity, delivering competitive performance with drastically reduced hardware compute time.
In Table we compare our model to recent state-of oh-the-art text-to-image models, based on available information and their reported values, acknowledging significant differences in training datasets, iteration counts, batch sizes, and model sizes. We also include a column detailing the compute used to train these models, with estimates derived from the papers and public sources. Our findings demonstrate that APTP achieves competitive results post-training while requiring several orders of magnitude less compute time.
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 'one-size- siee--iiiztttlalll fils-all' approach is not well-suited for T2I models. For CC3M, down wown-sampling blocks generally are retained with higher ratios. Overall, APTP tends to prune Resnet Blocks more frequently and aggressively compared to Attention Blocks.
24

---

Published as a conference paper at ICLR 2025
Dowel
Drent
Drank
Repert
Pagwrt
Brwnl
Crond
Brwwl
Brand
Drant
Cowwr
Sepert
Repert
Rrwelt
96
Sxpert
Rxpart
MAGEAA stach
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.
25

---

Published as a conference paper at ICLR 2025
Expert
Expert
Expert
Expert
Attention
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 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_on _out, and skip connections between corresponding down and blocks are omitted for brevity.
Figure 10:
26

---

Published as a conference paper at ICLR 2025
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 Kim et al. (002)). We measure
models' MACs and Latency with the input resolution of 768 on A100 GPU. 30/50k shows the model's fine-tuning iterations after pruning.
OC3M
MS COCO
Cumplexit
Performanes:
Comploxity
Prrfomuane
Method
Latency 231 (SeciSample) (S7788)
Methed
Latency (Sec/Sample) (47763)
CLIP
CLIP
2017
2017
3.4
Nexn
1077.4G
47.35
28.51
1.136
APTP10.00 (6)
6:30k
916.3G
2.6
60.04
28.64
APTP(0-6))
8900G
25
29.12
29.98
0.867
APTP10.56
6 50k
916.3G
2.6
54.95
29.08
APTP/0
SP5
890,0G
255
36.17
30.21
0.729
APTPPI. 55)
00.00
1182.80
3.4
36.77
30.54
0.675
APTP/O.75 75)
1076.6
3.1
22.60
51,32
0.569
APTP:0.55 55)
6.50k
1182.8
3.4
36.09
30,90
D.669
APTP10.75)
107666
3.1
22.26
51,35
0.561
SD
1384 30
4.0
$2.08
31.12
0.567
SD 2.1
1384.00
4.0
15.47
$1.33
0.500
Table Results on PartiPrompts. We report performance metrics using samples generated at resolution of 768. We measure models' MACs/Latency with the input resolution of 768 on an A100 GPU. @3/550k shows lie-ttnnng iterations after pruning.
Train on CC3M
Train on MS-COCO
Complexity
PatipPrompts
Complexity
PartiPrompts
Method
MAC's (6) 768)
Latency (See/Sample) (6768)
PickScore (1)
Method
(6 768)
MACs
Latency (Sec/Saimle) (62 768)
PickScore (*)
Norm
ct
2017
1185.3G
3.4
18,114
Norm
ct al.
2017
1077.4G
3.1
18.563
0 S0K
90%
SP Pang
2024
1192.1G
18:727
SP
al
2023
1071.40
3.3
19.317
B0K
B530K
BKSDM
Kim al. 2023
1180.0G
3.3
19.491
BKSDM
Kim al.
2023
1085.4G
3.1
19.94]
(PTP(0.66)
APTP(0.64)
6r 30k
916.3G
2.6
19.597
30K
890.0G
2.5
20.626
APTP(0.8S)
APTP(0.78)
@ 30k
1182.8G
3.4
21.049
30k
1076.6G
3.1
21.150
SD 2.1
1384.2G
4.0
21.316
SD
1384.2G
4.0
21.316
(a)
(b)
Table 7: Comparison of APTP and SOTA Text-to-lmage architecture design methods.
Midels
Type
Sampling
Steps
FID-30(())
CLII((T)
Params (B)
Images (R)
Compuse (GPLL/TPU days)
GigaGAN
Kang et 2023
GAN
1-step
9.09
D.9
0.98
4783 A 100
Cogview-22
24.0
DALLE-2
Imagen
S 20123
Diffusion
DDPM
256
7:27
3.60
0.45
7132 TPU
S132.1
Diffusion
DDIM
50
9.62
0.304
0.86
>-2
8334 A
Pixart
Diffusion
DDPM
20
10.65
0.6
0.025
753 A
SnapFusior
Diffusion
Disilled
Mobil-Difuuisn (/hpo al.
2023
Difiusion
DDIM
50
0.325
0.39
0.15
7680 TPU
RAPHAEL [ue at 2023
Difiusion
DDIM
6.61
3.0
>>
60000 1
APTP-Rise (D)3
Diffasion
DDIM
50
19.14
0.318
0.65
0.0001
6.5 1
27

---

Published as a conference paper at ICLR 2025
E.4.6 MORE VISUAL RESULTS
Table 8 displays the prompts for the images Fig. 3 from CC3M and COCO validation sets.
CC3M Prompts
Saw this beautiful sky my way home. A silhouetted palm tree with boat tied to it rests on beach that a man walks across during golden hour. Sketch of retro photo camera drawn by hand on a white background. From left: person person, the dress on display.
Never saw doll like this before but she sure is sweet looking.
The team on the summit.
water drop falls towards a splash already made by another water drop. Husky dog in a new year's 's interior.
Old paper with picture of flowers, ranked in moist environment.
People on new year's eve!
drive over stuff pretty cool jeep, 4x4, or truck t-shirt.
The crowds arrive for day of festival.
scary abandoned house under starry sky.
Freehand fashion illustration with a lady with decorative hair.
Introduce some flavors 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 piece of chocolate covered cake. Decorated coffee cup and knife sitting on a patterned surface.
A desk topped with laptop computer and speakers.
A man sitting the beach behind his surfboard.
A pizza type dish with vegetables and a lemon wedge.
The browned cracked crust of a baked berry pie.
tennis player in an orange skirt walks off the court.
The skier in the helmet moves through thick snow.
A giraffe walks leisurcly through the tall grass. Several brown horses are standing in field.
A red fire hydrant on a concrete block.
A bus driving in city
with traffic signs.
There is man
surf board the ocean.
A boat parked on top of a beach in crystal blue water.
A small kitchen with low ceiling.
calico cat curls up inside a bowl to sleep.
Table 8: Prompts for Fig.
3
We also provide more samples from APTP-Base pruned on CC3M and COCO. Fig. I presents samples from the validation set of CC3M generated by each 16 experts of APTP-Base 0.85 MACs budget. Fig. 12 shows samples from the validation set of COCO from each of the 8 experts of APTP-Base pruned 0.78 MACs budget.
We provide samples from APTP pruned models on various styles in Fig. 13 APTP preserves the original model's versatility and generalizes well to out-of-distribution prompts.
We provide samples comparing APTP to the best baseline, namely BKSDM et al. 2023 )) on different prompts from the PartiPrompts Yu et al. (2022) datasets Fig. 14 APTP produces better images compared to the baseline.
28

---

Published as a conference paper at ICLR 2025
AB
49
Figure 11: Samples of the APTP-Base experts after pruning the Stable Diffusion V2.1 using ( et al. (2018) as the target dataset. Each row corresponds to a unique expert. Please refer Table 2 for the groups of prompts assigned to each expert.
29

---

Published as a conference paper at ICLR 2025
Figure 12: Samples of the APTP-Base experts after pruning the Stable Diffusion V2.1 using MS- COCO Lin al. (2014) as the target dataset. Each row corresponds unique expert. Please refer to Table 4 for the groups of prompts assigned to each expert.
APTP-Base COCO APTP-Base CCM
S.21
landscape with bold, unnatural colors fauvism style
A serene scene cherry blos- soms by Hokusai
A forest by Sherry Akrami
A painting a woman in the style Picasso
watercolor painting of forest the style Gogh
Anime of kaiju monster
Photo of sad man by Nobuy- oshi Araki
Figure 13: APTP generalizes to various styles, even if they are not present in the target dataset.
30

---

Published as a conference paper at ICLR 2025
SD
BKSDM
COCO
APTP
SD
CC3M
BKSDM
APTP
Wkpeoson
Figure 14: A visual comparison of randomly sampled generated images from Stable Diffusion 2.1, BKSDM (the best baseline), and APTP. The top section displays models trained the COCO dataset, while the bottom section showcases samples from models trained on CC3M. Overall, APTP produces higher-quality images compared to the baseline. For prompts, Table 9
Prompts
handpalm with leaves growing from it close-up of the keys of piano the Eiffel Tower in desert a view of the Earth from the moon the silhouette of the Milllenium Wheel at dusk
an avocado on a table
glass of orange juice to the right of plate with buttered toast on it
Styracosaurus
view of clock tower from below a hat a trash bin
a woman with sunglasses and red hair
a close-up of an ostrich's 's face black t-shirt
A helicopter flies over Yosemite.
a spaceship hovering over The Alamo
a corgi
A bar of chocolate without a wrapper that has the word "WRAPPER" printed on it.
a laptop screen showing bunch of photographs
street
Table 9: Prompts for Fig. 14
31


Reviews:
[{"review_id": "jV6b3XNAXC", "reviewer": "Reviewer_Ff7o", "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.", "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? \n\n2. Training diffusion models often incorporate classifier-free guidance. Is the proposed method compatible with training under this manner?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 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,\n2. 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: \n\n1. 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.\n\n2. 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.\n\n3. 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.\n\n\n\n[1] Zhao, Yang, et al. \"Mobilediffusion: Subsecond text-to-image generation on mobile devices.\" arXiv preprint arXiv:2311.16567 (2023).\n\n[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).\n\n[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.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gqlahOrGQQ", "reviewer": "Reviewer_uiaa", "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.", "questions": "See Weaknesses", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The idea behind the paper is technically sound and novel.\n2. The paper is well written. \n3. 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 \n - different types of architecture - It is currently unknown if the proposed methods is generalizable across architectures (DiT, MMDiT etc). \n - Small datasets - How does the method perform when data is limited?\n - Fine grained concepts - How does their method handle expert assignment when concepts are fine-grained (breeds of different animals)\n2. 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].\n3. 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\n\n\n[1] RAPHAEL: Text-to-Image Generation via Large Mixture of Diffusion Paths, Xue et al", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "srKDIZnTyU", "reviewer": "Reviewer_vGxh", "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.", "questions": "- Did the authors try their approach on other architectures? even other backbones of Stable Diffusion?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "- The idea of pruning different parts of the network for each prompt is non-trivial and interesting.\n- 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:\n\n- *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. \n\n- *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.\n\n- *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?\n\nAs step skipping approaches became more popular recently, I believe including some of these as baselines is a must for such an approach.\n\n**(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.\n\n**(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). \n\n**(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.\n\n**(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). \n\n[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.\n\n[2] Sauer, Axel, et al. \"Adversarial diffusion distillation.\" European Conference on Computer Vision. Springer, Cham, 2025.\n\n[3] Meng, Chenlin, et al. \"On distillation of guided diffusion models.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mWeQsMRDRR", "reviewer": "Reviewer_xWxf", "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.", "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?\n* Did you do any additional experiments concerning varying resolutions and aspect ratios and their impact on the pruned-image quality?\n* 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?\n* How specific is the APTP-model to its training data? How do out-of-distribution prompts impact the quality of the generations?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Good Illustrations and concise explanation\n* The core idea is an interesting combination of known concepts in a new, but highly relevant setup.\n* 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.\n\nKirstain et al. 2023 https://proceedings.neurips.cc/paper_files/paper/2023/hash/73aacd8b3b05b4b503d58310b523553c-Abstract-Conference.html\n[Pernias et al. 2024] https://proceedings.neurips.cc/paper_files/paper/2023/hash/73aacd8b3b05b4b503d58310b523553c-Abstract-Conference.html", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)