BitsFusion: Cuantización de pesos de modelos de difusión a 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

Resumen

Los modelos de generación de imágenes basados en difusión han logrado un gran éxito en los últimos años al demostrar la capacidad de sintetizar contenido de alta calidad. Sin embargo, estos modelos contienen una enorme cantidad de parámetros, lo que resulta en un tamaño de modelo significativamente grande. Guardar y transferirlos es un cuello de botella importante para diversas aplicaciones, especialmente aquellas que se ejecutan en dispositivos con recursos limitados. En este trabajo, desarrollamos un novedoso método de cuantización de pesos que cuantiza el UNet de Stable Diffusion v1.5 a 1.99 bits, logrando un modelo 7.9 veces más pequeño mientras exhibe una calidad de generación incluso mejor que el original. Nuestro enfoque incluye varias técnicas novedosas, como asignar bits óptimos a cada capa, inicializar el modelo cuantizado para un mejor rendimiento y mejorar la estrategia de entrenamiento para reducir drásticamente el error de cuantización. Además, evaluamos exhaustivamente nuestro modelo cuantizado en varios conjuntos de datos de referencia y a través de evaluaciones humanas para demostrar su calidad de generación 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.