BitsFusion: 1.99 bits Weight Quantization of Diffusion Model

TL;DR

BitsFusion tackles the large storage footprint of diffusion models by introducing mixed-precision quantization to the Stable Diffusion v1.5 UNet, achieving $1.99$-bit weights and a $7.9\times$ reduction in size with maintained or improved image quality. It combines per-layer error analysis, a parameter-size-aware bit allocation strategy, initialization techniques (time-embedding caching and balance integers), and a two-stage training pipeline that includes CFG-aware distillation and time-step sampling adjustments. The approach is validated extensively on benchmarks like MS-COCO, TIFA, GenEval, and human assessments, showing BitsFusion can outperform the full-precision model after fine-tuning and is robust to different sampling strategies. Practically, this enables efficient deployment of high-quality diffusion models on resource-constrained devices while preserving semantic alignment and visual fidelity.

Abstract

Diffusion-based image generation models have achieved great success in recent years by showing the capability of synthesizing high-quality content. However, these models contain a huge number of parameters, resulting in a significantly large model size. Saving and transferring them is a major bottleneck for various applications, especially those running on resource-constrained devices. In this work, we develop a novel weight quantization method that quantizes the UNet from Stable Diffusion v1.5 to 1.99 bits, achieving a model with 7.9X smaller size while exhibiting even better generation quality than the original one. Our approach includes several novel techniques, such as assigning optimal bits to each layer, initializing the quantized model for better performance, and improving the training strategy to dramatically reduce quantization error. Furthermore, we extensively evaluate our quantized model across various benchmark datasets and through human evaluation to demonstrate its superior generation quality.

Figures (23)

• Figure 1: Top: Images generated from full-precision Stable Diffusion v1.5. Bottom: Images generated from BitsFusion, where the weights of UNet are quantized into 1.99 bits, achieving $7.9\times$smaller storage than the one from Stable Diffusion v1.5. All the images are synthesized under the setting of using PNDM sampler pndm with $50$ sampling steps and random seed as $1024$. Prompts and more generations are provided in App. \ref{['app:more_images']}.
• Figure 2: $1$-bit quantization error analysis for all the layers from the UNet of SD-v1.5.
• Figure 3: Overview of the training and inference pipeline for the proposed BitsFusion.Left: We analyze the quantization error for each layer in SD-v1.5 (Sec. \ref{['sec:mixed_analysis_per_layer']}) and derive the mixed-precision recipe (Sec. \ref{['sec:deciding-precision']}) to assign different bit widths to different layers. We then initialize the quantized UNet by adding a balance integer, pre-computing and caching the time embedding, and alternately optimizing the scaling factor (Sec. \ref{['sec:initialization']}). Middle: During the Stage-I training, we freeze the teacher model (i.e., SD-v1.5) and optimize the quantized UNet through CFG-aware quantization distillation and feature distillation losses, along with sampling time steps by considering quantization errors (Sec. \ref{['sec:multi-stage-train']}). During the Stage-II training, we fine-tune the previous model with the noise prediction. Right: For the inference stage, using the pre-cached time features, our model processes text prompts and generates high-quality images.
• Figure 4: More time steps are sampled towards where larger quantization error occurs.
• Figure 5: Comparison between our $1.99$-bits model vs. SD-v1.5 on various evaluation metrics with CFG scales ranging from $2.5$ to $9.5$. Ours-I denotes the model with Stage-I training and Ours-II denotes the model with Stage-II training.