Efficient Low-Rank Diffusion Model Training for Text-to-Image Generation

Thanks for your hard work, careful reviews and useful comments! We would like to clarify several major misunderstandings about this work. Please see our responses below:

Background section 2.1: In our Background section 2.1, we refer to the definition and formulation of re-parameterization from Aghajanyan et al [1] to lay the groundwork for our proposed method. The philosophy behind our work and the referred work is the same, this is the reason for including the background section 2.1. In Stage I, we employ low-rank decomposition to break down the U-Net model, enabling training in a reduced-dimensional space. In contrast to Aghajanyan et al [1], where the Fastfood transform is utilized for optimization in a lower dimension (d) and then projected back to a higher dimension (D), our approach differs significantly. Specifically, in Stage I, our model undergoes direct decomposition and training in this lower-dimensional space without any subsequent projection back. Stage II involves fine-tuning a smaller subset of parameters compared to the original full model. This fine-tuning occurs initially through Kronecker decomposition followed by low-rank decomposition, and these modifications are directly added to the original parameters in D dimensions. Hence, we introduce the operator $\mathcal{P}$ for projecting from D to d dimensions, distinguishing our notation from that used in Aghajanyan et al [1]. To avoid confusion, we will revise our phrasing regarding 'a set of D parameters' to clearly indicate a set of parameters in D-dimensional space. Additionally, we will clarify the initial sentence "Let $\theta_D = [\theta_0, \theta_1, ..., \theta_m]$ be a set of D parameters that parameterize some model $f(·, \theta)$" for better understanding.

Technical novelty Our contribution lies in the novel combination of low-rank decomposition with knowledge distillation in an end-to-end approach, which to our knowledge, is the first of its kind. This integration not only reduces memory costs but also preserves the effectiveness of the diffusion models. Besides, our application in the context of text-conditioned diffusion models, especially the controllable text-to-image generation tasks presents unique challenges. We make the Kronecker decomposition first and then the low-rank decomposition, with the decomposed weights shared across different conditions. We will further elucidate this aspect in our revised manuscript, highlighting the unique challenges and solutions in this context.

Comparison with LoRA We will add baseline comparisons with LoRA in the revision. We also want to clarify the 10GB memory reduction happens on the Uni-ControlNet model rather than the vanilla Stable Diffusion, which contains more modules and parameters.

Writing Issues: We will be thoroughly reviewed and corrected to improve the clarity and accuracy of the manuscript in the revision.

Low-rank decomposition for linear transformations: We want to clarify that equation (3) applies the decomposition for attention-layers, specifically for $W_k$ and $W_v$ matrices. Furthermore, as shown in the LoRA paper [1], they freeze the MLP modules from the consideration of simplicity and parameter-efficiency and only apply it on attention-matrices. Apply the decomposition on more weights can bring additional costs.

[1] LoRA: Low-Rank Adaptation of Large Language Models

Time to get a task-specific model: We train our model and baselines for the same number of iterations and that’s why we report the training time per iteration. Table II compares the training time for the individual Stage II, where the standard Stable Diffusion without distillation can also be used as the backbone. We propose Stage I to distillate a lightweight model that can serve as an off-the-shelf pretrained model and it’s not necessary to perform Stage I and II each time so it is meaningless to sum them together for comparison.

Difference in the experimental setup: Table 2 compares the cost only for Stage II while Table 3 shows the overall cost.

Qualitative Comparison: We will include additional experiments and analyses that compare our method with other baselines in additional datasets. We will add more qualitative results. We will also display more examples for the qualitative comparison in the revision.

We are grateful for the detailed feedback from the reviewer. Please find our responses below.

Novelty of Approach: Our contribution lies in the novel combination of low-rank decomposition with knowledge distillation in an end-to-end approach, which to our knowledge, is the first of its kind. This integration not only reduces memory costs but also preserves the effectiveness of the diffusion models. Besides, our application in the context of text-conditioned diffusion models, especially the controllable text-to-image generation tasks presents unique challenges. We will further elucidate this aspect in our revised manuscript, highlighting the unique challenges and solutions in this context.

Improvement in Controllability: We want to clarify that the target of this work is not to improve the controllability of the generation model by the proposed approach but to improve the overall efficiency of training such a model.

Model Distillation and Fine-Tuning Sequence: We distill the model first followed by fine-tuning because the downstream task requires a pre-trained Stable Diffusion backbone, and we perform distillation first to obtain a lightweight backbone so that we can perform fine-tuning based on it. However, your suggestion to reverse this sequence is intriguing. We will explore this alternative approach and include the findings in our revised manuscript.

Low-Rank Assumption: We appreciate your observation on the low-rank assumption. To address this, we will conduct and include ablation studies in our revised paper to validate this assumption and examine its impact on performance.

Comparison with LoRA: We would like to clarify that we didn’t assume the matrix is low-rank across both our two stages. In the first stage, rather than assume the U-Net is low rank, we perform low-rank decomposition over the model. In the second stage, we perform low-rank finetuning which can work on dense matrices as well. In the revised manuscript, we will include these comparisons to demonstrate the relative strengths and weaknesses of our method.

