VQ4DiT: Cuantificación Vectorial Eficiente Después del Entrenamiento para Transformadores de Difusión

VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

August 30, 2024

Autores: Juncan Deng, Shuaiting Li, Zeyu Wang, Hong Gu, Kedong Xu, Kejie Huang

cs.AI

Resumen

Los Modelos de Transformadores de Difusión (DiTs) han transitado la arquitectura de red desde los UNets tradicionales a los transformadores, demostrando capacidades excepcionales en la generación de imágenes. Aunque los DiTs se han aplicado ampliamente a tareas de generación de videos de alta definición, su gran tamaño de parámetros dificulta la inferencia en dispositivos periféricos. La cuantización vectorial (VQ) puede descomponer el peso del modelo en un libro de códigos y asignaciones, permitiendo una cuantización extrema del peso y reduciendo significativamente el uso de memoria. En este documento, proponemos VQ4DiT, un método de cuantización vectorial rápido post-entrenamiento para DiTs. Descubrimos que los métodos tradicionales de VQ solo calibran el libro de códigos sin calibrar las asignaciones. Esto lleva a que los subvectores de peso se asignen incorrectamente a la misma asignación, proporcionando gradientes inconsistentes al libro de códigos y dando como resultado un resultado subóptimo. Para abordar este desafío, VQ4DiT calcula el conjunto de asignaciones candidatas para cada subvector de peso basado en la distancia euclidiana y reconstruye el subvector basado en el promedio ponderado. Luego, utilizando el método de calibración de datos cero y por bloques, se selecciona eficientemente la asignación óptima del conjunto mientras se calibra el libro de códigos. VQ4DiT cuantiza un modelo DiT XL/2 en una sola GPU NVIDIA A100 en 20 minutos a 5 horas dependiendo de los diferentes ajustes de cuantización. Los experimentos muestran que VQ4DiT establece un nuevo estado del arte en el equilibrio entre el tamaño del modelo y el rendimiento, cuantizando pesos a una precisión de 2 bits mientras se mantiene una calidad aceptable en la generación de imágenes.

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.