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.