BitsFusion: 1.99 bits Weight Quantization of Diffusion Model

Q1. Comparison with QAT-based approaches such as Q-DM and TDQ.

A1. Thanks for the suggestion. We have compared the QAT-based approach LSQ and EfficientDM in Table 3 of the main paper. Here, we provide more results for Q-DM and TDQ on PartiPrompts with a CFG scale of 7.5. Compared to Q-DM with 2-bit weight quantization, our BitsFusion achieves a CLIP score of 0.3212, which is 0.228 higher than Q-DM. For a fair comparison with TDQ, which optimizes activation quantization through dynamic quantization, we also apply 8-bit activation quantization to our 1.99-bit BitsFusion. As shown in the table below, our BitsFusion achieves a CLIP score of 0.3193, surpassing TDQ with a score of 0.2907.

| Methods | Weight Bit-width | Activation Bit-width | CLIP score | | :---- | :---- | :---- | :---- | | Stable Diffusion v1.5 | 32 | 32 | 0.3175 | | Q-DM | 2 | 32 | 0.2984 | | BitsFusion | 1.99 | 32 | 0.3212 | | TDQ | 2 | 8 | 0.2907 | | BitsFusion | 1.99 | 8 | 0.3193 |

Q2. How to determine the sensitivity threshold in Section 3.

A2. Thanks for the comment. We empirically set the sensitivity threshold to achieve an average of 1.99 bits. The sensitivity threshold only affects the average bits. It can be adjusted as needed for other models with different bit requirements, thus it is generalized to other models.

We would like to kindly note here the sensitive threshold is not a hyperparameter that influences the model training. Thanks for mentioning this. We will modify our manuscript accordingly to emphasize this distinction.

Q3. Activation quantization.

A3. Thanks for the valuable question. Since this paper is focusing on reducing the model size, we explore the weight quantization on the diffusion model and keep activation as floating point value. However, 8-bit activation quantization can be directly applied to our approach. Based on 1.99 bits weight quantization, we further apply 8-bit activation quantization to all layers except for the first and last convolutional layer and train the corresponding scaling factors. The experimental settings are the same as those described in Section 5 of the main paper. With a CFG scale of 7.5, our stage-I model and stage-II model can achieve CLIP scores of 0.3172 and 0.3193 on the PartiPrompts, respectively. For the TIFA score, our stage-I model and stage-II model can achieve 0.786 and 0.805 with a CFG scale of 7.5. With two-stage training, our BitsFusion model with a 1.99-bit weight and 8-bit activation quantization can still outperform the full-precision SDv1.5.

| Methods | Weight Bit-width | Activation Bit-width | CLIP score | TIFA | | :---- | :---- | :---- | :---- | :---- | | Stable Diffusion v1.5 | 32 | 32 | 0.3175 | 0.788 | | Stage-I | 1.99 | 8 | 0.3172 | 0.786 | | Stage-II | 1.99 | 8 | 0.3193 | 0.805 |

Q4. Datasets used in the stage-II training.

A4. Thanks for the question. Our training data consists of synthetic data and real images, and there is no overlap between the training data and evaluation data. To clarify, our model performs zero-shot evaluation, which is the same as SDv1.5. Specifically, we perform different evaluations to verify the effectiveness of BitsFusion, such as CLIP scores on the MS-COCO dataset and PartiPrompts, GenEval score, and TIFA score. PartiPrompts provides prompts from different categories. GenEval provides prompts to evaluate compositional image properties such as object co-occurrence, position, count, and color. TIFA provides various prompts and uses the VQA model to measure the text-to-image model if it can accurately answer the questions. Evaluating CLIP scores, TIFA score, and GenEval score do not include any reference images, and we do not have the same prompts from these datasets in the training.

Q5. Is it necessary to incorporate an optimal encoding method such as Huffman encoding to reach the averaged bit-width?

