扩散模型低位量化突破！有效扩散量化的极限推向2-4位，W2A4位宽下FID降低58%，超越SOTA方法

中科院提出混合精度量化

原标题：扩散模型低位量化突破！有效扩散量化的极限推向2-4位，W2A4位宽下FID降低58%，超越SOTA方法
文章来源：量子位
内容字数：11110字

极低位扩散模型量化：MPQ-DM方法详解

本文总结了MPQ-DM团队发表在arXiv上的最新研究，该研究提出了一种混合精度量化方法MPQ-DM，用于降低扩散模型的计算成本，同时保持高水平的性能。现有扩散模型量化方法在极低位宽（2-4位）下性能严重下降，主要原因是激活值的显著离散化。MPQ-DM通过结合离群值驱动的混合量化(OMQ)和时间平滑关系蒸馏(TRD)两种技术来解决这个问题。

1. 扩散模型与量化

扩散模型通过逐步添加噪声然后去除噪声来生成图像。然而，其计算成本很高。量化通过降低参数的位宽来节省存储和计算资源，但现有方法在低位宽下性能下降严重。

2. MPQ-DM的核心技术

MPQ-DM的核心在于OMQ和TRD两项技术：

1. 离群值驱动的混合量化 (OMQ): 该方法利用峰度(kurtosis)来识别权重通道中的离群值。离群值显著的通道被分配更高的位宽，而其他通道分配较低的位宽，从而在保持平均位宽不变的情况下提高精度。这种层内混合精度量化避免了传统逐层分配方法的不足。
2. 时间平滑关系蒸馏 (TRD): 为了解决低位量化导致的特征表示不一致性，TRD在量化模型和全精度模型之间构建了一个时间平滑的关系蒸馏方案。它通过融合多个连续时间步长的中间特征，并使用KL散度来衡量特征相似性分布之间的差异，而不是直接比较数值，从而提高了模型的鲁棒性。

3. 实验结果

实验结果表明，MPQ-DM在LSUN-Bedrooms、LSUN-Churches和ImageNet数据集上，以及Stable Diffusion模型上都显著优于现有方法。尤其是在极低位宽(例如W2A4)下，MPQ-DM取得了巨大的性能提升，而其他方法甚至无法生成正常的图像。消融实验也验证了OMQ和TRD的有效性。

4. 主要发现

研究发现：

• 层内混合精度量化比逐层量化更有效。
• 基于峰度的离群值选择方法优于随机选择方法。
• 关系蒸馏比直接数值对齐更适合处理离散和连续特征的差异。

5. 总结

MPQ-DM通过巧妙地结合OMQ和TRD，有效地解决了极低位扩散模型量化中的性能下降问题，为扩散模型在资源受限场景中的应用提供了新的可能性。该方法在多个数据集和模型上的优异表现，以及全面的消融实验结果，都证明了其有效性和先进性。

联系作者

文章来源：量子位
作者微信：
作者简介：追踪人工智能新趋势，关注科技行业新突破