VQ4DiT: Quantização Vetorial Eficiente Pós-Treinamento para Transformadores de Difusão

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

Resumo

Os Modelos de Transformadores de Difusão (DiTs) fizeram a transição da arquitetura de rede dos UNets tradicionais para os transformers, demonstrando capacidades excepcionais na geração de imagens. Embora os DiTs tenham sido amplamente aplicados em tarefas de geração de vídeo em alta definição, seu grande tamanho de parâmetro dificulta a inferência em dispositivos de borda. A quantização vetorial (VQ) pode decompor o peso do modelo em um livro de códigos e atribuições, permitindo uma quantização extrema do peso e reduzindo significativamente o uso de memória. Neste artigo, propomos o VQ4DiT, um método de quantização vetorial pós-treinamento rápido para os DiTs. Descobrimos que os métodos tradicionais de VQ calibram apenas o livro de códigos sem calibrar as atribuições. Isso faz com que subvetores de peso sejam atribuídos incorretamente à mesma atribuição, fornecendo gradientes inconsistentes para o livro de códigos e resultando em um resultado subótimo. Para enfrentar esse desafio, o VQ4DiT calcula o conjunto de atribuições candidatas para cada subvetor de peso com base na distância euclidiana e reconstrói o subvetor com base na média ponderada. Em seguida, usando o método de calibração por dados zero e por blocos, a atribuição ótima do conjunto é selecionada de forma eficiente enquanto calibra o livro de códigos. O VQ4DiT quantiza um modelo DiT XL/2 em uma única GPU NVIDIA A100 em 20 minutos a 5 horas, dependendo das diferentes configurações de quantização. Experimentos mostram que o VQ4DiT estabelece um novo estado da arte em trocas de tamanho de modelo e desempenho, quantizando pesos com precisão de 2 bits enquanto mantém uma qualidade aceitável na geração de imagens.

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.