A5. Thanks for the comment. We calculate the average bits according to Section I in the Appendix with \frac{\sum_{i} log 2^{b_{i}^*}+1* N_i + 16*N_{tf}}{N_w}, where $b_{i}^{\ast}$ is the calculated bit-width in the $i_{th}$ layer, $N_i$ is the number of weights of the $i_{th}$ layer, $N_{tf}$ is the number of parameters for pre-cached time features, and $N_w$ is the total number of weights in linear and convolutional layers. We do not need the Huffman encoding to reach this average bit-width. However, as suggested by the reviewer, Huffman encoding is a valuable tool to further reduce our storage size.

Q1. About the motivation of this paper.

A1. The overall motivation for performing quantization on the diffusion model is to reduce significant burdens for transferring and storing the model due to its large size. To this end, we propose several methods with corresponding motivations for model quantization:

Mixed-precision Strategy:

• Mixed-precision Strategy: Various layers exhibit different sensitivities during quantization. To achieve a lower quantization error, it is crucial to employ a mixed-precision approach where more sensitive layers are quantized at higher bit-widths to retain essential information, while less sensitive layers can be quantized at lower bit-widths. In Section 3.2, we conduct the quantization error of layers for different bit-widths via appropriate metrics and obtain valuable observations to guide our mixed-precision approach.

Initialization Schemes:

• Time Embedding Pre-computing and Caching: During the inference, the embedding for a fixed time step is always the same (Line 177~179). Therefore, we can pre-compute the embedding offline and load cached values during inference, instead of computing the embedding every time.
• Adding Balance Integer: We observe that weight distributions in the diffusion model are symmetric around zero. However, the existing quantization works disrupt this symmetric property (Line 182~189). Therefore, we introduce an additional value to balance the original values to maintain this property.
• Scaling Factor Initialization via Alternating Optimization: Min-Max initialization leads to a large quantization error and the increased difficulty to converge in extremely low-bit quantization settings like 1-bit (Line 198~201). Therefore, we propose using an alternating optimization method to improve scaling factor initialization.

Training Schemes:

• CFG-aware Quantization Distillation: Distillation is an effective approach to recover the performance degradation caused by quantization. Additionally, considering that CFG enhances the performance of the diffusion model, we leverage CFG-aware distillation to effectively recover the quantization error (Line 217~218).
• Feature Distillation: To further improve the generation quality of the quantized model, we distill the full-precision model at a more fine-grained level (Line 220~221). Therefore, we adopt the feature distillation.
• Quantization Error-aware Time Step Sampling: We analyze the quantization error during various time steps during training. As shown in Fig. 4 in the main paper, the quantization error does not distribute equally across all time steps in the quantized model (Line 230~240). Therefore, we propose a quantization error-aware time step sampling based on Beta distribution to reduce the overall quantization error.

We believe that our comprehensive analysis and the proposed methods will contribute significantly to the field of quantization on diffusion models.

Q2. About training computation.

A2. Thank you for the suggestions. Here we provide more details for our two-stage training pipeline. For the stage-I training, we use 8 NVIDIA A100 GPUs with 40GB memory, and a total batch size of 256 to train the quantized model for 20K iterations. The total training time is within 40 hours. For the stage-II training, we use 32 NVIDIA A100 GPUs with 40GB memory, and a total batch size of 1024 to train the quantized model for 50K iterations. The total training time is within 100 hours.

Note that the stage-II training aims to further improve the performance of the quantized model such that it is better than the full-precision Stable Diffusion v1.5. If reducing the training cost would be the goal, we can remove stage-II training and only use stage-I training. Stage-I training is able to achieve similar performance as the full-precision Stable Diffusion v1.5, as shown in Fig. 5 of the main paper.

