VQ4DiT: Эффективная посттренировочная векторная квантизация для трансформеров диффузии

VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

August 30, 2024

Авторы: Juncan Deng, Shuaiting Li, Zeyu Wang, Hong Gu, Kedong Xu, Kejie Huang

cs.AI

Аннотация

Модели трансформеров диффузии (DiTs) перевели архитектуру сети с традиционных UNet на трансформеры, продемонстрировав исключительные возможности в генерации изображений. Хотя DiTs широко применялись в задачах генерации видео высокой четкости, их большой размер параметров затрудняет вывод на периферийных устройствах. Векторное квантование (VQ) может разложить весовую модель на кодовую книгу и назначения, позволяя крайне квантовать веса и значительно сокращать использование памяти. В данной статье мы предлагаем VQ4DiT, быстрый метод векторного квантования после обучения для DiTs. Мы обнаружили, что традиционные методы VQ калибруют только кодовую книгу, не калибруя назначения. Это приводит к неправильному назначению подвекторов весов одному и тому же назначению, обеспечивая несогласованные градиенты для кодовой книги и приводя к неоптимальному результату. Для решения этой проблемы VQ4DiT вычисляет набор кандидатов для каждого подвектора веса на основе евклидового расстояния и восстанавливает подвектор на основе взвешенного среднего. Затем, используя метод калибровки нулевых данных и блоковую калибровку, оптимальное назначение из набора эффективно выбирается при калибровке кодовой книги. VQ4DiT квантует модель DiT XL/2 на одном графическом процессоре NVIDIA A100 за 20 минут до 5 часов в зависимости от различных настроек квантования. Эксперименты показывают, что VQ4DiT устанавливает новый уровень в компромиссе между размером модели и производительностью, квантуя веса с точностью 2 бита, сохраняя приемлемое качество генерации изображений.

English

The Diffusion Transformers Models (DiTs) have transitioned the network architecture from traditional UNets to transformers, demonstrating exceptional capabilities in image generation. Although DiTs have been widely applied to high-definition video generation tasks, their large parameter size hinders inference on edge devices. Vector quantization (VQ) can decompose model weight into a codebook and assignments, allowing extreme weight quantization and significantly reducing memory usage. In this paper, we propose VQ4DiT, a fast post-training vector quantization method for DiTs. We found that traditional VQ methods calibrate only the codebook without calibrating the assignments. This leads to weight sub-vectors being incorrectly assigned to the same assignment, providing inconsistent gradients to the codebook and resulting in a suboptimal result. To address this challenge, VQ4DiT calculates the candidate assignment set for each weight sub-vector based on Euclidean distance and reconstructs the sub-vector based on the weighted average. Then, using the zero-data and block-wise calibration method, the optimal assignment from the set is efficiently selected while calibrating the codebook. VQ4DiT quantizes a DiT XL/2 model on a single NVIDIA A100 GPU within 20 minutes to 5 hours depending on the different quantization settings. Experiments show that VQ4DiT establishes a new state-of-the-art in model size and performance trade-offs, quantizing weights to 2-bit precision while retaining acceptable image generation quality.