浙江大学和vivo的研究人员推出一种适用于 DiTs 的快速后训练向量量化方法 VQ4DiT，它是一种针对扩散变换器模型（Diffusion Transformers，简称DiTs）的高效后训练矢量化方法。DiTs是一种在图像生成领域表现出色的模型，但它们庞大的参数量使得在边缘设备上进行推理变得困难。VQ4DiT通过将模型权重分解成码本和分配，实现了极端的权重量化，显著减少了内存使用量。
研究人员发现传统的 VQ 方法只校准代码本而不校准索引值，这会导致权重的子向量被错误地指向同一个索引值，从而向代码本传递了不一致的梯度，造成次优的结果。为了克服这个难题，VQ4DiT 根据欧氏距离计算每个权重子向量的候选索引值集，并基于加权平均重构子向量。接着，通过无数据和分块校准的方法，从候选集中高效地选出最佳索引值，并同时校准代码本。 VQ4DiT 能够在一个英伟达A100 GPU 上，根据不同的量化设置，在 20 分钟到 5 小时内完成对 DiT XL/2 模型的量化。实验结果显示，VQ4DiT 在模型大小与性能的平衡方面达到了新的最先进水平，实现了 2 位精度的量化，同时保持了图像生成的质量。
主要功能：
VQ4DiT的主要功能是对DiTs模型进行后训练量化，以减少模型大小并提高在资源受限设备上的推理效率，同时尽量保持图像生成质量。
主要特点：
- 极低比特宽度量化： VQ4DiT能够将模型权重量化到2比特精度，这在以往的研究中是具有挑战性的。
- 快速量化： 该方法能够在单个NVIDIA A100 GPU上在20分钟到5小时内完成DiT XL/2模型的量化。
- 无需校准数据集： VQ4DiT采用了零数据和分块校准方法，无需额外的校准数据集即可实现模型量化。
工作原理：
VQ4DiT的工作原理包括以下几个步骤：
- 初始化码本和候选分配集： 使用K-Means算法将模型权重映射到初始码本和分配。
- 计算候选分配集： 对于每个权重子向量，基于欧几里得距离计算其候选分配集。
- 重建权重子向量： 基于候选分配集的加权平均值重建权重子向量。
- 零数据和分块校准： 利用零数据和分块校准方法，同时校准码本和候选分配集的比率，最终选择最优分配。
具体应用场景：
- 边缘设备部署： 在智能手机、嵌入式系统等边缘设备上部署DiTs，用于图像和视频生成。
- 资源受限环境： 在计算资源有限的环境中，如某些云计算场景或物联网设备，使用VQ4DiT来部署高效的图像生成模型。
- 实时图像处理： 在需要实时图像处理的应用中，如视频监控、自动驾驶，使用量化后的模型以减少延迟。
评论0