EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models

Thanks to the reviewer for the valuable comments.

Q1: Clarify the difference in contribution from the combination of existing works (scale-aware optimization from LSQ, common distillation technique, and LoRA). In this paper, we propose a data-free, quantization-aware and parameter-efficient fine-tuning framework for low-bit diffusion models, which seamlessly integrates three tightly coupled components, QALoRA, scale-aware optimization and data-free distillation. Overall, our work has been well recognized by other reviewers, such as "the idea of QALoRA is novel" (Reviewer 5Wot), "this paper introduces low rank adapter and distillation loss ...is the first to achieve very good performance..." (Reviewers ec6g) and "the proposed QALoRA is data-free and efficient." (Reviewer BLfX).

To elaborate, our scale-aware optimization differs from the ``STEP SIZE GRADIENT SCALE" used in LSQ in both formula and function. From a formula standpoint, we adjust the gradient of QALoRA weight using the averaged quantization scale, whereas LSQ adjusts the gradient of quantization scales using the number of weights and precision. The goal of our technique is to improve the learning effectiveness of QALoRA, a critical issue in its optimization. As referred to Table 3 in the paper, our scale-aware optimization significantly improves EfficientDM's performance, leading to a 10.95 decrease in sFID.

Common distillation methods require access to the training dataset, which may be challenging due to massive size or copyright concerns. The recent work LLM-QAT [ii] proposes a distillation method that utilizes generations produced by full-precision (FP) models. However, it requires generating massive amounts of data using the teacher model beforehand and training the entire quantized model, incurring significant training costs. Our distillation scheme stands out by not requiring the generation and storage of any data. Instead, we input the same random Gaussian noise to both FP and quantized models and minimize the mean squared error between their denoising results. The distillation process is conducted in a data-free manner, which is efficient and easy to implement.

The proposed QALoRA is also distinct from the existing work QLoRA [iii], which applies LoRA to quantized models. As referred to ``Quantization-aware low-rank adapter" part in the paper, QLoRA incurs limitations when both weights and activations are quantized. The full-precision LoRA weights can not be merged with quantized model weights, leading to increased model size and massive floating-point calculations during inference. In contrast, the proposed QALoRA is a new variant of LoRA that can be seamlessly integrated with quantized model weights. After fine-tuning, there is no additional storage needed and the calculations between quantized weights and activations can be efficiently implemented for faster inference.

Q2: What happens if other quantization models also employ LoRA and distillation, which are common techniques to use? As referred to Q1, both QALoRA and our distillation scheme are distinct from previous methods and they are seamlessly integrated within the proposed data-free, quantization-aware fine-tuning framework for low-bit diffusion models.

[i] Esser, Steven K., et al. ``Learned step size quantization." ICLR 2020.

[ii] Liu, Zechun, et al. ``LLM-QAT: Data-Free Quantization Aware Training for Large Language Models." arXiv 2023.

[iii] Dettmers, Tim, et al. "Qlora: Efficient finetuning of quantized llms." NeurIPS 2023.

I appreciate the authors' response. I acknowledge the differences in the distillation technique. However, a part of my concerns remain regarding the novelty. In the main text, it says "Inspired by LSQ (Esser et al., 2019), a technique where quantization scales are optimized alongside other trainable parameters through the gradient descent algorithm, we extend the activation quantization scale into the temporal domain." According to the text, it seems like the proposed scale-aware quantization is a simple extension of the one proposed in LSQ. I do not think simple extention of existing work to temporal domain is novel enough to be considered as one of main contributions. Please clarify if there is any misunderstanding.

Dear Reviewer SYoS,

Thank you for your feedback. It appears there may be some misunderstandings regarding our proposed techniques. The scale-aware optimization and Temporal LSQ (TLSQ) are distinct techniques proposed in our paper. As referred to ``Variation of weight quantization scales across layers" part of the paper, the scale-aware optimization aims to improve the learning effectiveness of QALoRA due to variations in weight quantization scales across different layers, which is not an extension of LSQ.

In addressing the temporal dynamics of activation distribution across steps, we introduce Temporal LSQ (TLSQ), which is outlined in the "Variation of activation distribution across steps" part of the paper. To contextualize our innovation, we highlight a comparison with a related approach TDQ [i], as referred to Q1 in the general response. Unlike TDQ, which employs additional MLP modules to approximate quantization scales, TLSQ optimizes temporal quantization scales directly, offering a more computationally efficient approach. As referred to Tables 1 and 3 in the paper, TLSQ greatly improves the performance of EfficientDM, outperforming TDQ in both performance and efficiency under W8A8 and W4A8 bit-widths on CIFAR-10 dataset. These results affirm the advantages and contributions of the proposed TLSQ technique.

