BitsFusion: Quantizzazione a 1.99 bit dei pesi nel modello di diffusione

BitsFusion: 1.99 bits Weight Quantization of Diffusion Model

June 6, 2024

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

cs.AI

Abstract

I modelli di generazione di immagini basati su diffusione hanno ottenuto un grande successo negli ultimi anni, dimostrando la capacità di sintetizzare contenuti di alta qualità. Tuttavia, questi modelli contengono un numero enorme di parametri, risultando in dimensioni del modello significativamente grandi. Salvarli e trasferirli rappresenta un collo di bottiglia per varie applicazioni, specialmente quelle eseguite su dispositivi con risorse limitate. In questo lavoro, sviluppiamo un nuovo metodo di quantizzazione dei pesi che quantizza la UNet di Stable Diffusion v1.5 a 1,99 bit, ottenendo un modello con dimensioni 7,9 volte più piccole e una qualità di generazione persino migliore rispetto a quella originale. Il nostro approccio include diverse tecniche innovative, come l'assegnazione di bit ottimali a ciascun livello, l'inizializzazione del modello quantizzato per migliorare le prestazioni e il perfezionamento della strategia di addestramento per ridurre drasticamente l'errore di quantizzazione. Inoltre, valutiamo ampiamente il nostro modello quantizzato su vari dataset di benchmark e attraverso valutazioni umane per dimostrare la sua qualità di generazione superiore.

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.