SnapFusion: Text-to-Image Diffusion Model on Mobile Devices within Two Seconds
Authors: Yanyu Li
Huan Wang
Qing Jin
Ju Hu
Pavlo Chemerys
Yun Fu
Yanzhi Wang
Sergey Tulyakov
Jian Ren
Venue: NeurIPS 2023 poster (2023)
Keywords: Text-to-Image
Diffusion model
mobile devices
distillation
TL;DR: A text-to-image model that runs in less than two seconds on mobile devices with a similar quality as the Stable Diffusion.

Abstract: Text-to-image diffusion models can create stunning images from natural language descriptions that rival the work of professional artists and photographers. However, these models are large, with complex network architectures and tens of denoising iterations, making them computationally expensive and slow to run. As a result, high-end GPUs and cloud-based inference are required to run diffusion models at scale. This is costly and has privacy implications, especially when user data is sent to a third party. To overcome these challenges, we present a generic approach that, for the first time, unlocks running text-to-image diffusion models on mobile devices in **less than 2 seconds**.  We achieve so by introducing efficient network architecture and improving step distillation. Specifically, we propose an efficient UNet by identifying the redundancy of the original model and reducing the computation of the image decoder via data distillation. 
Further, we enhance the step distillation by exploring training strategies and introducing regularization from classifier-free guidance. Our extensive experiments on MS-COCO show that our model with $8$ denoising steps achieves better FID and CLIP scores than Stable Diffusion v$1.5$ with $50$ steps. Our work democratizes content creation by bringing powerful text-to-image diffusion models to the hands of users.

SnapFusion: Text-to-lmage Diffusion Model on Mobile Devices within Two Seconds
Yanyu Li
Huan Wang 22 122
Qing Jin 1. Ju Hu'
Pavlo Chemerys' Jian Ren -f
Yun Fu?
Yanzhi Wang²
Sergey Tulyakov'
*Snap Inc.
Northeastern University
Project Page: https//snap-rrssaachh. wwt.h..iio/Saaapii.
Figure 1: Example generated images by using our efficient text-to-image diffusion model.
Abstract
Text-to-image diffusion models can create stunning images from natural language descriptions that rival the work of professional artists and photographers. However, these models are large, with complex network architectures and tens of denoising iterations, making them computationally expensive and slow to As result, high-end GPUs and based inference are required to run diffusion models at scale. This is costly and has privacy implications, especially when user data is sent to a third party. To overcome these challenges, we present generic approach that, for the first time, unlocks running text-to-image diffusion models mobile devices in less than seconds. We achieve so by introducing efficient network architecture and improving step distillation. Specifically, we propose an efficient UNet by identifying the redundancy of the original model and reducing the computation of the image decoder via data distillation. Further, enhance the step distillation by exploring training strategies and introducing regularization from classifier-free guidance. Our extensive experiments on MS-COCC show that our model with 8 denoising steps achieves better FID and CLIP scores than Stable Diffusion v1.5 with 50 steps. Our work democratizes content creation by bringing powerful text-to-image diffusion models to the hands of users.
Introduction
Diffusion-baast text-to-image models [1, 2, 3, 4] show remarkable progress in synthesizing pho- torealistic content using text prompts. They profoundly impact the content creation [5, 6], image
Equall contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

editing and in- in-painting [7, 8, 9. 10, 11], super-resolution [12], video synthesis [13, 14], and 3D assets generation [15, 16, 17], to name few. This impact comes at the cost of the substantial increase the computation requirements to run such models [18, 19, 20, 21J. As a result, to satisfy the necessary latency constraints large scale, often cloud-based inference platforms with high-end GPU are required. This incurs high costs and brings potential privacy concerns, motivated by the sheer fact of sending private images, videos, and prompts to a third-party service.
Not surprisingly, there are emerging efforts to speed up the inference of text-to-image diffusion models on mobile devices. Recent works use quantization [22. 23] or GPU-aware optimization to reduce the run time, i.e., accelerating the diffusion pipeline to 11.5s on Samsung Galaxy S23 Ultra [24]. While these methods effectively achieve a certain speed-up on mobile platforms, the obtained latency does not allow for a seamless user experience. Besides, none of the existing studies systematically examine the generation quality of on-device models through quantitative analysis.
In this work, we present the first text-to-image diffusion model that generates an image on mobile devices in less than 2 seconds. To achieve this, we mainly focus on improving the slow inference speed of the UNet and reducing the number of necessary denoising steps. First, the architecture of UNet, which is the major bottleneck for the conditional diffusion model (as we show in Tab. 1), is rarely optimized in the literature. Existing works primarily focus on post-training optimizations [25, 26]. Conventional compression techniques, e.g., model pruning [27, 28] and architecture search [29, 30], reduce the performance of pre-trained diffusion models [31], which is difficult to recover without heavy fine-tuning. Consequently, the architecture redundancies are not fully exploited. resulting in limited acceleration ratio. Second, the flexibility of the denoising diffusion process is not well explored for the on-device model. Directly reducing the number of denoising steps impacts the generative performance, while progressively distilling the steps can mitigate the impacts [32, 33]. However, the learning objectives for step distillation and the strategy for training the on-device model have yet to be thoroughly studied, especially for models trained using large-scale datasets.
This work proposes a series of contributions address the aforementioned challenges:
We provide an in-depth analysis of the denoising UNet and identify the architecture redundancies.
We propose novel evolving training framework to obtain an efficient UNet that performs better than the original Stable Diffusion v1.5 while being significantly faster. We also introduce data distillation pipeline compress and accelerate the image decoder.
We improve the learning objective during step distillation by proposing additional regularization, including losses from the v-prediction and classifier- free guidance [34].
Finally, we explore the training strategies for step distillation, especially the best teacher-student paradigm for training the on-device model.
Through the improved Step distillation and network architecture development for the difFusion model, our introduced model, SnapFusion, generates a 512 x 512 image from the text on mobile devices in less than 2 seconds, while with image quality similar to Stable Diffusion v1.5 [4] (see example images from our approach in Fig. 1).
Model Analysis of Stable Diffusion
2.1 Prerequisites of Stable Diffusion
Diffusion Models gradually convert the sample x from real data distribution Pasta (x) into a noisy version, i.e., the diffusion process, and learn to reverse this process by denoising the noisy data step by step [35]. Therefore, the model transforms a simple distribution, egg., random Gaussian noise, to the desired more complicated distribution, e.g., real images. Specifically, given a (noise-prediction) diffusion model éø (-) parameterized by 0, which is typically structured as a UNet [36, 1], the training
can be formulated as the following noise prediction problem [35, 1, 2]:
min E,, (x)) ||ēo(t, 2.) <l12.
(1)
where t refers to the time step: e is the ground-truth noise; z; = a,x + ore is the noisy data; a and og are the strengths of signal and noise, respectively, decided by noise scheduler. A trained

---