If there are any additional concerns you have, please feel free to let us know.

Thank you once again for your time, guidance, and consideration.

Best regards,

Authors of #3007.

[i] So J, Lee J, Ahn D, et al. "Temporal Dynamic Quantization for Diffusion Models." NeurIPS, 2023.

Thanks to the reviewer for the valuable comments.

Q1: How to get $\mathbf{x}_t$ in Eq. (7)? The Eq. (7) is not written properly. Please refer to Q2 in general response. The input data $\mathbf{x}_t$ is derived from denoising random Gaussian noise $\mathbf{x}_T \sim \mathcal{N}(0,1)$ with the FP model through $T-t$ iterations. We have revised the explanation of Eq. (7) in the revised paper.

Q2: Can the authors show the learned scales and analyze how the scale factors of certain layers change on different datasets? As referred to Figure B to C in the supplementary material of revised paper, we visualize the learned quantization scales and the variance of the activations collected from W4A4 LDM-4 model over ImageNet $256\times256$ and LSUN-Bedrooms. Layers exhibiting significant shifts in activation distribution show considerable variation in the learned scale across different steps, and vice versa.

Q3: Can the authors compare TLSQ with deciding the time-step-wise activation quantization scale using some calibration data or dynamic quantization? We conduct ablation experiments to compare the performance of TLSQ with dynamic quantization and time-step-wise post-training quantization, as referred to Table D in the supplementary material of revised paper. The naive Min-Max dynamic quantization fails to denoise images, even with QALoRA weights fine-tuned. Time-step-wise quantization scales, obtained through post-training techniques with calibration sets, outperform non-time-step-wise quantization by incorporating temporal information. Utilizing TLSQ further improves performance, resulting in a reduction of $1.64$ in FID and $2.03$ in sFID, showcasing the superiority of learned temporal quantization scales.

Q4: Does scale-deciding technique work well in the fewer-step regime? As referred to Table 2 in the paper, the performance on ImageNet $256\times256$ is evaluated with 20 steps, which demonstrates the strong performance of EfficientDM in the fewer-step regime. Moreover, we conduct experiments on LSUN dataset with 20 steps, as referred to Table B in the supplementary material of revised PDF. The proposed EfficientDM consistently outperforms previous quantized diffusion models in the fewer-step regime, which demonstrates the effectiveness of the interpolated quantization scales.

