VQ4DiT: Efficient Post-Training Vector Quantization for Diffusion Transformers

本文介绍了一种名为Diffusion Transformers Models (DiTs)的网络结构，它将传统的UNet转换为transformers，并在图像生成方面表现出色。虽然DiTs已被广泛应用于高清视频生成任务，但其大量的参数大小阻碍了在边缘设备上的推理。向量量化（VQ）可以将模型权重分解为码本和分配，从而实现极端权重量化并显著减少内存使用。本文提出了一种名为VQ4DiT的快速后训练向量量化方法，用于DiTs。我们发现传统的VQ方法只校准码本而不校准分配。这导致权重子向量被错误地分配到相同的分配中，为码本提供不一致的梯度，导致结果次优。为了解决这个挑战，VQ4DiT基于欧几里得距离计算每个权重子向量的候选分配集，并基于加权平均值重构子向量。然后，使用零数据和块状校准方法，在校准码本的同时高效地选择最佳分配。VQ4DiT在单个NVIDIA A100 GPU上对DiT XL/2模型进行量化，根据不同的量化设置，量化时间为20分钟至5小时不等。实验表明，VQ4DiT在模型大小和性能权衡方面建立了新的最优状态，将权重量化到2位精度，同时保持可接受的图像生成质量。

本文旨在解决DiTs模型在边缘设备上推理时参数量大的问题，提出了一种快速的后训练向量量化方法VQ4DiT。

VQ4DiT通过计算每个权重子向量的候选分配集合，并基于加权平均重构子向量，然后使用零数据和块状校准方法，高效地选择最优分配并校准码本，解决了传统VQ方法只校准码本而不校准分配的问题。

VQ4DiT在单个NVIDIA A100 GPU上将DiT XL/2模型量化为2位精度，保持了可接受的图像生成质量，同时在模型大小和性能之间建立了新的权衡。实验结果表明，VQ4DiT在模型大小和性能之间建立了新的最优点，同时开源了代码。

在这个领域中，最近的相关研究包括《High-Fidelity Image Generation With Fewer Labels》、《Generative Adversarial Networks》等。