BitsFusion: 1.99 bits Weight Quantization of Diffusion Model
Authors: Yang Sui, Yanyu Li, Anil Kag, Yerlan Idelbayev, Junli Cao, Ju Hu, Dhritiman Sagar, Bo Yuan, Sergey Tulyakov, Jian Ren
NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | 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. [...] 5 Experiments |
| Researcher Affiliation | Collaboration | 1Snap Inc. 2Rutgers University |
| Pseudocode | Yes | We provide the detailed algorithm as outlined in Alg. 1. |
| Open Source Code | Yes | Project Page: https://snap-research.github.io/Bits Fusion [...] We plan to release our code and trained models to facilitate the research efforts towards extreme low-bits quantization. |
| Open Datasets | Yes | We include results on various benchmark datasets, i.e., TIFA [25], Gen Eval [13], CLIP score [66] and FID [19] on MS-COCO 2014 validation set [46]. Additionally, we perform human evaluation on Parti Prompts [86]. |
| Dataset Splits | Yes | We perform QAT over each candidate on a pre-defined training sub dataset, and validate the incurred quantization error of each candidate by comparing it against the full-precision model (more details in App. B). |
| Hardware Specification | Yes | For Stage-I, we use 8 NVIDIA A100 GPUs with a total batch size of 256 to train the quantized model for 20K iterations. For Stage-II, we use 32 NVIDIA A100 GPUs with a total batch size of 1024 to train the quantized model for 50K iterations. |
| Software Dependencies | No | The paper mentions using the 'diffusers library' and 'Adam W optimizer' but does not specify version numbers for these or other software dependencies. |
| Experiment Setup | Yes | We develop our code using diffusers library2 and train the models with Adam W optimizer [33] and a constant learning rate as 1e 05 on an internal dataset. For Stage-I, we use 8 NVIDIA A100 GPUs with a total batch size of 256 to train the quantized model for 20K iterations. For Stage-II, we use 32 NVIDIA A100 GPUs with a total batch size of 1024 to train the quantized model for 50K iterations. During inference, we adopt the PNDM scheduler [49] with 50 sampling steps to generate images for comparison. |