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

Аннотация

Модели генерации изображений на основе диффузии в последние годы достигли больших успехов, показав способность синтезировать контент высокого качества. Однако эти модели содержат огромное количество параметров, что приводит к значительному увеличению размера модели. Сохранение и передача их являются основными узкими местами для различных приложений, особенно тех, которые работают на устройствах с ограниченными ресурсами. В данной работе мы разрабатываем новый метод квантования весов, который квантует UNet из Stable Diffusion с 1.5 по 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.