BitsFusion: 拡散モデルの1.99ビット重み量子化

BitsFusion: 1.99 bits Weight Quantization of Diffusion Model

June 6, 2024

著者: Yang Sui, Yanyu Li, Anil Kag, Yerlan Idelbayev, Junli Cao, Ju Hu, Dhritiman Sagar, Bo Yuan, Sergey Tulyakov, Jian Ren

cs.AI

要旨

拡散モデルに基づく画像生成技術は近年、高品質なコンテンツ合成能力を示すことで大きな成功を収めています。しかし、これらのモデルは膨大なパラメータ数を有しており、結果として非常に大きなモデルサイズとなっています。そのため、保存や転送が主要なボトルネックとなっており、特にリソースが制約されたデバイス上での運用において大きな課題となっています。本研究では、Stable Diffusion v1.5のUNetを1.99ビットに量子化する新たな重み量子化手法を開発し、モデルサイズを7.9倍小さくしながら、元のモデルよりも優れた生成品質を実現しました。私たちのアプローチでは、各層に最適なビット数を割り当てる、量子化モデルをより良い性能で初期化する、量子化誤差を大幅に低減するための学習戦略を改善するなど、いくつかの新技術を導入しています。さらに、量子化モデルを様々なベンチマークデータセットで広範に評価し、人間による評価を通じてその優れた生成品質を実証しました。

English

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.