BitsFusion: Quantização de Pesos de Modelos de Difusão em 1.99 bits

BitsFusion: 1.99 bits Weight Quantization of Diffusion Model

June 6, 2024

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

cs.AI

Resumo

Modelos de geração de imagens baseados em difusão têm alcançado grande sucesso nos últimos anos ao demonstrar a capacidade de sintetizar conteúdo de alta qualidade. No entanto, esses modelos contêm um número enorme de parâmetros, resultando em um tamanho de modelo significativamente grande. Salvar e transferi-los é um grande gargalo para diversas aplicações, especialmente aquelas executadas em dispositivos com recursos limitados. Neste trabalho, desenvolvemos um novo método de quantização de pesos que quantiza o UNet do Stable Diffusion v1.5 para 1,99 bits, alcançando um modelo 7,9 vezes menor em tamanho, ao mesmo tempo que exibe uma qualidade de geração ainda melhor do que o original. Nossa abordagem inclui várias técnicas inovadoras, como a atribuição de bits ótimos para cada camada, a inicialização do modelo quantizado para melhor desempenho e a melhoria da estratégia de treinamento para reduzir drasticamente o erro de quantização. Além disso, avaliamos extensivamente nosso modelo quantizado em diversos conjuntos de dados de referência e por meio de avaliação humana para demonstrar sua qualidade de geração superior.

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.