Theoretical Clarifications: Thank you for pointing out the need for clarification in our theoretical analysis, particularly concerning Equation 8. We would like to clarify that in Equation 8, \theta_0^d is indeed intended to represent a transformation of \theta_0, but not strictly a linear one. Instead, \theta_0^d is derived through a more general transformation operator, which we denote as 'T'. This operator 'T' encompasses a broader range of transformations, which includes but is not limited to Singular Value Decomposition (SVD).

We thank the reviewer for the constructive suggestions. Your suggestions will be incorporated in our next revision. Please see our clarifications and responses below:

Qualitative and Quantitative Results: In the revised manuscript, we will include additional qualitative examples to better illustrate the capabilities of our method. Furthermore, we will expand our quantitative analysis to include KID and Inception Score, alongside FID and CLIP Score, to provide a more comprehensive evaluation.

Novelty of Approach: Our contribution lies in the novel combination of low-rank decomposition with knowledge distillation in an end-to-end approach, which to our knowledge, is the first of its kind. This integration not only reduces memory costs but also preserves the effectiveness of the diffusion models. Besides, our application in the context of text-conditioned diffusion models, especially the controllable text-to-image generation tasks presents unique challenges. [a][b][f] propose knowledge distillation and low-rank decomposition methods mainly for convolutional neural networks, intelligent edge applications, and language models, while ours propose specific methods for diffusion-based text-to-image generation models. [c] proposes to use LoRA to compose multiple customized concepts, while ours propose an end-to-end efficient training strategy that can deal with different input conditions such as segmentation maps and depth maps. [e] proposes a low-rank matrix factorization technique for panchromatic (PAN) image, where the task is also completely different from ours.

[a] Li, Tianhong, et al. "Few sample knowledge distillation for efficient network compression." CVPR 2020.

[b] Dai, Cheng, et al. "A tucker decomposition based knowledge distillation for intelligent edge applications." Applied Soft Computing 101 (2021): 107051.

[c] Gu, Yuchao, et al. "Mix-of-Show: Decentralized Low-Rank Adaptation for Multi-Concept Customization of Diffusion Models." arXiv preprint arXiv:2305.18292 (2023).

[d] Yao, Xufeng, et al. "Distill Vision Transformers to CNNs via Low-Rank Representation Approximation." (2022).

[e] Rui, Xiangyu, et al. "Unsupervised Pansharpening via Low-rank Diffusion Model." arXiv preprint arXiv:2305.10925 (2023).

[f] Ryu, Simo, Seunghyun Seo, and Jaejun Yoo. "Efficient Storage of Fine-Tuned Models via Low-Rank Approximation of Weight Residuals." arXiv preprint arXiv:2305.18425 (2023).

Experimental Details: We will provide a more detailed description of both the teacher and student architectures, the extent of fine-tuning, and the specific hyperparameters used in the revision.

Responding to the Question: Regarding the fewer steps and data required for fine-tuning due to the reduced number of parameters in ELR compared to Uni-ControlNet, this aspect was considered in our experiments. As can be seen from our experimental results in Table 2, the training time required is less than half of the Uni-ControlNet and we train EIR-diffusion the same number of iterations with Uni-ControlNet. As for data, we only require 5 million text-image pairs, while Uni-ControlNet requires 10 million. Ours saves 50%, demonstrating the data efficiency of EIR-Diffusion.

We thank the Reviewer for the constructive feedbacks. We value your specific suggestions, which we have carefully considered in our responses below and we will include them in our next revision.

Novelty of Low-Rank Decomposition: Our contribution lies in the novel combination of low-rank decomposition with knowledge distillation in an end-to-end approach, which to our knowledge, is the first of its kind. This integration not only reduces memory costs but also preserves the effectiveness of the diffusion models. Besides, our application in the context of text-conditioned diffusion models, especially the controllable text-to-image generation tasks presents unique challenges.

Extensive Evaluations: We will include additional experiments and analyses that compare our method with other baselines in diverse domains, datasets, and resolutions. This will provide a clearer understanding of the effectiveness and versatility of our method.

Detailed Experiment Settings: In our revised manuscript, we will provide more comprehensive details of our experimental setup, ensuring that other researchers can replicate our work accurately.

Additional Ablations: we used DDIM scheduler in the inference as the standard scheduler for accelerating the text-to-image generation. In the revised manuscript, we will include results and observations related to the use of various schedulers in our method. We are currently conducting experiments on other datasets and will include these results in our revised submission. These additional experiments will provide a broader view of how our method performs across different task setups.

Effects on compression: we want to clarify that the qualitative results contain the task of controllable text to image generation, where the generation is regularized by the additional input condition. In this context, our method produces results that are comparable in terms of diversity and quality. We have observed some interesting phenomena in this regard, both positive and negative. For instance, we noted a decrease in the performance of image quality at higher compression levels when increasing the number of decomposed matrices to 16. We will add more qualitative results in the revision.