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.