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)는 모델 가중치를 코드북과 할당으로 분해하여 극도의 가중치 양자화를 허용하고 메모리 사용량을 크게 줄일 수 있습니다. 본 논문에서는 DiTs를 위한 빠른 사후 훈련 벡터 양자화 방법인 VQ4DiT을 제안합니다. 우리는 기존의 VQ 방법이 코드북만 보정하고 할당을 보정하지 않는 것을 발견했습니다. 이는 가중치 하위 벡터가 잘못된 할당에 배정되어 코드북에 일관되지 않은 그래디언트를 제공하고 최적의 결과를 얻지 못하게 합니다. 이러한 도전에 대처하기 위해 VQ4DiT은 유클리드 거리를 기반으로 각 가중치 하위 벡터에 대한 후보 할당 세트를 계산하고 가중 평균을 기반으로 하위 벡터를 재구성합니다. 그런 다음, 제로 데이터 및 블록별 보정 방법을 사용하여 세트에서 최적의 할당을 효율적으로 선택하면서 코드북을 보정합니다. VQ4DiT은 다양한 양자화 설정에 따라 단일 NVIDIA A100 GPU에서 20분에서 5시간이 소요되는 시간 내에 DiT XL/2 모델을 양자화합니다. 실험 결과, 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.