Performance of unconditional image generation on LSUN-Bedrooms with 20 steps. `N/A' denotes failed image generation.

Q5: Is QALoRA applied for all the weights, including the convolutions and the attention layers? As referred to ``Quantization settings" in Section 5.1 in the paper, the input embedding and output layers are fixed to 8-bit, whereas other convolutional and fully-connected layers are applied with QALoRA modules.

Q6: Discuss limitations and raise questions worth future studying. Thanks for your valuable comment. We have updated the ``Limitations and future work" part in the revised paper. Specifically, while EfficientDM can attain QAT-level performance using PTQ-level data and time efficiency, it employs a gradient descent algorithm for optimizing QALoRA parameters. Compared to PTQ-based methods, this places higher demands on GPU memory, particularly with diffusion models featuring massive parameters. To address this, memory-efficient optimization methods can be integrated. Moreover, efficient diffusion models for tasks such as video or 3D generation are still under explored.

Thanks to the reviewer for the valuable comments.

Q1: The experimental results listed in this paper are confusing and do not align well. In Table 2, the FID score of W4A4 model is much lower than FP model. While in Table A, the FID score of W4A4 model is much worse than the FP model. How to explain the gap in these two sets of experiments? There might be a misunderstanding of our results. To clarify, while FID provides an informative metric, it might not holistically capture the improved image quality. Therefore, referring to Table 2 in the paper, we thoroughly evaluate EfficientDM's performance by reporting the results of four standard metrics: IS, FID, sFID, and Precision. It's crucial to emphasize the comprehensive integration of these metrics when evaluating the model, rather than relying solely on any individual metric. Through a thorough evaluation with these metrics, FP models consistently outperform W4A4 models. Moreover, we have provided visualizations in Figures D-J in the supplementary material, convincingly showcasing the quality of images generated by EfficientDM.

Q2: Why is there a gap between the FID results reported in this paper and LDM [i]? The settings are different. LDM paper uses 250 DDIM steps to obtain an FID of 3.60, while we use 20 DDIM steps to obtain an FID of 11.28.

Q3: The paper should clarify the number of time steps used in TLSQ and discuss the settings. In practice, we set the number of time steps as 100 during fine-tuning and interpolate the learned temporal quantization scales when sampling with fewer time steps. We include the settings in Section 5.1 in the revised paper.

Q4: Does the QALoRA use LSQ algorithm to fine-tune the low rank adapter parameters and quantization step-size parameters? Yes, the weight quantization step-size parameters are optimized during the fine-tuning process.

Q5: Additional experimental results of using LSQ [ii] method on quantized diffusion models. As referred to Table A in the supplementary material, we conduct experiments using LSQ method on LSUN dataset. Despite consuming $15.8\times$ less time, our approach surpasses LSQ at W8A8 bit-width, demonstrating superior FID performance. At W6A6 bit-width, our approach achieves an FID that is only 0.07 higher than LSQ, striking a better balance between performance and efficiency.

Q6: Is $\mathbf{x}_t$ in Eq. (7) sampled from Gaussian noise with an FP model? Yes. Please refer to Q2 in the general response. The input data $\mathbf{x}_t$ is derived from denoising random Gaussian noise $\mathbf{x}_T \sim \mathcal{N}(0,1)$ with the FP model through $T-t$ iterations. We have revised the explanation of Eq. (7) in the revised paper.

Q7: Is the data-free fine-tuning in Eq. (7) fixed for 20 steps? As referred to Q2, we set the denoising step for fine-tuning as 100 in practice and interpolate the learned scales for sampling with fewer steps.

Q8: In Eq. (3), the quantization scheme has three parameters, are they all trainable? No. Only quantization scale $s$ is trainable and optimized with LSQ algorithm, while $l$ and $u$ are determined by the target bit-width. We have revised it in the revised paper.

[i] Rombach, Robin, et al. ``High-resolution image synthesis with latent diffusion models." CVPR, 2022.

[ii] Esser, Steven K., et al. ``Learned step size quantization." ICLR 2020.

Thanks for the authors' response. It well addresses my concerns.

Dear Reviewer ec6g,

In our method, during each training iteration corresponding to a single denoising step, only the step-size parameters associated with this denoising step are updated. To achieve this, we store a copy of the Adam optimizer's statistics for all other step sizes in a separate memory space and then temporarily zero them out to ensure they remain unaffected during the update process. For the step size parameters associated with this step, their statistics will be loaded back into the optimizer.

Best regards,

Authors of #3007.

Dear Reviewer ec6g,

As referred to "Variation of activation distribution across steps" part of the paper, we allocate temporal step-size parameters for activations and optimize them individually for each step. Therefore, during each training iteration, which corresponds to a single denoising step, only the step-size parameters associated with this step are updated, while other step sizes have zero gradients because they are not involved in forward propagation. After a complete denoising process spanning all denoising steps, all temporal step-size parameters undergo optimization.

Best regards,

Authors of #3007.

Thanks to the reviewer for the valuable comments.

Q1: Evaluate EfficientDM over recent text-to-image diffusion models. As referred to Figures K-M in the supplementary material, we include visualization results to evaluate EfficientDM over text-to-image Stable Diffusion models.

Q2: The differences from TDQ should be discussed. Please refer to Q1 in the general response. TDQ incorporates additional MLP modules for each quantized layer to approximate quantization scales. This approach entails training these extra MLP modules, leading to a significant surge in training expenses. Conversely, TLSQ introduces a limited number of adjustable quantization scales per layer, making TLSQ more computationally efficient than TDQ. Empirically, TLSQ achieves a 20.9% faster training iteration speed compared to TDQ over DDIM model on CIFAR-10 dataset.

Q3: Formulating the gradient of LoRA weights. Referring to Eq.~(6) in the paper and considering the impact of STE [i], the gradient of LoRA weights $\mathbf{BA}$ can be expressed as:

$

\frac{\partial \mathbf{Y}}{\partial (\mathbf{BA})}= \frac{\partial \mathbf{Y}}{\partial \hat{\mathbf{W}}} \frac{\partial \hat{\mathbf{W}}}{\partial (\mathbf{BA})} = \hat{\mathbf{X}}^\top

$

The key insight is that the gradient magnitudes of LoRA weights remain unaffected by the quantization scale, which is exactly the step size of the quantization step-like function. Consequently, in layers with larger quantization scales, the minor updates to LoRA weights will be diminished by the round operation, which causes the ineffective learning of QALoRA.

Q4: The notation for scale-aware optimization in Fig. (2) is inconsistent with Eq. (8). Thanks for the valuable comments. We have corrected it in the revised version.

[i] Bengio, Yoshua, Nicholas Léonard, and Aaron Courville. "Estimating or propagating gradients through stochastic neurons for conditional computation." arXiv, 2013.