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.
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.
* The results of the proposed architectural compression are impressive, although some questions regarding training computations and resources exist (described below).
* The improved step distillation is well-motivated and novel. I think this can be broadly applicable to diffusion-based large models beyond Stable Diffusion.
* 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).
* The paper is very well-written and easy-to-follow. The terms and equations are properly described.
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.
* [49] Coyo-700m: Image-text pair dataset. https://github.com/kakaobrain/coyo-dataset, 2022.
* [38] Laion-5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08402, 2022.
* [+] an internal dataset with high-resolution images to fine-tune our model for more pleasing visual quality.
* 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)
* 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.
* 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.
* On Architectural Compression of Text-to-Image Diffusion Models, https://arxiv.org/abs/2305.15798
* 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.
* 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?
* 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).
* 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.
Questions
Please check the above weakness section.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.