To better understand the training cost of our model, we show the training cost comparison between our model and Stable Diffusion v1.5 (the training cost is obtained from the public resource https://huggingface.co/runwayml/stable-diffusion-v1-5). We use the total trained samples to compare the training costs as the training hours depending on the actual hardware used. As shown in the following table, our stage-I training only requires 0.16% training cost compared with training the Stable Diffusion v1.5 from scratch, and stage-II training requires 1.6% training cost compared with training the Stable Diffusion v1.5 from scratch. For the training time, based on the public information (Huggingface), SD-v1.5 is trained for 30-60 days on 256 GPUs, while our model takes 5~6 days on at most 32 GPUs.

| Stage | A100 | Batch size | Gradient accumulation | Total batch size | 256x256 iters (K) | 512x512 iterations (K) | Total training samples (M) | Percentage of total training samples compared to SD-v1.5 | | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | | Stable Diffusion v1.5 | 256 | 4 | 2 | 2048 | 237 | 1304 | 3155.97 | 100% | | Stage-I | 8 | 8 | 4 | 256 | - | 20 | 5.12 | 0.16% | | Stage-II | 32 | 8 | 4 | 1024 | - | 50 | 51.2 | 1.6% |

Q1. About training computation.

A1. Thanks for the suggestions. For the stage-I training, we use 8 NVIDIA A100 GPUs with a total batch size of 256 to train the quantized model for 20K iterations. The training time is within 40 hours. For the stage-II training, we use 32 NVIDIA A100 GPUs with a total batch size of 1024 to train the quantized model for 50K iterations. The total time is within 100 hours.

Note that the stage-II training aims to further improve the performance of the quantized model so that it is better than the full-precision SDv1.5. If reducing the training cost would be the goal, we can remove stage-II training. Stage-I training is able to achieve similar performance as the full-precision SDv1.5 (Fig. 5 of the main paper).

Additionally, we compare training costs between our model and SDv1.5 from scratch (https://huggingface.co/runwayml/stable-diffusion-v1-5). We use the total trained samples to measure the training costs as the training hours depend on the actual hardware used. As in the following table, our stage-I training only requires 0.16% training cost compared with training the SDv1.5, and stage-II training requires 1.6% training cost compared with SDv1.5.

| Stage | A100 | Batch size | Gradient accumulation | Total batch size | 256x256 iters (K) | 512x512 iterations (K) | Total training samples (M) | Percentage of total training samples compared to SD-v1.5 | | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | | Stable Diffusion v1.5 | 256 | 4 | 2 | 2048 | 237 | 1304 | 3155.97 | 100% | | Stage-I | 8 | 8 | 4 | 256 | - | 20 | 5.12 | 0.16% | | Stage-II | 32 | 8 | 4 | 1024 | - | 50 | 51.2 | 1.6% |

Q2.1. Why certain metrics are correlated.

A2.1. Thanks for the insightful comments! MSE, PSNR, and LPIPS are highly correlated since they measure the changes of image quality. Namely, they calculate the perceptual similarity between two images that align with human perception. More specifically, MSE quantifies the pixel-wise difference between images. Similarly, PSNR, derived from MSE, evaluates the peak error in images via $PSNR = 10 * log_{10}(Max^2 / MSE)$. LPIPS has a different calculation as MSE and PSNR. Yet, it still focuses on significant perceptual differences beyond pixel-level comparisons. Given that perceptual changes often accompany measurable pixel-wise changes, LPIPS correlates with PSNR and MSE.

On the other hand, the CLIP score measures the alignment between the image and prompt, leading to a low correlation with MSE, PSNR, and LPIPS metrics.

Q2.2. Why quantization affects some layers more than others.

A2.2. We notice that cross-attention layers and convolutional shortcut layers are more sensitive to low-bit quantization than other layers. Cross-attention layers combine textual and image information. Therefore, any quantization-induced errors for cross-attention layers might greatly change the semantics and content of the generative images, leading to substantial semantic shifts and the loss of important content. For example, in Line 147~149 and Fig. 2 of the main paper, quantizing a cross-attention key layer changes the image content from "a teddy bear" to "a person". Convolutional shortcut layers link different layers and scales of features across the UNet. They facilitate the information flow between layers, enabling the network to learn complex patterns. Quantization of these shortcut connections disrupts the well-established connectivity of features, leading to a breakdown in the coherence of the learned representations (as in Fig. 2).

Q3. Some evaluation datasets, like MS-COCO 2014, may not fully represent the diversity in real-world.

A3. We agree with the reviewer that most of the datasets used for evaluating text-to-image models have their limitations. In this paper, we evaluate our model on most of the commonly used benchmark datasets to get relatively accurate and fair comparisons. Specifically, we use Parti Prompts, TIFA, GenEval scores, and Human evaluation. These benchmark datasets cover a wide variety of prompts and are commonly adopted in evaluating text-to-image models in the literature, like SDXL. Additionally, we would like to kindly clarify that our human preference evaluation includes complex scenes, abstract concepts, and specific artistic styles (as in Fig. 13 and Fig. 14).

Regarding our claim that our quantized model outperforms the full-precision SDv1.5 across all the evaluations, we intend to mention that our BitsFusion can outperform full-precision SDv1.5 across commonly used metrics evaluated in the paper. Based on the feedback from the reviewer, we will modify our writing on Line 61 as our model outperforms the full-precision model across various evaluation metrics used in the paper. We hope it can make the paper more rigorous.

We thank the reviewer for suggesting new metrics. We will discuss them in the revised paper. These metrics have not been widely adopted for evaluating text-to-image generation. For instance, Visual Genome is for visual question answering [A], Flickr30k Entities is for evaluating image description [B], Open Images Dataset V4 is for classification, detection, and visual relationship detection [C], and ADE20K is for understanding scene parsing [E]. Adapting these metrics to evaluate text-to-image models requires a series of experiments to understand how to use them and correlate them with human perception. We take this as a future work.

During the rebuttal, we run the suggested evaluation metric from Conceptual 12M with the CLIP score [D]. By randomly selecting 1K prompts (captions) and using the same CFG scale of 7.5, our 1.99 bits model outperforms full-precision SD-v1.5 (CLIP scores of 0.3084 vs. 0.3067).

Q4. More related methods.

A4. Thanks for suggesting these papers, which are indeed insightful and relevant! We will discuss these papers in the revised paper.

**Q1. Comparison with adaround $1$ , which is used in Diffusion Model quantization research $2$ **.

A1. Thank the reviewer for suggesting the adaround, which is a relevant method $1$ . However, we would like to kindly emphasize that adaround $1$ focuses on post-training quantization by proposing a better weight-rounding operation, while our scaling factor initialization is used for training-aware quantization. These two technologies are applied during different quantization pipelines. As highlighted by Reviewer WzcG, our scaling factor initialization is novel.

Indeed, as mentioned by the reviewer, the Q-Diffusion $2$ uses the adaround, implemented by MinMax and Brute-force MSE searching. As analyzed in the paper (Line 198~201), MinMax is able to maintain the outlier weights but overlook the overall quantization error in the extreme low-bit setting, thus sacrificing the performance of the quantized model. To further prove this, we have conducted an ablation analysis in Table 4 of the main paper. With the alternating optimization (indicated by the row labeled “+Alternating Opt.”), the CLIP score is improved from 0.3054 (indicated by the row labeled “+Balance”, which uses MinMax as the default initialization method) to 0.3100.

As for the brute-force MSE, it takes lots of time to find the ideal scaling factor than our initialization approach. Following the same setting, we measure the CLIP score where the brute-force MSE in Q-Diffusion $2$ can achieve a score of 0.3098, while our alternating optimization achieves 0.3100. However, the brute-force MSE requires 13 minutes to initialize the scaling factor by traversing 80 different scaling factors in each layer. In contrast, our alternating optimization significantly reduces the initialization time, which only needs 5 minutes with 10 optimization iterations.

In fact, we have compared our approach with Q-Diffusion $2$ . As shown in Table 3 of the main paper, our 1.99 bits model can achieve a 0.3175 CLIP score, even outperforming the 4 bits Q-Diffusion model which has a CLIP score of 0.3137 evaluated on the PartiPrompts.

Thanks again for recommending the comparison with adaround. We will revise the manuscript to include the above discussion.

Q2. Quantization of the diffusion transformer (DiT).

A2. We agree with the reviewer that besides UNet, DiT is another promising architecture used as the backbone for the diffusion model. This work targets the quantization of UNet for several reasons.

• First, UNet is a widely adopted architecture for diffusion models. Besides the text-to-image generation, there are numerous applications like ControlNet and IP-Adapter that are built upon the UNet-based diffusion models. However, there is no previous effort to study the extremely low-bit quantization for such an architecture. Therefore, it is important to understand how to quantize the UNet to low-bits.
• Second, UNet is more efficient than DiT during inference, especially on resource-constrained devices. Using quantization to compress the efficient architecture has important practical values, as agreed by Reviewer WzcG and W2xs. Even if we apply the quantization to DiT models, how to deploy such models on resource-constrained devices is unclear, e.g., we are not able to run DiT-based text-to-image models on the iPhone.
• Third, although there is quite some work studying the improvement of DiT-based diffusion models, the UNet-based model still shows competitive performance for text-to-image generation. For instance, the open-sourced Kolors (https://github.com/Kwai-Kolors/Kolors/tree/master?tab=readme-ov-file#-evaluation), which is built upon SDXL, achieves the SOTA results for text-to-image generation.

With all that being said, we believe that our approach can be applied to quantize DiT models, especially considering the quantization error analysis and training pipelines proposed in this paper are generic methods.

Thanks again for suggesting the quantization of DiT, which is an interesting topic that we would like to explore by using our proposed methods.

Q3. Adding the FID comparison to the main paper.

A3. Thanks for the suggestion. We will move the FID comparison from the Appendix to the main paper.

Q4. About the activation quantization for the 1.58-bit LLM quantization approach.

A4. This work mainly targets the size reduction of the diffusion model. Therefore, we did not show the activation quantization. Nevertheless, 8-bit activation quantization can be directly applied to our approach. Based on 1.99 bits weight quantization, we further apply 8-bit activation quantization to all layers except for the first and last convolutional layer and train the corresponding scaling factors. The experimental settings are the same as those described in Section 5 of the main paper. With a CFG scale of 7.5, our stage-I model and stage-II model can achieve CLIP scores of 0.3172 and 0.3193 on the PartiPrompts, respectively. For the TIFA score, our stage-I model and stage-II model can achieve 0.786 and 0.805 with a CFG scale of 7.5. With two-stage training, our BitsFusion model with a 1.99-bit weight and 8-bit activation quantization can still outperform the full-precision SD-v1.5.

| Methods | Weight Bit-width | Activation Bit-width | CLIP score | TIFA score | | :---- | :---- | :---- | :---- | :---- | | Stable Diffusion v1.5 | 32 | 32 | 0.3175 | 0.788 | | Stage-I | 1.99 | 8 | 0.3172 | 0.786 | | Stage-II | 1.99 | 8 | 0.3193 | 0.805 |

$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.

Dear authors,

Thanks for your valuable information. It is worth noting that the findings in prior diffusion model quantization research are mostly based on the U-Net structure, like shortcut splitting quantization in Q-Diffusion [1] and temporal information aware reconstruction in TFMQ-DM [2]. However, the findings in BitFusion, like how to decide the optimal precision, are not limited to the U-Net diffusion model. So I still suggest BitFusion's effectiveness be further verified with DiT as most of the SOTA diffusion models are DiT-based.

Best, Reviewer tjXw

[1] Li, Xiuyu, et al. "Q-diffusion: Quantizing diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023. [2] Huang, Yushi, et al. "Tfmq-dm: Temporal feature maintenance quantization for diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.