BitsFusion: 1.99 bits Weight Quantization of Diffusion Model

Summary

This paper proposes a weight quantization method to quantize the UNet of SDv1.5 to 1.99 bits while maintaining model performance comparable to the floating-point model. The approach includes a series of techniques such as bit-width allocation for mixed-precision quantization, low-bit diffusion model initialization, and a two-stage training pipeline. Extensive experiments demonstrate the effectiveness of the proposed method. Results show that the quantized diffusion model can even outperform the floating-point SDv1.5.

Strengths

1. The paper is well-written and easy to follow. 2. The per-layer sensitivity analysis is comprehensive, considering both quality and contextual information, which is helpful for future research. 3. The effectiveness of the proposed approach is validated across multiple datasets and evaluation metrics. 4. The generated results look impressive.

Weaknesses

1. It is well-known that QAT-based approaches can achieve superior performance at the expense of additional training effort. While this work involves training the weights via QAT, it does not include comparisons with other QAT-based methods such as Q-DM and TDQ. 2. It is unclear how the sensitivity threshold in Section 3 is determined and whether this hyperparameter generalizes to other models.

Questions

1. The setting of activation is not mentioned in the paper. Are the activations kept in floating point? 2. What is the dataset used for stage-II training? Suppose the training set of the corresponding dataset is used. In that case, claiming the model outperformed the floating-point SDv1.5 model may be inappropriate since the SDv1.5 model is evaluated under a zero-shot setting. 3. The average bits uses $\log(2^n + 1)$ bits for calculation. Is it necessary to incorporate an optimal encoding method such as Huffman encoding to reach this average bit-width?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

See weaknesses and questions.

Summary

The paper presents BitsFusion, a novel method for weight quantization of diffusion models, specifically applied to the UNet architecture in Stable Diffusion v1.5. The approach quantizes weights to 1.99 bits, achieving a model size reduction of 7.9 times while enhancing or maintaining image generation quality. The experiment evaluates the quantized model across various benchmarks, including MS-COCO, TIFA, GenEval, and human evaluations, and demonstrates superior performance of the quantized model compared to full-precision Stable Diffusion v1.5.

Strengths

1.The paper is well-organized and presents its content in an accessible manner. The organization of the paper is clear. 2.This paper mainly focuses on quantizing large-scale models like Stable Diffusion v1.5 with fewer bits than 4, which has significant implications for the industry and has achieved satisfactory results. 3.The paper methodologically addresses the challenges associated with low-bits quantization in SD and provides a robust solution through the BitsFusion framework. The experiments demonstrate the effectiveness of BitsFusion in achieving superior performance compared to existing methods on various large-scale datasets.

Weaknesses

1.The motivation of the paper seems to be heuristic and lacks necessary theoretical analysis, but the impact of this method in the industry is worth looking forward to.

Questions

What's the total resource consumption for this entire process? For instance, how many GPUs with what memory capacity were utilized, and for how many hours was the training conducted? How does this compare to the resources required for a baseline that trains the original SD v1.5?

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

See weakness.

Summary

This paper demonstrates a quantized diffusion model called BitFusion, which successfully quantize the Stable Diffusion (SD) v1.5 to 1.99 bits with 7.9x smaller size. They first analyse the SD model in a layer perspective and assign the optimal bit based on the analysis. Then they propose a training pipeline to perform quantization-aware training (QAT) based on the analysis result. The experiment result shows that the quantized diffusion model BitFusion outperforms the full-precision SD v1.5 model.

Strengths

1. The result is significant. BitFusion successfully compressed the SD model while maintaining its performance. 2. The authors successfully combine novel quantization techniques and existing quantization techniques, showing an effective way to quantize SD model to extreme low bit. 3. The paper provides valuable observations about different layers' behaviour under quantization.

Weaknesses

1. The analysis and experiment according to scaling factor initialization is not comprehensive. The authors propose alternating optimization for scaling factor initialization which is based on the MSE between quantized weight and full-precision weight. However, quantized weight initialization is well-studied and there are existing weight initialization technique that optimize the same objective. Nevertheless, there are some advanced initialization technique like adaround [1]. Those techniques are adapted by other Diffusion Model quantization research [2]. An experiment to compare the alternating Opt. and those initialization techniques is necessary. 2. Stable Diffusion 1.5 is no longer the SOTA in Diffusion text-to-image model. Diffusion transformer (DiT) is the new trend in diffusion model research. The conclusion on UNet-based diffusion model like SD 1.5 might not be precise for the DiT diffusion model. 3. FID is widely adapted by most of the diffusion model quantization research. Although the authors provide human evaluation and justify that FID is not accurate compared with human evaluation, the FID comparison between BitFusion and other baseline models should be included in the main body of the paper. 4. In this paper [3], the authors propose a 1.58-bit LLM quantization approach comparable to Bit Fusion at model size reduction. However, they can perform activation quantization to 8-bit while Bit Fusion doesn't. [1] Nagel, Markus, et al. "Up or down? adaptive rounding for post-training quantization." International Conference on Machine Learning. PMLR, 2020. [2] Li, Xiuyu, et al. "Q-diffusion: Quantizing diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023. [3] Ma, Shuming, et al. "The era of 1-bit llms: All large language models are in 1.58 bits." arXiv preprint arXiv:2402.17764 (2024).