diffusion model can generate samples from noise with various samplers. In our experiments, we use DDIM [37] to sample with the following iterative denoising process from to previous time step t',
0, Zi) + (t, 21):
ae
(2)
at
where 2y will be fed into (0(.) again until 1' becomes O, i.e. the denoising process finishes.
Latent Diffusion Model / Stable Diffusion. The recent latent diffusion model (LDM) [4] reduces the inference computation and steps by performing the denoising process in the latent space, which is encoded from pre-trained variational autoencoder (VAE) [38,39]. 39]. During inference, the image is constructed through the decoder from the latent. LDM also explores the text-to-image generation, where text prompt embedding с is fed into the diffusion model as the condition. When synthesizing images, an important technique, classifier-free guidance (CFG) [34], is adopted to improve quality.
ét Ze, c) = wép (f, Zt, c) (w 1) Zz, Ø),
where (t, Ze, 2) represents the unconditional output obtained by using null text Ø. The guidance scale w can adjusted to control the strength of conditional information on the generated images to achieve the trade-off between quality and diversity. LDM is further trained on large-scale datasets [40], delivering a series of Stable Diffusion (SD) models [4]] We choose Stable Diffusion v1.5 (SD-v1.5) as the baseline. Next, we perform detailed analyses to diagnose the latency bottleneck of SD-v1.5.
2.2 Benchmark and Analysis
Here we comprchensively study the parameter and computation intensity of the 1.5. The in in-tepth analysis helps us understand the bottleneck to deploying text-to-image diffusion models on mobile devices from the scope of network architecture and algorithm paradigms. Meanwhile, the micro-level breakdown of the networks serves as the basis of the architecture redesign and search.
RDD
400
Table 1: Latency Comparison between Stable Diffu-
sion v1.5 and our proposed efficient diffusion models
Up-1/6
(UNet and Image Decoder) on iPhone 14 Pro.
Up-1/2
Stable Diflusion v1.5
Text Engoser
UNet
VAE Decoder
Mi8-1/8
Iaput Resolution
tokens
x
61
Down: 1/32
#Parametors (M)
123
860
50
Latency ims)
4
~1,700'
369
Dowe: 1/16
Inference Stegs
50
Total Lotency
85,000
369
Down- 1/wn-1/8
Our Model
Teat Fncoder
Our UNet
Our Irange Becoder
100
400
Iapul Resolution
tokens
61%
CA parms
ResNet pamms
lattryy
Ros lonene
#Parmeterss (M)
$48
Figure 2: Latency (iPhone 14 Pro. ms) and pa- rameter (M) analysis for cross-attention (CA) and ResNet blocks in the UNet of Stable Diffusion.
Latency ims)
4
230
116
Inference Steps
2
$
1
Total Lalency
1,840
116
Macro Prospective. As shown in Tab. and Fig. 3, the networks of stable diffusion consist of three major components. Text encoder employs ViT-H model [41] for converting input text prompt into embedding and is executed in two steps (with one for CFG) for each image generation process, constituting only tiny portion of inference latency (8 ms). The VAE decoder takes the latent feature to generate an image, which runs as 369 ms. Unlike the above two models, the denoising UNet is not only intensive in computation (1.7 seconds latency) but also demands iterative forwarding steps to ensure generative quality. For instance, the total denoising timesteps is set to 50 for inference in SD-v1.5, significantly slowing down the on-device generation process the minute level.
Breakdown for UNet. The time-conditional (t) UNet consists of cross-attentiin and ResNet blocks. Specifically, cross-attention mechanism is employed at each stage to integrate text embedding
(c) into spatial features: Cross-Attention Qz, Ke, Vz) Softmax( Qz K < Ve. where Q is projected Val from noisy data 2,, K and V are projected from text condition, and dl is the feature dimension. UNet also uses ResNet blocks to capture locality, and we can formulate the forward of UNet as:
(t, 2)) II (rros--Attention((,, c), ResNet(z,, ()) ())-
We notice the latency varies depending on the phones and use three phones get the average speed.
3

---

Our Step Distillation
Diffusion
Image
E
Teacher:
Z0
2, ax
oye
Zr
Vo(C)
vo(c)
Iterative Denoising
(o(
Prediction
Eol
Z
Ze''
Ze
Zr
Bq.7/10
Bq.7/10;
x
DDIM
Studenf.
UNet
Vn(c)
VAE
Loss:Cai
0
Prompt
ViT
Zt
Latents
Eq.8/10
Figure 3: Workflow of text-to-image diffusion model (left) and the proposed step distillation (right).
The distribution of parameters and computations of UNet illustrated Fig. 2, showing that parameters concentrated the middle (downsampled) stages because of the expanded channel dimensions, among which ResNet blocks constitute the majority. contrast, the slowest parts of UNet are the input and output stages with the largest feature resolution, spatial cross-attentions have quadratic computation complexity with respect to feature size (tokens).
3 Architecture Optimizations
Here we investigate the architecture redundancy SD-v1.5 to obtain efficient neural networks. However, is non-trivial to apply conventional pruning [42, 43, 44, 45] architecture search [46, 47, 30] techniques, given the tremendous training cost SD. Any permutation in architecture may lead to degraded performance that requires fine-tuning with hundreds thousands of GPUs days. Therefore, we propose an architectureeeoolingg method that preserves the performance of the pre-trained UNet model while gradually improving its efficacy. As for the deterministic image decoder, apply tailored compression strategies and a simple yet effective prompt-driven distillation approach.
3.1 Efficient UNet
From our empirical observation, the operator changes resulting from network pruning or searching lead to degraded synthesized images, asking for significant training costs to recover the performance. Thus, we propose a robust training, and evaluation and evolving pipeline to alleviate the issue.
Robust Training. Inspired by the idea of clastic depth [48, 49], we apply stochastic forward propagation to execute each cross-altention and ResNet block by probability p(-, where refers to identity mapping that skips the corresponding block. Thus, we have Eq. (4) becomes as follows:
(t 2.) II [p(Cross Tr(Cross-AAtentioonn.. c), I), p,,
(5)
With this training augmentation, the network is robust to architecture permutations, which enables an accurate assessment each block and stable architectural evolution (more examples in Fig. 5).
Evaluation and Architecture Evolving, We perform online network changes of UNet using the model from robust training with the constructed evolution action set: E (A+.-. { (nuss-nnetio(t..)"" A Resheff.gg where A+.- denotes the action to remove (-) or add (+) a cross-attention or ResNet block at the corresponding position (stage i, block j). Each action evaluated by its impact on execution latency and generative performance. For latency, we use the lookup table built in Sec. 2.2 for each possible configuration of gross-attention and ResNet blocks. Note we improve the UNet for on-device speed; the optimization of model size can be performed similarly and is left as future work. For generative performance, we choose CLIP score [41] to measure the correlation between generated images and the text condition. We use small subset (2K images) of MS COCO validation set [50], fixed steps (50), and CFG scale as 7.5 to benchmark the score, and takes about 2.5 A100 GPU hours to test each action. For simplicity, the value score of each action is defined as ALatency where a block with
ACUP
lower latency and higher contribution to CLIP tends to be preserved, and the opposite is removed in

---

architecture evolving (more details in Alg. 1). To further reduce the cost for network optimization, we perform architecture evolving, i.e,, removing redundant blocks or adding extra blocks at valuable positions by executing group of actions at time. Our training paradigm successfully preserves the performance of pre-trained UNet while tolerating large network permutations (Fig. 5). The details of our final architecture is presented in Sec. A.
3.2 Efficient Image Decoder
For the image decoder, we propose a distilla- tion pipeline that uses synthetic data to learn the efficient image decoder obtained via channel re- duction, which has 3.8x fewer parameters and 3.2x faster than the one from S.-v1.5.. The efficient image decoder is obtained by applying 50% uniform channel pruning to the original image decoder, resulting in compressed ef- ficient image decoder with approximately 1/4 size and MACs of the original one. Here we only train the efficient decoder instead of fol- lowing the training of VAE [4, 38, 39] that also learns the image encoder. We use text prompts to get the latent representation from the UNet of SD-v1.5 1.5 after 50 denoising steps with DDIM and forward it our efficient image decoder and the one of SD-v1.5 to generate two images. We then optimize the decoder by minimizing the mean squared error between the two images. Using synthetic data for distillation brings the advantage of augmenting the dataset on-the-fly where each prompt be used obtain unlimited images by sampling various noises. Quantita- tive analysis of the compressed decoder can be
Algorithm Optimizing UNet Architecture
Require: UNet: éø; validation set: Deati lateney lookup
table T : [CrossAArentiin(i.]] ResNerli, j]}.
Ensure: Eg converges and satisfies latency objective
while ég not converged do
Perform robust training.
Architecture optimization:
if perform architecture evolving at this iteration then
+ Evaluate blocks:
for each block do
ACLIP + eval(Ea, Asock( Beat),
A Latency + eval(Eo. A TT) end
+ Sort actions based on ACUP
Mateney execute ac-
tion, and evolve architecture get latency T:
latency objective S is not satisfied then {Á-) } arg minA ACUP
else
ACLIP {Ā copy(arg max Moteny
êg + evalve(éo, (A})
end if end if
end while
found in Sec. B.2.
4 Step Distillation
Besides proposing the efficient architecture of the diffusion model, we further consider reducing the number of iterative denoising steps for UNet to achieve more speedup. We follow the research direction step distillation [33], where the inference steps are reduced by distilling the teacher, egg, at 32 steps, to a student that runs fewer steps, e.g., 16 steps. This way, the student enjoys 2x speedup against the teacher. Here we employ different distillation pipelines and learning objectives from existing works [33, 32] to improve the image quality, which we elaborate on as follows.
4.1 Overview of Distillation Pipeline
Citing the wisdom from previous studies [33, 32], step distillation works best with the v-prediction type, i.e., UNet outputs velocity v [33] instead of the noise c. Thus, we fine-tune SD-v1.5 1.5 to v- prediction (for notation clarity, we use vo to mean the SD model in v-prediction vs. its e-prediction counterpart éq) before step distillation, with the following original loss Lori:
Corl = E,-U (x-
|/vo Z;, c) - v
(6)
where v is the ground-truth truth target velocity, which can be derived analytically from the clean latent and noise e given time step t: v = a4E o,x.
Our distillation pipeline includes three steps. First, we do step distillation on SD-v1.5 to obtain the UNet with 16 steps that reaches the performance of the 50-step model. Note here we use a 32-step SD-v1.5 1.5 to perform distillation directly, instead of doing il progressively, e.g., using 128-step model as teacher to obtain the 64-step model and redo the distillation progressively. The reason is that we empirically observe that progressive distillation is slightly worse than direct distillation (see Fig. 6(a)
5

---

for details). Second, we use the same strategy to get our 16-step efficient UNet. Finally, we use the 16-step SD-v1.5 as the teacher to conduct step distillation on the efficient UNet that is initialized from its 16-step counterpart. This will give us the 8-step efficient UNet, which is our final UNet model.
4.2 CFG-Aware Step Distillation
We introduce the vanilla step distillation loss first, then elaborate more details on our proposed CFG-aware step distillation (Fig. 3).
Vanilla Step Distillation. Given the UNet inputs, time step t, noisy latent z,, and text embedding the teacher UNet performs two DDIM denoising steps, from time f to t' and then to t" (0 S t" < t' < t S 1). This process can be formulated as (see the Sec. C for detailed derivations).
V = Vol Zz, c) = ze = ar (a - o.V.) + oy (0 + as Vy = vo 20, + Ze" (ay - ve) ) + or (oe + Ve).).
The student UNet, parameterized by 17. performs only one DDIM denoising step,
&. vn z: = x a. v.v.
(8)
where the super-script (s) indicates these variables are for the student UNet. The student UNet is supposed to predict the teacher's 's noisy latent 2,r from 24 with just one denoising step. This goal translates the following vanilla distillation loss objective calculated in the x-space [33, 32],
Zr T Ze
L danii_ds] = = (Ae)
12
Oegir
2 where (A:) max( 1) is the truncated SNR weighting coefficients [33].
CFG-Aware Step Distillation. The above vanilla step distillation can improve the inference speed with no (or only little) FID compromised. However, we do observe the CLIP score turns obviously worse. As remedy, this section introduces classifier-free guidance-aware (CFG-aware) distillation loss objective function, which will be shown to improve the CLIP score significantly.
We propose perform classifier free guidance to both the teacher and student before calculating the loss. Specifically, for Eq. (7) and (8), after obtaining the v-prediction output of UNet, we add the
CFG step. Take Eq. (8) for an example, (x) is replaced with the following guided version,
wv, (t, z,, c) (w - (t, zt. Ø),
(10)
where is the CFG scale. In the experiments, is randomly sampled from uniform distribution over a range ((2, 14) by default) - this range is called CFG range, which will be shown to provide a way to tradeoff FID and CLIP score during training.
After replacing the UNet output with its guided version, all the other procedures remain the same for both the teacher and the student. This gives us counterpart version of Levani dstll which we term CFG distillation loss, denoted as detg
Total Loss Function. Empirically, we find Lvani dstl helps to achieve low FID while Lett_ helps achieve high CLIP score (see Fig. 6(c)). To get the best of both worlds, we introduce loss mixing scheme to use the two losses at the same time- - A predefined CFG probability p is introduced, indicating the probability of using the CFG distillation loss in each training iteration (so with 1 - p probability, the vanilla distillation loss is used). Now, the overall loss can be summarized:
L Lastl + Loris
(11)
Last = Lote if P U10, 1] < else LvanLdsth,
where Cori represents the original denoising loss in Eq. (6) and 7 is its weighting factor; and U represents the uniform distribution over range (0, 1).
1]
Discussion. As far as we know, only one very recent work [32] studies how to distill the guided diffusion models. They propose to distill CFG into a student model with extra parameters (called w--conitiion) to mimic the behavior of CFG. Thus, the network evaluation cost is reduced by 2x when generating image. Our proposed solution here is distinct from theirs [32] for at least four

---

perspectives. (1) The general motivations are different. Their u-condition model intends to reduce the number of network evaluations of UNet, while ours aims to improve the image quality during distillation. (2) The specific proposed techniques are different they integrate the CFG scale as an input to the UNet, which results in more parameters, while we do not. (3) Empirically, w-condition model cannot achieve high CLIP scores when the CFG scale is large (as in Fig. 6(b)), while our method is particularly good at generating samples with high CLIP scores. (4) Notably, the trade-off of iiversityqqualigy is previously enabled only during inference by adjusting the CFG scale, while our scheme now offers a nice property to realize such trade-off during training (see Fig. 6(d)), which w-condition cannot achieve. This can be very useful for model providers to train different models in favor of quality or diversity.
Experiment
Implementation Details. Our code is developed based on diffusers library". Given step distillation mostly conducted on v-prediction models [33, 32], we fine-tune UNet in our experiments to prediction. Similar to SD, we train our models on public datasets [51, 40] to report the quantitative results, iee., FIID and CLIP scores (ViT-g/14), on MS-COCK 2014 validation set [50] for zero-shot evaluation, following the common practice [20, 19, 6, 18]. In addition, we collect an internal dataset with high-resolution images to fine-tune our model for more pleasing visual quality. We use 16 or 32 nodes for most of the training. Each node has 8 NVIDIA A100 GPUs with 40GB or S0GB memory. We use AdamW optimizer [52], set weight decay as 0.01, and apply training batch size as 2, 048.
Stallo Eithation vL1., 50 etope 1. E.iis
Table 2: Zero-shot evaluation on MS-COCO 2017 5K subset. Our efficient model is compared against recent arts in the s-step configura- tion. Note the compared works use
Clert Sittrpe, 1.54
37-
Clare Ie crepl prestacted
Gark shree proipeed doeit
Gare Beeree (Bilde oor rear
Gare ctiqe CBG diell sp-ct mosiie
Glare inpee CaI dialll CPG mocke!
the same model as SD-v1.5, which
22
is much slower than our approach.
Method
Steps
FID
CLIP
DPM [53]
0.32
Figure 4: FID vx. CLIP on MS-COCO 2014 validation set with CFG scale
DPM++ [54]
8
25.6
0.32
from 1.0 to 10.0. Left: Comparison with SD-v 1.5 on onffull set (30K). Right:
Meng et [32]
26.9
0,30
Ours
24.2
0,30
Different settings for step and teacher models tested on 6K samples.
5.1 Text-to-Image Generation
We first show the comparison with SD-v1.5 on the full MS-COCO 2014 validation set [50] with 30K image-caption pairs. As in Fig. 4 (left), thanks to the architecture improvements and the dedicated loss design for step distillation, our final 8-step. 230ms per step UNet outperforms the original SD-v1.5 in terms of the trade-off between FID vs. CLIP. For the most use--preeerable guidance scales (ascending part of the curve), our UNet gives about 0.004 0.010 higher CLIP score under the same FID level. In addition, with an aligned sampling schedule (8 DDIM denoising steps), our method also outperforms the very recent distillation work [32] by 2.7 FID with on-par CLIP score, in Tab. 2. Example synthesized images from our approach are presented in Fig. 1. Our model can generate images from text prompts with high fidelity. More examples are shown in Fig. 9.
We then provide more results for performing step distillation on our efficient UNet. As in Fig. (right), we demonstrate that our 16-step undistilled model provides competitive performance against SD-v1.5. However, we can see a considerable performance drop when the denoising step is reduced to 8. We apply progressive (vanilla) distillation [33, 32] and observe improvements in scores. Though mostly comparable to the SD-v1.5 baseline, the performance of the 8-step model gets saturated for the CL.IP score as the guidance scale increases, and is capped at 0.30. Finally, we use the proposed CFG-aware step distillation and find it consistently boosts the CLIP score of the 8-step model with varied configurations. Under the best-observed configuration (CFG distilled 16-step teacher), our 8-step model is able to surpass SD-v1.5 by 0.002 - 0.007 higher CLIP under similar FID. Discussions on the hyperparameters can be found in ablation studies.
7

---

(h) (c) Remove
(d) Remove CA in mid.
(c) Remove CA in US
(f) Remove RB in DS
(g) Remove
(h) Remove
(a) SD
removes CA
CA DS
RB in mid.
RB in US
Figure 5: Advantages robust training. Prompts of top row: a photo of an astronaut riding a horse mars and bottom row: A pikachu fine dining with view to the Diffel Tower. (a) Images from SD-v1.5. (b) Removing cross-attention (CA) blocks in downsample stage of SD-v1.5. (c) (e) Removing cross-attention (CA) blocks in [downsample (DS), middle (mid.), upsample (US) using our model after robust training. (f) - (h) Removing ResNet blocks (RB) in different stages using our model. The model with robust training maintains reasonable performance after dropping blocks.
5.2 Ablation Analysis
Here we present the key ablation studies for the proposed approach. For faster evaluation, we test the settings on 6K image-caption pairs randomly sampled from the MS-COCO 2014 validation set [50].
Robust Training, As in Fig. 5, we verify the effectiveness of the proposed robust training paradigm. The original model is sensitive to architecture permutations, which makes it difficult to assess the value score of the building blocks (Fig. 5(b)). In contrast, our robust trained model can be evaluated under the actions of architecture evolution, even if multiple blocks are ablated at a time. With the proposed strategy; we preserve the performance of pre-trained SD and save the fine-tuning cost to recover the performance of candidate offspring networks. In addition, we gather some insights into the effect of different building blocks and ensure the architecture permutation is interpretable. Namely, cross-attention is responsible for semantic coherency (Fig. 5(c)-(e)), while ResNet blocks capture local information and are critical the reconstruction of details (Fig. 5(D-(h)), especially in the output upsampling stage.
Step Distillation. We perform comprehensive comparisons for step distillation discussed in Sec. 4. For the following comparisons, we use the same model as SD-v1.5 to study step distillation.
Fig. 6(a) presents the comparison of progressive distillation to 8 steps direct distillation to steps. As seen, direct distillation wins in terms of both FID and CLIP score. Besides, it is procedurally simpler. Thus, we adopt direct distillation in our proposed algorithm.
Fig. 6(b) depicts the results of iu-conditioned models [32] different inference steps. They are obtained through progressive distillation, i.e., 64 + 32 - 16 + 8. As seen, there is clear gap between w-coonditioeed models and the other two, especially in terms of CLIP score. In contrast, our 8-step model can significantly outperform the 50-step SD-v1.5 1.5 in terms of CLIP score and maintain a similar FID. Comparing ours (8-step model) to the w-conditioned 16-step model, one point of particular note is that, these two schemes have the same inference cost, while ours obviously wins terms of both FID and CLIPP score, suggesting that our method offers better solution to distilling CFG guided diffusion models.
Fig. 6(c) shows the effect of our proposed CFG distillation loss vs. the vanilla distillation loss. As seen, the vanilla loss achieves the lowest FID, while the CFG loss achieves the highest CLIP score. To get the best of both worlds, the proposed loss mixing scheme (see "vanilla CFG distill") successfully delivers a better tradeoff: achieves the similar highest CLIP score as the CFG loss alone and the similar lowest FID as the vanilla loss alone.
There are two hyper-parameter in the proposed CFG distillation loss: CFG range and CFG probability. Fig. 6(d) shows the effect of adjusting them. Only using the vanilla loss (the blue line) and only using the CFG loss (the purple line) lay down two extremes. By adjusting the CFG range and probability, we can effectively find solutions in the middle of the two extremes. As rule of thumb, higher CFG probability and larger CFG range will increase the impact of CFG loss, leading to better CLIP score but worse FID. Actually, for the 7 lines listed top to down in the legend, the
8

---

02
clict dedl
pengmeend of (rogge)
yadla mnde
deo sodd jesiig leed
(M0
CR dell
yodd a Cit clett aaps, aed
21
3.25
(a) Direct vs. progressive (b) w-cconditioned Es. ours (c) Vanilla vs. CFG distill (d) CFG hyper-parameterss
Figure 6: Ablation studies in step distillation (best viewed in color). For each line, from left to right,
the CFG scales starts from 1.0 to 10.5 with interval 0.5. (a) To obtain the same 8-step student model,
in direct distillation, the teacher only distills once (16 + 8), while progressive distillation [33, 32] starts from the 64-step teacher, distills 3 times to 8 steps (64 + 32 + 16 + 8). (b) w-conditioned model [32] struggles at achieving high CLIP scores (such as over 0.30) while the original SD-v1.5 and our distilled 8 8-step SD-v1.5 can easily achieve so. (c) Comparison between vanilla distillation
loss Lvani_dsil - the proposed CFG distillation loss Lat dau- and their mixed version Ldst)- (d) Effect of adjusting the two hyper-pane evper-parameters CFG range and CFG probability, in CFG distillation. As seen, these pepeepprreeters can effectively tradeoff FID and CLIP score.
impact of CFG loss is gradually raised, and we observe the corresponding lines move steadily to the upper right, fully in line with our expectation, suggesting these two very reliable way tradeoff FID and CLIP score during training this feature, as far as we know, has not been reported by any previous works.
provide
Analysis of Original Loss for Distillation. In Eq. (11), we apply the original denoising loss (Cori in Eq. (6)) during the step distillation. Here we show more analysis for the using Cori in step distillation.
Fig. 7(a) shows the comparison between using and not using the original loss in our proposed CFG distillation method. To our best knowledge, existing step distillation approaches [33, 32] do not include the original loss in their total loss objectives, which is actually sub-optimal. Our results in Fig. 7(a) suggest that using the original loss can help lower the FII) at no loss of CLIPP score.
Fig. 7(b) provides detailed analysis using different y to balance the original denoising loss and the CFG distillation loss in Eq. (11). We empirically set dynamic gamma to adjust the original loss into a similar scale to step distillation loss.
Analysi for the Number of Inference Steps of the Teacher Model. For the default training setting of the step distillation, the student runs one DDIM step while the teacher runs two steps, e.g., distilling a 16-step teacher to an 8-step student. At the first glance, if the teacher runs more steps, it possibly provides better supervision to the student, e.g., distilling 32-step teacher to the 8-step student. Here we provide empirical results to show that the approach actually does not perform well.
Fig. 7(c) presents the FID and CLIP score plots of different numbers of steps of the teacher model in vanilla step distillation. As seen, these teachers achieve similar lowest FID, while the 16-step teacher (blue line) achieves the best CLIP score. A clear pattern is that the more steps of the teacher model, the worse CLIP score of the student. Based on this empirical evidence, we adopt the 16-step teacher setting in our pipeline to get 8-step models.
Applying Step Distillation to Other Model. Lastly, we conduct the experiments by applying our proposed CFG-aware distillation on SD-v2, where the student model has the same architecture as SD-v2. The results are provided in Fig. 7(d). As can be seen, our 8-step distilled model achieves comparable performance to the 50-step SD-v2 model. We use the same pypper-parmettes from the training of SD-v1.5 for the step distillation of SD-v2, and further tuning might lead better results.
6 Related Work
Recent efforts on text-to-image generation utilize denoising diffusion probabilistic models [55, 35, 56, 1, 2, 4] to improve the synthesis quality by conducting training on the large-scale dataset [40]. However, the deployment of these models requests high-end GPUs for reasonable inference speed due to the tens or hundreds of iterative denoising steps and the huge computation cost of the diffusion

---

12
31-
Den (f (hodens seew borl
Cap oidia ioe)
aad cihed
12
call coall saei
CTC
24
22-
50-2
638
1.31
8.1
622
0.28
620
6.21 100
(a) Cori in Eq. (11).
(b) y in Eq. (1I).
(c) Number of teacher steps. (d) Distillation on SD-v2.
Figure 7: FID and CLIP results on the 6K samples from the MS-COCO 2014 validation set [50] for various models and experimental settings. (a) Comparison between using (red line) and not using (orange line) the original loss in our proposed CFG distillation method. The hyper-parameter setup of "ours" experiments: CFG range [2, 14] and CFG probability 0.1. (b) Analysis of loss scaling y in Eq. (11). Note that we employ dynamic scaling adjust original loss (Cori) into a similar scale of step distillation loss (Cast). We show 7 as 0.01, 0.2, 1.0. Our choice (0.2) gives slightly better FID, despite all dynamic scalings resulting in very similar results. We further show results of constant scaling. Here 0.0 indicates no Lort added, while 1.0 refers to non-scaled Lort where Lori dominates the optimization and degrades the effect of step distillation. (c) Analysis for the number of steps for the teacher model in vanilla step distillation. The student is supposed to run at 8 steps, and we can actually employ different teachers that run different numbers of steps during the step distillation. The default setting in our experiment is that teacher 16 steps, student 8 steps, i.e., the blue line, which turns out to be the best. (d) Results of our proposed CFG-aware step distillation applied on SD-v2.
model. This limitation has spurred interest from both the academic community and industry to optimize the efficiency of diffusion models, with two primary approaches being explored: improving the sampling process [57, 58, 59, 60, 53, 61] and investigating on-device solutions [62].
One promising area for reducing the denoising steps is through progressive distillation, where the sampling steps gradually reduced by distillation that starts from pre-trained teacher [33]. The later work further improves the inference cost of classifier-fre free guidance [34] by introducing the w-condition [32]. Our work follows the path of step distillation while holding significant differences with existing work, which is discussed above (Sec. 4). Another direction studies the methods for optimizing the model runtime on devices [63], such as post-training quantization [22, 23] and GPU- aware optimization [24]. Nonetheless, these works require specific hardware or compiler support. Our work is orthogonal to post optimizations and can be combined with them for further speed up. We target developing a generic and efficient network architecture that can run fast on mobile devices without relying on specific bit width or compiler support. We identify the redundancy in the SD and introduce one with a similar quality while being significantly faster.
7 Discussion and Conclusion
This work proposes the fastest on-device text-to-image model that runs denoising in 1.84 seconds with image quality on par with Stable Diffusion. To build such a model, we propose a series of novel techniques, including analyzing redundancies in the denoising UNet, proposing the woolving-rrainingg framework to obtain the efficient UNet model, and improving the step distillation by introducing the CFG-aware distillation loss. We perform extensive experiments and validate that our model can achieve similar or even better quality compared to Stable Diffusion while being significantly faster.
Limitation. While our approach is able to run the large-scale text-to-image diffusion model on mobile devices with ultra-fast speed, the model still holds a relatively large number of parameters. Another promising direction is to reduce the model size to make it more compatible with various edge devices. Furthermore, most of our latency analysis is conducted on iPhone 14 Pro, which has more computation power than many other phones. How to optimize devices to achieve fast inference speed is also an interesting topic to study.
models for other mobile
Broader Impacts. Similar to existing studies on content generation, our approach must be applied cautiously so that will not be used for malicious applications. Such concerns can also be alleviated by approaches that could automatically detect image content that violates specific regulations.
10

---

References
[1] Jonathan Ho, Ajay Jain, and Pieter Abheel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems. 33:6840-660 6851, 2020. 2.
[2] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint
arXiv:2011.1346,, 2020. 1, 2,
[3] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[4] Robin Rombach, Andreas Blattmann, Dominik Lorenz. Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent difffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10684-10695, 2022. 1, 3, 5,9, 9, 14, 16
[5] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning. pages 8821-8831, PMLR, 2021.
[6] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.066 06125, 2022. 1,7
[7] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorcalistic image generation and editing with ext--uuiedd diffusion models. arXiv preprint arXiv: 1074), 2021. 2
[8] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-inage diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, pages 1-10, 2022.
[9] Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang. Ming- Husu- Yang. Kevin Murphy, William Freeman, Michael Rubinstein, et al. Muse: Text-to-image generation via masked generative transformers. arXiv preprint axXiv:231.007004 2023. 2
[10] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Lue Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11461-11471, 2022. 2
[11] Zhixing Zhang, Ligong Han, Arnb Ghosh, Dimitris N Metaxas, and Jian Ren. Sine: Single image editing with cext-to-image diftusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6027-600 -037, 2023.
[12] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEFF Tronsactions on Pattern Analysis and Machine
Intelligence, 2022. 2
[13] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David Fleet, et al. Imagen video: High definition vidco generation with diffusion models. arXiv preprint arXiv:2210.0230: 2022. 2
[14] Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang 7hang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, et al. Make-a-video: Text-to-video generation without text-viden data. arXiv preprint ar 22v: 14792, 2022. 2
[15] Xiaohui Zeng, Arash Vahdat, Francis Williams, Zan Gojeic, Or Litany, Sanja Fidler, and Karsten Kreis. Latent point diffuusion models for 3d shape generation. arXiv preprint arXiv:2210.0697 2022. 2
Lion:
[16] Chen Hsuan Lin, Jun Gao, Luming Tang. Towaki Takikawa, Xiaohui Zeng. Xun Huang, Karsten Kreis, Sanja Fidler, Ming- Yu Liu, and Tsung-Yr Yi Lin. Magic3d: High-resolution text-to-3d content creation. arXiv preprint ar 2rx: I. 10440, 2022. 2
[17] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint 2209 14988, 2022. 2
[18] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong. Gunjan Baid, Zirui Wang. Vijay Vasudevan, Alexander Ku, Yinfei Yang. Burcu Karagol Ayan, et al. Scaling autoregressive models for content-rich text-to-image generation. arXiv preprint arXiv:2206. 10789, 2022. 2.7
[19] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorcalistic text-to- image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494, 2022. 2. 7, 15
[20] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song. Karsten Kreis, Miik Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, et al. ediffi: Text-to- -omiea diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01324 2022. 2. 7, 15

---

[21] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 4195-4205, 2023.
[22] Xiuyu Li, Long Lian, Yijiang Liu, Huanrui Yang. Zhen Dong, Daniel Kang. Shanghang Zhang, and Kurt Keutzer. Q-diffusion: Quantizing diffusion models. arXiv preprint Stt::2302.043004 2023. 2, 10
[23] Jilei Hou and Ziad Asghar. World's first on-device demonstration of stable diffusion on an android phone,
2023, 2, 10
[24] Yu-Hui Chen, Raman Sarokin, Juhyun Lee, Jiuqiang Tang. Chuo-Ling Chang, Andrei Kulik, and Matthias Grundmann. Speed is all you need: On-device acceleration large diffusion models via gpu-aware optimizations. arXiv preprint arXiv:2304. 1267, 2023. 2, 10
[25] Yuzhang Shang, Zhihang Yuan, Bin Xie, Bingzhe Wu, and Yan Yan. Post-training quantization on diffusion models. arXiv preprint arXiv:2211.1772 2022. 2
[26] Tri Dao. Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficiint exact attention with io-awareness. Advances in Neural Information Processing Systems, 15:16344-16359
2022. 2
[27] Qing Jin, Jian Ren, Oliver J Woodford, Jiazhuo Wang, Geng Yuan, Yanzhi Wang, and Sergey Tulyakow. Teachers do more than teach: Compressing image-to-image models. In Proceedings of the IEEE/CVF
Conference on Computer Vision and Pattern Recognition, pages 13600 13611, 2021. 2
[28] Geng Yuan, Yanyu Li. Sheng Li. Zhenglun Kong. Sergey Tulyakow. Xulong Tang. Yanzhi Wang, and Jian Ren. Layer freezing & data sieving: Missing pieces of generic framework for sparse training, arXiv preprint arXiv:2209. 12044 2022. 2
[29] Yanyu Li. Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, and Jian Ren. Efficientformer: Vision transfomers mobilenet speed. arXiv preprint arXiv:2206.01 191, 2022.
[30] Yanyu Li, Ju Hu, Yang Wen, Georgios Evangelidis, Kamyar Salahi, Yanzhi Wang, Sergey Tulyakov, and Jian Ren. Rethinking vision transformers for mobilenet size and speed. arXiv preprint wxXi:2222.000052 2022. 2,4
[31] Bo- Bo-Kyeong Kim, Hyoung- Kyu Song, Thibault Castells, and Shinkook Choi. On architectural compression of text-to-image diffusion models. arXiv preprint arXiv:2305. 2023. 2
[32] Chenlin Meng, Ruiqi Gao. Diederik P Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. arXiv preprint arXiv:2210.03/4 2022. 2, 5, 6, 7, 8, 9, 10
[33] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00 00512, 2022. 2. 5, 6, 7, 10,
[34] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 2022.
2, 3, 10
[35] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256 2265. PMLR, 2015. 2,9
[36] Olaf Ronneberger, Philipp Fischer, and Thomas Brox, U-net: Convolutional networks for biomedical image segmentation. In MICAI, 2015. 2
[37] Jaaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint mxxi:2010.005000 2020, 3
[38] Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv: 131.:11312.6114 2013. 3.5
[39] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference machine learning, pages 1278-1286. PMLR. 2014. 3,
[40] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Sb: An open large- scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.0844 0800002 2022. 3, 7,9
[41] Alec Radford, Jong Wook Kim, Chris Hallacy. Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-8763. PMLR, 2021. 3, 4
[42] Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang, Learning efficient convolutional networks through network slimming. In ICCV, 2017. 4
[43] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. In ICCV, 2017. 4
12

---

[44] Huan Wang, Can Qin, Yulun Zhang, and Yun Fu, Neural pruning via growing regularization. In ICLR.
2021. 4
[45] Huan Wang and Yun Fu. Trainability preserving neural pruning. In /CLR, 2023.
[46] Barret Zoph and Quoc Le. Neural architecture search with reinforcement learning. In ICLR, 2017. 4
[47] Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: A survey. JMLR,
20(55):1-21, 2019. 4
[48] Gao Huang. Yu Sun. Zhuang Liu, Daniel Sedra, and Kilian Weinberger. Deep networks with stochastic depth. In Computer Vision 2016: 14th European Conference, Amsterdam, The Netherlands, October 11 14, 2016, Proceedings, Part IV 14, pages 646 661. Springer, 2016. 4
[49] Jiahui Linjie Yang. Ning Xu, Jianchao Yang. and Thomas Huang. Slimmable neural networks. arXiv preprint or
2018. 4
[50] Tsung- Yi Lin, Michacl Maire, Serge Belongie, James Hays, Pictro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitmick. Microsoft coco: Common objects in context. In Computer Vision ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part 13, pages 740-755. Springer, 2014. 4, 7, 8, 10, 15
[51] Minwoo Byeon, Beomhee Park, Haccheon Kim, Sungjun Lec, Woonhyuk Back, and Sachoon Kim. Coyo-700m: Image-text pair dataset. https: bttps://github. httpchhbb.ooomeeeeeeeeeeeeeeeeee http//kkabbbaaannooooooaa wwtprraaiicooooaattoooo cat.//toooodddett...ooaoaaaaa.ee 2022.
[52] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015. 7
[53] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.00922,
2022. 7, 10
[54] Cheng Lu, Yuhao Zhou, Fan Bao. Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffiusion probabilistic models. arXiv preprint
2022. 7
[55] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv preprint arXiv:2206. 00364, 2022. 9
[56] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. 9
[57] Alexia Lolicocur-Martincau, Ke Li, Rémi Piché-Taillefer, Tal Kachman, and loannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2 14080, 2021. 10
[58] Eric Luhman and Troy Luhman. Knowledge distillation iterative generative models improved sampling speed. arXiv preprint arXiv:2 101.0288, 2021.
[59] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. arXiv
preprint 22 2022. 10
[60] Luping Liu, Yi Ren, Zhijic Lin, and Zhou Zhao. Pscudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:220..09777. 2022. 10
[61] Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers for diffusion models by differentiating through sample quality. In International Conference on Learning Representations, 2022. 10
[62] Atila Orhon, Michael Siracusa, and Aseem Wadhwa. Stable diffusion with core on apple silicon, 2022.
10
[63] Muyang Li, Ji Lin, Chenlin Meng, Stefano Emoon, Song Han, and Jun-Yan Zhu, Efficient spatially sparse inference for conditional gans and diffusion models, arXiv preprint arXiv:2211.02044 2022.
[64] Tensorr.
https://develog: nvidia. com/ com/tensorrt. 14
13

---

A Efficient UNet
We provide the detailed architecture of our efficient UNet in Tab. 3. We perform denoising diffusion in latent space [4]. Consequently, the input and output resolution for UNet g g which 64 x 64 for generating an image of 512 x 512.
W.
In the main paper, mainly benchmark the lateney on iPhone 14 pro. Here we provide the runtine of the model on more mobile devices in Tab. 4.
In addition to mobile phones, show the latency and memory benchmarks Nvidia A100 40G GPU, as in Tab, 5. We demonstrate that our efficient UNet achieves over 12x speedup compared to the original SD-v1.5 performed via the public TensorRT [64] library in single precision.
server-level GPU and shrinks 46% rumning memory. The analysis is
Table 3: Detailed architecture of our efficient UNet model.
UNet Model
Stage
Resolution
Type
Config
Origin
Ours
Cross Attention
Dimension
320
Blocks
0
Down-1
x 8
ResNet
Dimension
320
Blocks
2
Cross
Dimension
640
Attention
Blocks
2
Down-2
X
W T6
ResNet
Dimension
640
Blocks
2
Cross
Dimension
1280
Attention
Blocks
2
Down-3
H X
W
Dimension
1280
ResNet
Blocks
Cross
Dimension
1280
Attention
Blocks
Mid
W
Dimension
1280
64
ResNet
Blocks
4
Cross
Dimension
1280
Attention
Blocks
6
Up-1
Dimension
1280
32
32
ResNet
Blocks
2
Cross
Dimension
640
Attention
Blocks
3
Up-2
16
x
W
16
ResNet
Dimension
640
# Blocks
3
3
Cross
Dimension
320
Attention
Blocks
0
Up-3
11
x
W
ResNet
Dimension
320
Blocks
3
Table Latency benchmark on iPhone12 Pro Max, iPhone13 Pro Max, and iPhone 14 Pro.
Device
Text Encoder (ms)
UNet (ms)
VAE Decoder (ms)
Overall (s)
iPhone 14 Pro
4.0
230
116
1.96
iPhone13 Pro Max
5.7
315
148
2.67
iPhone 12 Pro Max
6.3
526
187
4.40
Table 5: Latency analysis on Nvidia A100 40G GPU with the TensorRT [64] library, tested with single precision (FP32).
UNet
Batch Size
Latency (ms)
Memory (MB)
Iters
Total Latency (ms)
Speedup
SD-v1.5
2
51.2
6634
50
2,560
Ours
2
26.2
3549
209.6
12.2x
14

---

B Discussions of Text Encoder and VAE Decoder
B.1 Text Encoder
Exiting works have explored the importance of the pre-trained text encoder for generating images [19, 20J. In our work, considering the negligible inference latency (4ms) of the text encoder compared to the UNet and VAE Decoder, we do not compress the text encoder in the released pipeline.
B.2 VAE Decoder
We provide qualitative visualizations and quantitive results of our compressed VAE decoder in Fig. 8. The main paper shows that the image decoder constitutes small portion of inference latency (369ms) compared to the original UNet from SD-v1.5. However, regarding our optimized pipeline (230ms x 8 steps), the decoder consumes a considerable portion of overall latency. We propose an effective distillation paradigm to compress the VAE decoder. Specifically, we obtain the latent-image pairs by forwarding the text prompts into the original SD-v1.5 model. The student, which the compressed decoder, takes the latent from the teacher model as input and generates an output image that is optimized with the ones from the teacher model by the mean squared error. Our proposed method wields the following advantages. First, our approach does not demand paired text-image samples, and can generate unlimited data on-they-fly, benefiting Second, the distillation paradigm is simple and straightforward, requiring minimal implementation efforts compared to conventional VAE training. As in Fig. 8, our compressed decoder (116ms) provides comparable generative quality, and the performance degradation compared to the original VAE decoder is negligible.
generalization of the compressed decoder.
Siarr
S0v110.4
Ours, stice
Ohrs, compresed óvoree
28
26
32
20
0.21
0.25
CLDP ac00s
Original (top) vs. compressed (bottom) decoder.
(b) FTD Vs. CL.IP score.
Figure 8: Evaluation using MS-COCO 2014 validation set [50]. (a) Generated images by using the
decoder from SD-v1.5 and our compressed image decoder. The UNet is our efficient UNet, and the guidance scale for CFG is 9.0. (b) Quantitative comparison on the 6K samples. Our compressed decoder performs similarly to the original one considering the widely used CFG scale, i.e., from to 9, and still performs better than the SD-v1.5.
C Detailed Derivations of Step Distillation
The following are the detailed derivations of Eq. (7) Eq. (9) in the main paper.
Given the UNet inputs, time step I, noisy latent 2,, and text embedding c, the teacher UNet performs nwo DDIM denoising steps, from time t to t' and then to t" (0 < t" < <t 1).
We first examine the process from f', which can be formulated as,
Vi = vo (t, Zi, c)
D> Teacher UNet first forward
3 & = a:2: oiVe,
>> Teacher predicted clean latent at time L
é = oi + V.ve,
I> Teacher predicted noise at time t
(12)
zy = a + arg.
b Teacher predicted noisy latent at time t!
= ar (a,2, v.) + oe (a + a,))).
15

---

The process from t' to t" can be derived just like the above, by replacing and t' with t' and t", respectively:
ve velt', 2P, D- Teacher UNet second forward
= Ey agby or VN, + Teacher predicted clean latent at time i'
Ey BN aNy Teacher predicted noise time t
(13)
2r" An oF Teacher predicted noisy latent at time t" (ayzz og ar" (azzy + arvy).
The student UNet, parameterized by 7) performs only one DDIM denoising step,
((0) 2t, Student UNet forward
(1) aiz -ov((") Student predicted clean latent at time
(6) (2. a.x; (0) /or, & Student predicted noise at time
(8) x + o
> Student predicted noisy latent at time t"
pp
x 0e" (z: x
(14)
= (a,"
au )x,
(a) +
de
x/
ai
0gd
where the super-script (s) indicates these variables are for the student UNet. The student UNet is supposed predict the noisy latent z from Ze of the teacher with just one denoising step, namely,
Zr
(15)
Replacing (t) with z in the final equation of Eq. (14), we arrive at the following loss objective,
Evani detl (A.) || x
Zyer
d de
Ze
2:
(16)
Ctter de Cte
where (At)) = max( o 1) is the truncated SNR weighting coefficients [33].
D Different Teacher Options for Step Distillation
is con-trivial decide the best teacher model distill our final 8-step efficient UNet. Fig. we conduct several experiments to explore different teacher options. As straightforward choices, self- distillation from our 16-step efficient UNet or distillation from the 16-step SD-v1.5 baseline model can effectively boost the performance of our 8-step model. Additionally, we investigate whether stronger teachers can further boost performance by training a CFG-aware distilled 16-step SD-v1.5 model, discussed in Sec. 4. We obtain significant improvements in CLIP scores, demonstrating the potential employing better teacher models. We would like to mention that we also experiment with SD-v2 the teacher model. Surprisingly, the different text embeddings used SD-v1. 1.5 and SD-v2 pipelines. Distillation between different infrastructures might be possible future direction explore.
observe much worse results. We attribute this to
E
Additional Qualitative Results
We provide more generated images from our text-to-image diffusion model in Fig. 9. As an accelera- tion work for generic Stable Diffusion [4], our efficient model demonstrates sufficient capability to synthesize various contents with high aesthetics, such as realistic objects (food, animals), scenery, and artistic and cartoon styles.
16

---

Sonor ddld mace bobdh pozdoar 20
Fog Elarason
dictar Al
Seauyy puneza paackg and cbtss excreeing
lage plas afficiices grewe dipin cualeds otvertingg preptee at
Swadik manball
Abaatifal auge calt anine fer deboni beaaty
Goad Tasa Landeape 30
(kelggack of raokers accredd jingle anl smmun Bowing kide tegyy atncaphric
The clench lanlecee siB lantay doakig vullagee. paina deall daip
undogiand Sinary viller dus
Waiiigg in tuffe palnting
bonase ran, ietmettcc beuehlalnggs, concipt ant Magl My
bouttdd topamatiat lark Donuns old lie Caniie lightng
2045 suis sates city traducal, dasted ansect. loper aaliti aolers anvanvaly
houattd picnss diruat
Aplanooia auveea r68ag
doan, kelicas dradl ingeadan
unmod lighting volauaticc clado spperdalee beading
Deoupt et dagd pilbol
fig chil lightiightng dipar higoo doalled mnedds
coaa chooe endked dacmmii habiito cranie.
ideecou Be mook Ariza
plansee ce weeds tbiic dasal
acceutic bunding austaso. body unrvek
fick bacck anmxephi
fass, kyper
Figure 9: Example generated images by using our efficient text-to-image diffusion model.
17


Reviews:
[{"review_id": "hHOEDq0lBK", "reviewer": "Reviewer_5sgC", "summary": "Text-to-image generation (T2I) is getting popular and has a vast application value. This paper explores efficient T2I by reducing computational redundancy. They learn an efficient U-Net via data distillation and then decrease the required diffusion steps. They can achieve T2I on an iPhone 14 Pro in 2 seconds.", "questions": "Please see the Weakness", "limitations": "N/A", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "+ This paper is well-written and easy to follow.\n+ The efficiency of T2I models is less explored in previous studies. The proposed efficient U-Net is valuable for future research and can be practical for real-world usage.\n+ They provide an attractive demo video in their supplementary.", "weaknesses": "+ Why use ViT to encode the input prompt (Fig. 3)? From my best knowledge, ViT is for encoding images, and StableDiffusion borrows the text encoder (similar to BERT) from CLIP.\n+ Since they aim at T2I on mobile devices, various devices should be compared/discussed. iPhone 14 Pro is currently one of the most powerful mobile phones. How about those mid-level phones? Or older versions of the iPhone. This study can make this paper more robust.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for the positive feedback and thoughtful comments. We appreciate the reviewer's acknowledgment that the paper proposes a valuable efficient U-Net for research and practical usage, is easy to follow, and provides an attractive demo video.**\n\n***\n**Q1. About using ViT for encoding the input prompt in Fig.3.**\n\nHere we follow the model notation from CLIP [a] that uses ViT to represent both image encoder and text encoder. In our model, we only use the text encoder from ViT to get the textual embedding from input prompts. Thank you for the suggestion and we will add more details in the revised version to better clarify the implementation.\n***\n**Q2. Latency on more devices.**\n\nThank you for the suggestion. We conduct latency analysis on more devices, such as iPhone12 Pro Max and iPhone13 Pro Max, and will add these results in the revision. We use the latency benchmark tool from Xcode to get the reproducible runtime on different devices.\n\n>|      Device      | Text Encoder (ms) | UNet (ms) | VAE Decoder (ms) | Overall (s) |\n>|----------------|:-----------------:|:---------:|:----------------:|:-----------:|\n>|   iPhone14 Pro   |        4.0        |    230    |        116       |     1.96    |\n>| iPhone13 Pro Max |        5.7        |    315    |        148       |     2.67    |\n>| iPhone12 Pro Max |        6.3        |    526    |        187       |     4.40    |\n\n***\nReferences:\n\n[a] Radford, Alec, et al. Learning transferable visual models from natural language supervision. ICML, 2021.\n"}, {"review_id": "6sZDKP5swJ", "reviewer": "Reviewer_vQqa", "summary": "This paper introduces a novel framework that significantly reduces the inference speed of text-to-image diffusion models to under two seconds on mobile devices. The authors propose two key techniques. Firstly, they introduce Efficient U-Net which enhances inference speed by eliminating redundant architectural components without compromising performance. Secondly, the authors present a novel step distillation method that incorporates a classifier-free guidance loss term, which further improves CLIP score. Extensive experiments are conducted on the MS-COCO dataset to validate the effectiveness of the proposed approach. \n", "questions": "- In Fig 5, can authors provide the generated images under SD without ResBlocks (before robust training)?\n- Just out of curiosity, did authors conduct experiments on transformer-based diffusion models (e.g., DiT [1]) as well?\n\n[1] Peebles, William, and Saining Xie. (2022) \"Scalable diffusion models with transformers.\"", "limitations": "The paper provided limitations along with future research directions in Section 7. \n", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper tackles a timely and practically-relevant problem supported by a fair amount of experiments, and stands as the pioneering study in attempting to reduce the latency of text-to-diffusion models on mobile devices.\n- The paper covers a good amount of relevant previous studies.", "weaknesses": "- In lines 232-233, the authors mention the use of an additional dataset for the fine-tuning purpose. It is crucial for the authors to provide further clarification regarding the fair comparison issue in relation to this additional dataset.\n- The proposed distillation scheme appears to involve a series of intricate choices, making it less readily applicable to other settings. For instance, can this distillation scheme be applied to SD-v2.0?\n- The clarity of Algorithm 1 should be improved. For instance, the condition “T is not satisfied” and $\\hat{A}$ needs clarification. Also, please write the details regarding removing a group of actions in line 136. \n- The readability of Section 4.1 can be enhanced, particularly by providing an explicit distillation algorithm. Further, the reason as to how step distillation and cfg loss trade off FID/CLIP score is unclear.\n- While the image generation speed has been enhanced, the model still carries a memory burden.\n\nI am willing to raise my score if above concerns are properly addressed during rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for the positive feedback and valuable suggestions. We appreciate that the reviewer acknowledges our paper introduces a novel framework with efficient Unet and novel step distillation to significantly reduce the inference speed of the text-to-image models, and stands as the pioneering study. We also thank the reviewer for agreeing the problem studied in this paper is timely and practically relevant with extensive experiments. We address the main questions in the following.**\n\n***\n**Q1. About the additional dataset.**\n\nWe report all the quantitative experiments in the paper by only using public datasets for fair comparisons with existing works, as mentioned in Lines 230-232.\n\nSome qualitative images in the paper, such as Fig. 1 in the main paper, are generated by using the model fine-tuned on additional dataset, as mentioned on Lines 232-233. Fine-tuning our model on the additional dataset is optional and actually achieves very similar quantitative results as using public datasets, while sometimes generating more visually-pleasing images. \n\n***\n**Q2. Distillation scheme on Stable Diffusion V2.**\n\nThanks for the suggestion. Our distillation pipeline is a generic approach without relying on specific models. We conduct the distillation experiments on SD-v2 and plot the FID and CLIP scores in the author response PDF (Fig.1.a). As we can see, our 8-step distilled model achieves comparable performance to the 50-step SD-v2 model. Please note that we use the exact same hyper-parameters from the training of SD-v1.5 for the step distillation of SD-v2, and further tuning might lead to better results.\n\n***\n**Q3. Clarification of Algorithm 1.**\n\nThanks for your suggestion and we will improve the clarity of Algorithm.1. $\\hat A$ refers to the desired action indicated by the scoring metric defined in Sec.3.1. Here \"$T$ is not satisfied\" should be \"latency objective $S$ is not satisfied\". We are sorry for the notation typo and will fix it in the revision. \n\nFor architecture optimization, the model is initialized from pretrained SD-v1.5, and only 'removing' actions ($\\hat A^-$ ) are executed, looping until target latency $S$ is satisfied. When the target latency is achieved, the evolving process comes to a vibrate equilibrium. If the latency falls below the objective, the 'adding' action ($\\hat A^+$ ) is executed, as in Line 158 in Alg.1. \n\nFor both actions ($\\hat A^-$ and $\\hat A^+$), we add or remove a group of blocks at a time to save evolving and validation costs. For instance, we remove 2 attention blocks in the first Downsampling stage at a time when these blocks are identified as less valuable (slow, contribute less CLIP score). \n\n***\n**Q4. About the detailed distillation algorithm in Sec. 4.1.**\n\nThanks for the suggestion. We provide an explicit step distillation algorithm in the author response PDF (Alg. 1) and will also include it in the revision. \n\n***\n**Q5. How step distillation and cfg loss trade off FID/CLIP score.**\n\nUsing our CFG loss for the step distillation, the student model has the probability to mimic two CFG steps of the teacher model with one CFG step. Notice that the CFG has the behavior of using the guidance scale to trade off the FID and CLIP score (as in Fig. 4 Left). Therefore, during the step distillation, the sampled guidance range of the CFG loss and the probability of using CFG loss can influence the FID/CLIP trade-off of the student model, which can be attributed to the FID/CLIP trade-off of the teacher model.   \n\nAs we can see in Fig. 6(d), using a small range of the CFG guidance scale, such as only 3 as the red line, the student CLIP score can not be improved, since the teacher model only provides supervision of good FID but weak CLIP score under CFG scale as 3. On the other hand, by only using the CFG loss without the vanilla distillation loss, as in the purple line, the student model has a weak FID score, since the teacher model has the behavior of overall good CLIP but weak FID score. Therefore, we use the probability of CFG loss and the range of CFG guidance scale to trade-off the FID/CLIP information from the teacher model, which will then be learned by the student model.\n\nThanks for the suggestion and we will include the discussion in the revised paper.\n\n***\n**Q6. About model memory.** \n\nThank you for raising the discussion. As mentioned in the Limitation section (Line 323), our current scope is the speed optimization of text-to-image models (within two seconds of runtime on mobile devices), and we leave storage size and running memory optimizations as future work. In Tab.6 in Appendix, we can still observe that our efficient model runs with almost 47% less memory on NVIDIA A100 with TensorRT than SD-v1.5. We further obtain the running memory on iPhone 13 Pro Max (iOS 16.5.1), benchmarked by Xcode. Our model requires 3.2G NPU memory (out of a total of 5.5G NPU memory) when using our demo App, while we can not get a consistent estimation of the original SD-v1.5 due to out-of-memory issues on NPU.\n\n***\n**Q7. Generated images under SD without ResBlocks.**\n\nThanks for the suggestion. We provide the generated images of SD-v1.5 without ResBlocks before robust training in the author response PDF (Fig.1.b). We will also include them in the revised paper.\n\n***\n**Q8. About experiments on DiT [a].**\n\nThanks for suggesting DiT, which is a relevant and great work, and we will discuss it in the revision. In our work, we focus on accelerating large-scale text-to-image models on mobile devices. As for now, DiT only has the model trained on ImageNet. Training DiT on large-scale datasets such as LAION for text-to-image generation requires re-designing the model architecture of DiT to incorporate text condition and many computational resources. We are interested in incorporating architecture improvements from DiT in our future work.       \n\n***\nReferences:\n\n[a] Peebles, William, and Saining Xie. Scalable diffusion models with transformers. arXiv, 2022.\n\n"}, {"review_id": "dbdANvcskB", "reviewer": "Reviewer_3k2m", "summary": "This work develops a lightweight Stable Diffusion model with architecture compression and step reduction. For architectural compression, an efficient UNet architecture is obtained from evaluating the importance of individual residual and attention blocks, and an efficient image decoder is obtained via channel reduction and conventional knowledge distillation. For step reduction, a distillation that applies classifier-free guidance (CFG) during the training phase is introduced. Unlike previous methods that apply CGF during the inference, the proposed approach achieves a better tradeoff between FID and CLIP score.\n", "questions": "Please check the above weakness section.", "limitations": "n/a", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* This work successfully compresses Stable Diffusion, one of the most famous foundation models, and enables their deployment on mobile devices with a two-second latency. I think this study can receive significant attention in both academia and industry fields.\n* The results of the proposed architectural compression are impressive, although some questions regarding training computations and resources exist (described below).\n* The improved step distillation is well-motivated and novel. I think this can be broadly applicable to diffusion-based large models beyond Stable Diffusion.\n* I appreciate Figure 2 and Table 1, which effectively explain the compute cost of SD. In particular, Figure 2 is very informative and well illustrates architectural bottlenecks in parameters (inner stages) and latency (outer stages).\n* The paper is very well-written and easy-to-follow. The terms and equations are properly described. \n", "weaknesses": "* According to the implementation details, this work seems to use the following datasets and compute machines. The training cost feels VERY HUGE, in contrast to the authors’ initial claim at the introductory paragraph of Section 3 (i.e., pruning and architecture search require significant training compute, whereas the authors propose some methods to alleviate the issue). It seems like a training cost comparable to building the original SD model from scratch.\n    * [49] Coyo-700m: Image-text pair dataset. https://github.com/kakaobrain/coyo-dataset, 2022.\n    * [38] Laion-5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08402, 2022.\n    * [+] an internal dataset with high-resolution images to fine-tune our model for more pleasing visual quality.\n    *  At least, 16 nodes x 8 GPUs per node = 128 A100 GPUs (We use 16 or 32 nodes for most of the training. Each node has 8 NVIDIA A100 GPUs with 40GB or 80GB memory)\n\n* Would you clarify the number of training pairs AND the training hours used for each of (1) Robust Training in Section 3.1, (2) Image Decoder Compression in Section 3.2, and (3) Step Distillation in Section 4? I think the detailed description of training data/hours is quite important for readers and future research, especially under the popularity of large foundation models.\n\n* Furthermore, are there any supporting materials or results to argue “From our empirical observation, the operator changes resulting from network pruning or searching lead to degraded synthesized images, asking for significant training costs to recover the performance.” at the introductory paragraph of Section 3? A recent, concurrent work* claims that it would be possible to obtain a small pruned UNet in SD with very small retraining cost.\n    * On Architectural Compression of Text-to-Image Diffusion Models, https://arxiv.org/abs/2305.15798\n\n* Is there any overlap between the data used for architecture evolving (a smaller subset (2K images) of MS-COCO validation set [48] in line 130 at page 4) and the zero-shot MS-COCO set for final evaluation? If yes, it cannot be considered as the zero-shot evaluation, because the development data for selecting proper blocks were already used and seen. Furthermore, I was wondering the type and quantity of data for architecture evolving would matter and affect the performance.\n\n* Robust Training in Section 3.1 - Would you clarify whether random block dropping is applied to training from scratch (i.e., random initialization) or to retraining from the pretrained SD-v1.5?\n\n* It would be better to describe the method details about compressing the image decoder. Could you describe the channel reduction and the type of synthetic data in detail? I was not able to find them although I have checked the supplementary material (e.g., Supple Section 2.2 VAE Decoder).\n\n* It would be good to analyze the impact of gamma in Eqn (11), the loss weighting between the vanilla step distillation loss and the proposed CFG-aware distillation loss. How did the authors set this parameter? How did different gamma values affect the generation performance? I have checked Figure (6)-b to analyze the effect of CFG range of w in Eqn (10) and the CFG probability p of the loss mixing in Eqn (11), but the effect of gamma in Eqn (11) is not explained.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for the positive feedback and thoughtful comments. We appreciate that the reviewer acknowledges that our successfully compressed two-second text-to-image model can achieve significant attention in academia and industry, our step distillation is novel and can be broadly applied to models beyond Stable Diffusion, our results are informative and impressive, and our paper is well-written.**\n\n***\n**Q1. Training dataset and detailed training cost per each stage.**\n\nThanks for raising the question. In our training pipelines, we use 189M text-image pairs from Laion-5b with the shortest edge as 1024, and 194M text-image pairs from Coyo-700M with the shortest edge as 512.  As mentioned in Lines 230-232, we report all the quantitative experiments in the paper by only using the above public datasets for fair comparisons with existing works. The data we used only accounts for 6.8% of the data in Laion-5b and Coyo-700M.\n\nSome qualitative images, such as Fig. 1 in paper, are generated by using the model fine-tuned on additional dataset, as in Lines 232-233, which has 160M text-image pairs.  Fine-tuning our model on the additional dataset is optional and achieves very similar quantitative results as public datasets, while sometimes generating more visually-pleasing images. \n\nWe provide detailed training costs in the following table, including the cost of robust training, efficient UNet fine-tuning, step distillation, and decoder compression. Overall, our model training only requires 8.6% training samples in comparison with the training of SD-v1.5 from scratch. Note here we report training iterations and total training samples because the GPU-hour estimation is highly dependent on GPU cluster specs, especially inter-node bandwidth, and thus inaccurate.  Based on public information (Huggingface), SD-v1.5 is trained with 30-60 days on 32 nodes (8 GPUs per node), while replicating our entire workflow takes $\\sim 4$ days on our 32-node cluster. \n\n\n>|           Stage           | A100 GPUs | Batch size | 256x256 iters (K) | 512x512 iters (K) | Total training samples (M) | \n>|-------------------------|:---------:|:----------:|:-----------------:|:-----------------:|:--------------------------:|\n>|      Robust training      |    128    |      4     |         -         |         25        |            12.8            |  \n>| Efficient UNet finetuning |    256    |      8     |         -         |         55        |            112.6           |  \n>|     Step distillation     |    128    |      4     |         -         |         10        |             5.1            |  \n>|    Decoder compression    |       96    |    6        |         -         |         10          |           5.7                 |  \n>|   SD-v1.5 (from scratch)  |    256    |      4     |        237        |        1304       |           1578.0           |\n\n***\n**Q2. Discussions on pruning cost and concurrent work.** \n\nWe tried to apply simple channel pruning on UNet. Under similar training time as robust training and efficient UNet fine-tuning (more than 80K iterations), the model is still hard to recover ($\\sim 0.02$ drop in CLIP score). More advanced channel pruning methods are required.\n\nThanks for mentioning the concurrent work [a], which is an interesting and inspiring work! We notice three major differences between our work and [a] that might cause the differences in training cost. First, we do step distillation to reduce the inference time, while [a] does not. Second, we use the automatic approach for finding and adding blocks, while [a] uses a manually designed pipeline. The cost of manual labor is not clear. Third, and most importantly, we aim to achieve a similar or even better FID/CLIP score than the Stable Diffusion models, while [a] performs slightly worse than SD baseline (FID degrades by $2.71\\sim 4.07$, CLIP degrades by $0.008\\sim 0.03$). We will cite and discuss [a] in the revised paper.\n\n***\n**Q3. About the COCO 2K validation data for architecture evolving.** \n\nThere is no overlap between the 2K data for architecture validation and the data used for our evaluation on MS-COCO (6K or 30K). \n\nAs suggested, we show the results of using different data, i.e., 2K subset from MS-COCO and 2K subset from LAION, to validate actions in architecture evolving. As shown in the response PDF (Fig.2.a), using different data leads to very similar conclusions ($\\Delta  CLIP$) for each block. We show stability regarding the validation set selection and strong correlations of the block importance estimation. For instance, both COCO and LAION validation sets suggest the importance of Up.2.Attention, i.e., we should apply more attention modules in Up.2 stage. \n\n***\n**Q4. About robust training.**\n\nRobust training in Section 3.1 is applied to pre-trained SD-v1.5. \n\n***\n**Q5. Compression details of image decoder.** \n\nWe apply 50% uniform channel pruning to the image decoder, resulting in a compressed efficient image decoder with approximately $\\frac{1}{4}$ size and MACs. \n\nIn order to train the efficient image decoder, we use the dataset that only includes the text prompts from the LAION dataset. For each text prompt, we get the latent representation from the UNet of SD-v1.5 after 50 denoising steps with DDIM. We then forward the latent to our efficient image decoder and the decoder from SD-v1.5, to generate two outputs. Then we distill the efficient decoder by minimizing the mean squared error between the two outputs. We provide the details on Line 148 and will improve it with more details in the revision.\n\n***\n**Q6. About gamma in Eq.11.**\n\nThank you for your valuable comment. We empirically set a dynamic gamma to adjust the original loss into a similar scale to step distillation loss. We provide a detailed ablation analysis of different scaling strategies in the response PDF (Fig.2.b) and will include it in the revision. \n\n***\nReferences:\n\n[a] Kim, et al., On Architectural Compression of Text-to-Image Diffusion Models, 2023.\n"}, {"review_id": "wlFUHq43KF", "reviewer": "Reviewer_oXj2", "summary": "This paper propose a system of model compression methods for text-to-image models such as stable diffusion. Authors propose to apply robust training (i.e. stochastic depth training) to train a network robust to architecture change and propose CFG-free knowledge distillation during the compression of the u-net. Besides, author also give in-depth study on the computation cost of different layers.", "questions": "Please refer to the weakness.", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. This paper focus on a hot and an interesting topic - compression of large generative models.\n2. This paper have good evalution and good experimental results.\n3. Good writing.\n3. I like the video in the supplementary material.", "weaknesses": "1. The main model compression methods are borrowed from previous methods. Robust training is totally same as stochastic depth training. Their evaluation methods for architecture search is also well used in NAS. In the knowledge distillation part, simply adding knowledge distillation loss CFG guidance does not show enough novelty.\n2. About the experimental results. Their have been many KD methods for stable diffusion (e.g. On Distillation of Guided Diffusion Models). No enough comparison is provided. Besides, although the appendix gives some qualitative results, it still will be better to provide more.\n3. One of the main obstacles in compression of large-scale generative models is the overlarge training cost. Although authors provide their training settings in the manuscript,  it will be better to provide more information about the training costs of robust training and the architecture optimization. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for the positive feedback and valuable comments. We appreciate that the reviewer acknowledges our paper studies a hot and interesting topic, and includes good evaluation and experimental results. We are glad to know the reviewer likes the video in the supplementary material. We address the raised concerns as follows.**\n\n***\n**Q1. About novelty.** \n\nWe humbly think our proposed architectural evolving and CFG distillation are both novel. First, we agree that the architectural evolving is inspired by the elastic depth, which we have cited on Line 117. However, to the best of our knowledge, we are the first to study how to efficiently optimize the architecture of large-scale text-to-image diffusion models, which is rarely studied before. The process requires careful design and evaluation metrics. Second, we comprehensively study the effects of step distillation and propose a new CFG distillation that can better trade off the FID/CLIP score than existing works. Therefore, we think our proposed approaches are valuable and novel, which is also agreed by Review 5sgC, Reviewer vQqa, and Reviewer 3k2m.\n\n***\n**Q2. About comparisons to KD methods for stable diffusion, and more qualitative results.** \n\nWe carefully reviewed the related literature and there are two works [a,b] that are highly related to our step distillation. The two works use progressive step distillation [a, b], and [b] proposes the w-conditioned model. In the paper, we have comprehensively compared our results with previous works [a,b]. For example, in Fig. 6(a), we compare our direction distillation with progressive distillation. In Fig. 6 (b), we compare our results with the w-conditioning method proposed in [b]. The extensive comparison demonstrates that the strategy for our step distillation is better than existing works. Additionally, in Tab. 2 of the main paper, we further compare our efficient UNet with the model in [b], which has a similar architecture as SD-v1.5, and demonstrate that our model can achieve better FID under the same DDIM scheduler. Additionally, we notice that in Fig.10 in [b], the distilled model is capped with less than 0.30 CLIP score, showing degraded text-image alignment, while our 8-step model overcomes this issue and achieves even better CLIP than SD-v1.5 (approaching 0.31). \n\nFor the qualitative results, we provide images in Fig. 1 of the main paper, Fig. 7 and Fig. 10 in the supplementary file, and in the demo video. We further provide more generated images in the author response PDF (Fig. 3).\n\n***\n**Q3. About training cost.** \n\nThanks for the question about training costs. We provide more details in the following table, including the cost for robust training, efficient UNet fine-tuning, step distillation, and decoder compression. Overall, our model training only requires 8.6% training samples in comparison with the training of SD-v1.5 from scratch. Note here we report training iterations and total training samples because the GPU hour estimation is dependent on GPU cluster specs, especially inter-node bandwidth, and thus inaccurate. Based on public information (by Huggingface), SD-v1.5 is trained with 30-60 days on 32 nodes, while replicating our entire workflow takes about 4 days on our 32-node cluster.\n\n\n>|           Stage           | A100 GPUs | Batch size | 256x256 iters (K) | 512x512 iters (K) | Total training samples (M) | \n>|-------------------------|:---------:|:----------:|:-----------------:|:-----------------:|:--------------------------:|\n>|      Robust training      |    128    |      4     |         -         |         25        |            12.8            |  \n>| Efficient UNet finetuning |    256    |      8     |         -         |         55        |            112.6           |  \n>|     Step distillation     |    128    |      4     |         -         |         10        |             5.1            |  \n>|    Decoder compression    |       96    |    6        |         -         |         10          |           5.7                 |  \n>|   SD-v1.5 (from scratch)  |    256    |      4     |        237        |        1304       |           1578.0           |\n\n***\nReferences:\n\n[a] Salimans, Tim, et al., Progressive distillation for fast sampling of diffusion models. ICLR, 2022.\n\n[b] Meng, Chenlin, et al., On distillation of guided diffusion models. CVPR, 2023. \n"}]

Author rebuttal:
Dear Reviewers,

Thank you all for your positive rating and valuable comments. We upload a one-page PDF to include additional figures and algorithms. All other questions are addressed in the following individual responses. We sincerely look forward to having further discussions with you during Reviewer-Author Discussions if there are any other questions.

Thanks,

Authors


Meta-review: All the reviewers recommend the acceptance of the work. Reviewers appreciated the timely work on compressing large generative models and good results. Reviewers raised some clarification questions and also raised concerns regarding the large training time. Authors addressed several of these questions in their responses. The reviewers did raise some valuable concerns and suggestions that should be addressed in the final camera-ready version of the paper, which include adding the relevant rebuttal discussions and revisions in the main paper. The authors are encouraged to make the necessary changes to the best of their ability.
Decision: Accept (poster)