SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $w$  is the CFG scale. In the experiments,  $w$  is randomly sampled from a 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 a counterpart version of  $\mathcal{L}_{\mathrm{vani\_dst1}}$  - which we term CFG distillation loss, denoted as  $\mathcal{L}_{\mathrm{cfg\_dst1}}$ .
Total Loss Function. Empirically, we find  $\mathcal{L}_{\mathrm{vani\_dst1}}$  helps to achieve low FID while  $\mathcal{L}_{\mathrm{cfg\_dst1}}$  helps to achieve high CLIP score (see Fig. 6(c)). To get the best of both worlds, we introduce a 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:
$$
\mathcal {L} = \mathcal {L} _ {\mathrm {d s t l}} + \gamma \mathcal {L} _ {\mathrm {o r i}},
$$
$$
\mathcal {L} _ {\mathrm {d s t l}} = \mathcal {L} _ {\text {c f g} - \mathrm {d s t l}} \text {i f} P \sim U [ 0, 1 ] <   p \text {e l s e} \mathcal {L} _ {\text {v a n i} - \mathrm {d s t l}}, \tag {11}
$$
where  $\mathcal{L}_{\mathrm{ori}}$  represents the original denoising loss in Eq. (6) and  $\gamma$  is its weighting factor; and  $U[0,1]$  represents the uniform distribution over range  $(0,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$ -condition) to mimic the behavior of CFG. Thus, the network evaluation cost is reduced by  $2 \times w$  when generating an image. Our proposed solution here is distinct from theirs [32] for at least four
perspectives. (1) The general motivations are different. Their  $w$ -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 diversity-quality 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.
Implementation Details. Our code is developed based on diffusers library<sup>3</sup>. Given step distillation is mostly conducted on v-prediction models [33, 32], we fine-tune UNet in our experiments to v-prediction. Similar to SD, we train our models on public datasets [51, 40] to report the quantitative results, i.e., FID and CLIP scores (ViT-g/14), on MS-COCO 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 80GB memory. We use AdamW optimizer [52], set weight decay as 0.01, and apply training batch size as 2, 048.
Table 2: Zero-shot evaluation on MS-COCO 2017 5K subset. Our efficient model is compared against recent arts in the 8-step configuration. Note the compared works use the same model as SD-v1.5, which is much slower than our approach.
Method	Steps	FID	CLIP
DPM [53]	8	31.7	0.32
DPM++ [54]	8	25.6	0.32
Meng et al. [32]	8	26.9	0.30
Ours	8	24.2	0.30
Figure 4: FID vs. CLIP on MS-COCO 2014 validation set with CFG scale from 1.0 to 10.0. Left: Comparison with SD-v1.5 on full set (30K). Right: 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 user-preferable 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, as 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. 4 (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 CLIP 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.
Figure 5: Advantages of robust training. Prompts of top row: a photo of an astronaut riding a horse on mars and bottom row: A pikachu fine dining with a view to the Eiffel 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 to the reconstruction of details (Fig. 5(f)-(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 vs. direct distillation to 8 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  $w$ -conditioned models [32] at different inference steps. They are obtained through progressive distillation, i.e.,  $64 \rightarrow 32 \rightarrow 16 \rightarrow 8$ . As seen, there is a clear gap between  $w$ -conditioned 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 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 in terms of both FID and CLIP score, suggesting that our method offers a 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: it 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-parameters 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 a 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
(a) Direct vs. progressive
(b)  $w$  -conditioned vs. ours
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 \to 8$ ), while progressive distillation [33, 32] starts from the 64-step teacher, distills 3 times to 8 steps ( $64 \to 32 \to 16 \to 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-step SD-v1.5 can easily achieve so. (c) Comparison between vanilla distillation loss  $\mathcal{L}_{\mathrm{vani\_dstd}}$ , the proposed CFG distillation loss  $\mathcal{L}_{\mathrm{cfg\_dstd}}$ , and their mixed version  $\mathcal{L}_{\mathrm{dstl}}$ . (d) Effect of adjusting the two hyper-parameters, CFG range and CFG probability, in CFG distillation. As seen, these hyper-parameters can effectively tradeoff FID and CLIP score.
(c) Vanilla vs. CFG distill
(d) CFG hyper-parameters
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 hyper-parameters provide a very reliable way to tradeoff FID and CLIP score during training – this feature, as far as we know, has not been reported by any previous works.
Analysis of Original Loss for Distillation. In Eq. (11), we apply the original denoising loss  $(\mathcal{L}_{\mathrm{ori}}$  in Eq. (6)) during the step distillation. Here we show more analysis for the using  $\mathcal{L}_{\mathrm{ori}}$  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 FID at no loss of CLIP score.
- Fig. 7(b) provides a detailed analysis using different  $\gamma$  to balance the original denoising loss and the CFG distillation loss in Eq. (11). We empirically set a dynamic gamma to adjust the original loss into a similar scale to step distillation loss.
Analysis 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 a 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 hyper-parameters from the training of SD-v1.5 for the step distillation of SD-v2, and further tuning might lead to better results.
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
(a)  $\mathcal{L}_{\mathrm{ori}}$  in Eq. (11).
(b)  $\gamma$  in Eq. (11).
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  $\gamma$  in Eq. (11). Note that we employ dynamic scaling to adjust original loss  $(\mathcal{L}_{\mathrm{ori}})$  into a similar scale of step distillation loss  $(\mathcal{L}_{\mathrm{dstl}})$ . We show  $\gamma$  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  $\mathcal{L}_{\mathrm{ori}}$  added, while 1.0 refers to non-scaled  $\mathcal{L}_{\mathrm{ori}}$  where  $\mathcal{L}_{\mathrm{ori}}$  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 at 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.
(c) Number of teacher steps.
(d) Distillation 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 are gradually reduced by distillation that starts from a pre-trained teacher [33]. The later work further improves the inference cost of classifier-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.
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 evolving-training 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 our models for other mobile devices to achieve fast inference speed is also an interesting topic to study.
Broader Impacts. Similar to existing studies on content generation, our approach must be applied cautiously so that it 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.
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 is  $\frac{H}{8} \times \frac{W}{8}$ , which is  $64 \times 64$  for generating an image of  $512 \times 512$ .
In the main paper, we mainly benchmark the latency on iPhone 14 pro. Here we provide the runtime of the model on more mobile devices in Tab. 4.
In addition to mobile phones, we show the latency and memory benchmarks on Nvidia A100 40G GPU, as in Tab. 5. We demonstrate that our efficient UNet achieves over  $12 \times$  speedup compared to the original SD-v1.5 on a server-level GPU and shrinks  $46\%$  running memory. The analysis is performed via the public TensorRT [64] library in single precision.
Table 3: Detailed architecture of our efficient UNet model.
Stage	Resolution	Type	Config	UNet Model
Origin	Ours
Down-1	H/8 × W/8	Cross Attention	Dimension	320
# Blocks	2	0
ResNet	Dimension	320
# Blocks	2	2
Down-2	H/16 × W/16	Cross Attention	Dimension	640
# Blocks	2	2
ResNet	Dimension	640
# Blocks	2	2
Down-3	H/32 × W/32	Cross Attention	Dimension	1280
# Blocks	2	2
ResNet	Dimension	1280
# Blocks	2	1
Mid	H/64 × W/64	Cross Attention	Dimension	1280
# Blocks	1	1
ResNet	Dimension	1280
# Blocks	7	4
Up-1	H/32 × W/32	Cross Attention	Dimension	1280
# Blocks	3	6
ResNet	Dimension	1280
# Blocks	3	2
Up-2	H/16 × W/16	Cross Attention	Dimension	640
# Blocks	3	3
ResNet	Dimension	640
# Blocks	3	3
Up-3	H/8 × W/8	Cross Attention	Dimension	320
# Blocks	3	0
ResNet	Dimension	320
# Blocks	3	3
Table 4: Latency benchmark on iPhone12 Pro Max, iPhone13 Pro Max, and iPhone 14 Pro.
Device	Text Encoder (ms)	UNet (ms)	VAE Decoder (ms)	Overall (s)
iPhone14 Pro	4.0	230	116	1.96
iPhone13 Pro Max	5.7	315	148	2.67
iPhone12 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	8	209.6	12.2×
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, 20]. 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 quantitative results of our compressed VAE decoder in Fig. 8. The main paper shows that the image decoder constitutes a small portion of inference latency (369ms) compared to the original UNet from SD-v1.5. However, regarding our optimized pipeline (230ms × 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 is 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 yields the following advantages. First, our approach does not demand paired text-image samples, and it can generate unlimited data on-they-fly, benefiting the generalization of the compressed decoder. 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.
(a) Original (top) vs. our compressed (bottom) decoder.
(b) FID vs. CLIP 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 7 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)  $\sim$  Eq. (9) in the main paper.
Given the UNet inputs, time step  $t$ , noisy latent  $\mathbf{z}_t$ , and text embedding  $\mathbf{c}$ , the teacher UNet performs two DDIM denoising steps, from time  $t$  to  $t'$  and then to  $t''$  ( $0 \leq t'' < t' < t \leq 1$ ).
We first examine the process from  $t$  to  $t'$ , which can be formulated as,