Questions

NA

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors honestly point out their limitation: not quantizing VAE and text encoder and not quantizing activation. For the first limitation, it is quite common that diffusion model quantization research doesn't quantize VAE and text encoder according to the literature. But for the second limitation, the authors leave it for future work and doesn't address it. Activation quantization is recommended, as the fact that activation usually has higher quantization difficulty.

Summary

This paper proposes a novel weight quantization method called "BitsFusion" for compressing the Stable Diffusion v1.5 model. The primary goal is to address the issue of large model sizes, which hinder the deployment of diffusion models on resource-constrained devices. The BitsFusion framework quantizes the UNet component of Stable Diffusion v1.5 to 1.99 bits, achieving a 7.9× reduction in model size while improving image generation quality. The approach involves mixed-precision quantization, novel initialization techniques, and an improved two-stage training pipeline. Extensive experiments on various benchmark datasets demonstrate that the quantized model outperforms the original full-precision model in terms of generation quality and text-image alignment. My detailed comments are as follows.

Strengths

1. The BitsFusion framework is an innovative approach to compressing the UNet component of Stable Diffusion v1.5 to 1.99 bits, achieving a significant reduction in model size while enhancing image quality by using mixed-precision quantization, novel initialization techniques, and an improved two-stage training pipeline. 2. The paper provides a thorough analysis of quantization errors and develops a mixed-precision strategy based on this analysis, which contributes to the theoretical understanding of low-bit quantization for large-scale diffusion models. 3. The results demonstrate the effectiveness of the proposed framework. The quantized model consistently outperforms the original full-precision model across various metrics, including CLIP score, TIFA score, and GenEval score. The model's ability to achieve a 7.9× reduction in size while maintaining or improving performance highlights its practical potential for deployment on resource-constrained devices. 4. The manuscript is well-written and easy to understand, with clear explanations of the methodology and experimental procedures. Sufficient experimental details are provided, ensuring that the results can be reproduced by other researchers, which enhances the paper's credibility and utility.

Weaknesses

1. The two-stage training pipeline, while effective, may introduce additional computational complexity and training time, which could be a drawback for some applications. However, the paper does not explicitly discuss the additional computational complexity and training time it introduces. To provide a more comprehensive evaluation of this method, it is recommended that the authors discuss the potential additional computational overhead introduced by the two-stage training pipeline and its impact on training time. 2. The study demonstrates that there is a high correlation between certain metrics, such as MSE and PSNR, and uses these correlations to validate the effectiveness of their quantization methods. However, correlation does not imply causation. The study does not delve into why these metrics are correlated or explain the intrinsic mechanisms that cause certain layers to be more sensitive to quantization (e.g., why quantization affects some layers more than others, and what intrinsic properties of these layers make them more prone to quantization-induced errors). Understanding these mechanisms is crucial for improving the quantization process, as knowing why certain layers are sensitive can lead to more targeted and effective quantization strategies. 3. The claim that the 1.99-bit quantized model outperforms the full-precision model in all evaluated metrics is a significant overgeneralization. This assertion, based on specific experimental settings and datasets, may not fully represent the diversity encountered in real-world applications. Primarily, the evaluation on the MS-COCO 2014 validation set, though extensive, does not cover a wide variety of prompts, limiting the assessment of the model's ability to handle different types of inputs (e.g., complex scenes, abstract concepts, specific artistic styles). Additionally, the tests do not include various real-world conditions such as different lighting, backgrounds, or object complexities. Therefore, the claim that the 1.99-bit quantized model outperforms the full-precision model in all evaluated metrics is premature. A more robust evaluation including varied datasets, prompts, and real-world conditions, such as [A-E], is necessary to substantiate such a broad claim. 4. The discussion regarding model compression techniques is insufficient. It would be better for the authors to present more model quantization methods [F-J]. [A] Visual Genome: Connecting Language and Vision Using Crowdsourced Dense Image Annotations. International Journal of Computer Vision. IJCV. 2017. [B] Flickr30k Entities: Collecting Region-to-Phrase Correspondences for Richer Image-to-Sentence Models. ICCV. 2015. [C] The Open Images Dataset V4: Unified Image Classification, Object Detection, and Visual Relationship Detection at Scale. IJCV. 2020. [D] Conceptual 12M: Pushing Web-Scale Image-Text Pre-Training To Recognize Long-Tail Visual Concepts. CVPR. 2021. [E] Scene Parsing through ADE20K Dataset. CVPR. 2017. [F] Binary quantized network training with sharpness-aware minimization. Scientific Computing 2023. [G] Network Quantization with Element-Wise Gradient Scaling. CVPR 2021. [H] Single-path bit sharing for automatic loss-aware model compression. TPAMI 2023. [I] Generative Data Free Model Quantization with Knowledge Matching for Classification. TCSVT 2023.

Questions

Please check the weakness section.

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

N/A