KOALA: Empirical Lessons Toward Memory-Efficient and Fast Diffusion Models for Text-to-Image Synthesis

Summary

The paper at hand performs a series of ablation studies evaluating the impact of three design decisions on the quality of Text to Image Generation models. Specifically, the authors look at the target of knowledge distillation, training data as well as choice of teacher for distillation. In the end, the authors present a model named KOALA, building upon the insights gained from those ablations. The model is smaller and more efficient than the comparisons presented in the paper.

Strengths

Improving the efficiency of Text To Image Generation models is of key importance to the research and industrial community, as the size an inference speed of current state of the art models is often a limiting factor for real world applications. The paper shows positive results and a good set of comparisons to prior art. The paper touches important aspects, such as distillation objectives as well as data selection.

Weaknesses

My main concern about the paper is its approachability. The presentation of the paper is subpar. Specifically, the paper is very verbose and repetitive, but then on the other hand imprecise and lacking important details. The paper could be half its size and get more information across. For example, the paper has an almost two page discussion on knowledge based distillation, but I could not a single precise definition how the proposed approach actually performs. Similarly, the entire discussion on the U-Net architecture could be a single table. This would make it much more approachable and easier to understand.

Questions

It would be great, if the authors could provide a short precise explanation how the proposed distillation loss works.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The limitations are discussed in the paper.

Summary

This paper presents KOALA, a pair of efficient text-to-image synthesis models that reduce computational and memory requirements compared to the base model. This paper achieves this through three key innovations: knowledge distillation into a compact U-Net, the strategic use of high-resolution images and detailed captions in training data, and the employment of step-distilled teachers to enhance the learning process. The resulting models, KOALA-Turbo and KOALA-Lightning, demonstrate faster inference times and the ability to generate high-quality images on consumer-grade GPUs.

Strengths

1. This paper is easy to follow. 2. This paper offers a practical solution for generating high-quality images with reduced computational requirements.

Weaknesses

1. Although KOALA achieves good visual quality, the extent of degradation in text rendering is unclear as it has not been compared with a baseline in terms of text rendering capabilities. 2. For each baseline, KOALA requires a carefully designed pruning network architecture followed by retraining, which entails a significant acceleration cost compared to training-free methods. However, this paper does not discuss or compare with training-free methods, nor does it compare with step distillation approaches.

Questions

In the "Lesson 3" section, it is puzzling that SDXL-Base, as a teacher model, performs the worst. Intuitively, a better teacher model should have greater potential to distill a superior student model. I am curious if the results presented in the paper could be attributed to a suboptimal distillation approach.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed both limitations and potential negative social impacts.

Summary

This paper presents a set of empirical guidelines to use when distilling Stable Diffusion XL (SDXL) when computational/data resources are limited. The presented guidelines focus on 1) identifying which transformer blocks to drop, 2) what features are the best to distill, and 3) which publicly-available datasets are the best to use. The paper presents analyses for all guidelines, including ablations and comparisons to what other popular distillation approaches do. Insights are extracted from these analyses and used to motivate the guidelines. Performance results are presented showing that the guidelines result in distilled models that are quite competitive with the parent/other distilled models yet are much more computationally efficient.

Strengths

The main topic of this paper -- how to more easily generate cheaper SDXL distillations -- is of much relevance to virtually all associations. This paper does a thorough job of exploring many distillation aspects. The analyses it reports give good insights into what aspects are most important. For example, the feature distillation ablation section shows that self-attention are the best feature to distill. The authors then use those results to motivate removing fewer transformer decoder blocks than encoder ones. Fig 3c is compelling and really drives home the primary role of self attention in distillation. The suggested recipe for removing blocks goes beyond just a basic approach of removing all blocks in a stage, as some other distillation approaches suggest. I also thought the insights into the various LAION datasets was interesting and that the conclusion of using few high res images with detailed prompts is better than more low res with with short (or detailed) prompts is useful. Further, the performance results from the distilled models are quite competitive with other variants.

Weaknesses

The guidelines this paper presents are very focused on specifically SDXL. It is unclear if any of them would apply to other transformer-based models. Whatever the next version after SDXL is created, it is unclear how relevant these guidelines would be.

Questions

1) Sec 3.1.1. mentions both "block removal" as well as "layer-wise removal". However, the text describes layer-wise removal as "reducing the number of transformer layers". I found this vernacular confusing. It is common practice to let a "layer" mean a "block", so both approaches sounds as if they are doing the same thing: removing blocks. I think (?) when the text refers to "block removal" (108) it means "removal of 100% of the blocks in a sequence" but when it refers to "layer-wise removal" (line 117) it means "removal of *less* than 100% of the blocks of the sequence". Is this understanding correct? If so, then I recommend re-writing this section to make it more clear that "block removal" is the special case of "layer-wise removal" when all blocks are removed. 2) Table 10d shows that adding DW-1 and UP-3 feature distillation to SA + LF further boosts performance. But the formal paper stops at suggesting SA + LF3 is optimal. Why? 3) Nit: The only time I see the "KOALA" acronym explained is implicitly in a couple of the figure captions. Adding it somewhere in the introduction would be nice. 4) Nit: Line 199 uses the phrase "in the second row of the table" before any table has been presented to the reader.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper presents the commonly suggested limitations of stable diffusion -- it doesn't generate text well and has problems with multiplicity.

Summary

1. The authors propose two different designs for an efficient denoising U-Net based on SDXL. 2. The authors present three empirical lessons for developing efficient U-Net designs.

Strengths

1. The authors present empirical findings to justify their design considerations and provide further analyses to examine the effects of those choices. 2. The authors conduct a systematic analysis with extensive quantitative and qualitative experiments, comparing their results with different baselines.

Weaknesses

1. The proposed method is specific to the SDXL U-Net and is based on heuristics, rather than presenting a generalizable approach. This makes the paper resemble a technical report more than an academic paper. It would be beneficial if the authors could further clarify their novelty and contributions in terms of academic research. 2. The authors present some failure cases of their methods, including rendering long legible text, complex prompts with multiple attributes, and human hand details. Do the teacher models and other baseline models also suffer from the same problems? Were any of these failure cases aggravated in the proposed model to some extent? If so, what could be the reasons?

Questions

Please refer to the "Weaknesses" section.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately address the limitations in section 5.

Response

I would like to thank the authors for their response to both my questions and the questions by the fellow reviewers. After going over the other reviews and considering all the answers, I come to believe that the contributions and presentation of the paper might just pass the bar for acceptance. However, I strongly encourage the authors to significantly improve the organization and conciseness of the paper.