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

wsx

降低扩散模型生成的计算成本，性能还保持在高水平！
最新研究提出一种用于极低位差分量化的混合精度量化方法。

△图1 W2A6位宽下基线和MPQ-DM生成的样本可视化
目前，扩散模型在生成任务中受到了广泛关注。然而，昂贵的计算成本阻碍了扩散模型在资源受限场景中的应用。
量化是一种实用的解决方案，通过减少参数的位宽来显著节省存储和计算。然而，现有的扩散模型量化方法仍然会导致性能严重下降，特别是在极低的位宽（2-4位）下。
性能的主要下降来自低位量化时激活值的显著离散化。太少的激活候选对离群值显著权重通道量化不友好，离散特征阻碍了扩散模型在不同时间步长上的稳定学习。
为此，来自中国科学院计算技术研究所、苏黎世联邦理工学院、北京交通大学的研究人员提出了由离群值驱动的混合量化（OMD）和时间平滑关系蒸馏（TRD）组成的极低位扩散模型（MPQ-DM）的混合精度量化。


△图2 MPQ-DM概述
MPQ-DM主要依赖于两种技术：

• 为减轻异常值权重信道引起的量化误差，团队提出了一种异常值驱动的混合量化（OMQ）技术，该技术使用峰度来量化异常值显著信道，并应用优化的层内混合精度比特宽度分配，以在目标效率范围内恢复精度性能。
  
• 为了鲁棒地学习跨越时间步长的表示，团队在量化扩散模型与其全精度对应物之间构建了一个时间平滑关系蒸馏（TRD）方案，将离散和连续的潜势转移到统一的关系空间，以减少表示的不一致性。
  

综合实验表明，与SOTA量化方法相比，MPQ-DM在极低的比特宽度下实现了显著的精度增益。与基线相比，在W2A4设置下，MPQ-DM的FID降低了58%，而所有其他方法甚至无法生成正常的图像。

相关工作

扩散模型

扩散模型通过逐步向数据分布x0∼q（x）添加噪声来执行前向采样过程。在DDPM中，扩散模型的前向噪声添加过程是马尔可夫链，其形式为：

其中αt=1−βt , βt 是时间相关的时间表。扩散模型通过对随机采样的高斯噪声 xT∼N(0, I)应用去噪过程来生成高质量图像，形式如下：

扩散模型量化

训练后量化（PTQ）和量化感知训练（QAT）是模型量化的两种主要方法。
常用的QAT方法，如LSQ和扩散模型Q-dm和Binarydm的方法，确保了模型在极低比特宽度甚至二进制量化下的性能，但与PTQ方法相比，它们需要大量额外的训练时间，从而导致更大的训练负担。对扩散模型PTQ4DM和Q-diffusion的PTQ方法进行了初步探索。
以下工作PTQ-D，TFMQ-DM，APQ-DM和QuEST在量化误差方向、瞬态特征、校准数据和校准模块方面进行了改进。量化后的扩散模型的性能得到了进一步的提升。
然而，基于PTQ的方法在极低的比特宽度下性能严重下降。为了结合QAT的优点并减少所需的训练时间，EfficientDM使用LoRA方法对量化扩散模型进行微调。然而，这两种有效的量化方法都不能保证差分模型在低位下的性能。因此，本文的重点是最大限度地提高极低位量化扩散模型的性能。
方法

模型量化

模型量化将模型权重和激活映射到低位整数值，以减少内存占用并加速推理。对于浮动向量xf，量化过程可以公式化为：

其中 表示整数中的量化向量，⌊·⌉是舍入函数，clip(·)是将值限制在[0，2N−1]范围内的函数，s是比例因子，z是量化零点。l和u分别是量化阈值的下限和上限。它们由xf和目标位宽决定。相反，为了将低位整数量化向量 恢复到全精度表示，反量化过程公式为：

其中 是用于正向处理的去量化向量。
离群值指导的混合精度量化

扩散模型的一些研究发现，存在一些具有显著权重值的通道，这些通道对量化误差和模型性能也至关重要。
团队在图3a中可视化了扩散模型的一层权重分布，在图3b中可视化了其对应的异常显著通道分布。可以发现，不仅不同通道的权重在数值范围内存在较大差距，而且一些权重通道也存在严重的异常值。


△图3 LDM-4 ImageNet 256x256的权重分布
在图3c和图3d中，团队可视化了两个通道的权重分布。
总体权重分布可以被视为正态分布加上异常值分布。对于模型量化，正态分布的直接量化得到了很好的研究，但少量异常值的存在可能会导致一些异常值被箝位或占据部分比特宽度，从而难以量化主正态分布，导致量化误差。
这种误差在极低的位宽（2-4位）下会大大放大，因此团队自然希望为这些通道分配更多的位宽，以确保它们的性能。然而，在特定通道中的这种异常现象可能发生在模型的各个层。
因此，使用传统方法将不同的比特分配给不同的模型层进行混合精度量化显然不是一个好的选择。但是由于权重的量化是逐信道的，团队可以在层内的不同信道之间重新分配量化比特宽度来解决这个问题，即层内混合精度量化。
为了确定具体的比特宽度分配方法，需要量化输出的重要性。
峰度κ可用于量化实值变量概率分布的“尾性”。这自然与团队的目标相一致，因为异常值将异常的长尾添加到正态分布中，如上所述。如图3c和图3d所示，与没有异常值的正态分布相比，具有更显著异常值现象的通道具有更高的κ值。因此，团队使用峰度κ作为量化不同权重通道难度的指标，为每层权重Wf找到最佳比特分配方法。
团队希望在混合比特宽度分配量化后，模型输出尽可能与全精度模型保持一致。由于团队的目标是模型输出Y=XWT，并且异常值是权重中的最大值或最小值，因此团队可以在量化之前使用矩阵乘法的特性来降低异常值的显著程度，而不会造成损失，如下所示：

其中Y、X和W表示输出激活、输入激活和模型权重。δ是一个通道平滑因子，通过将异常值缩放到更接近正态分布来平衡权重和激活的量化难度。经过预缩放后，离群显著通道更平滑以进行量化。团队为每个权重通道计算κ，并相应地对其进行排序。然后，团队使用以下优化公式来确定异常突出和非客户端通道，如下所示：

其中N是目标平均比特宽度，n是权重中的通道数。  表示量化和去量化过程。
团队采用通道混合比特量化为 ，其中ci表示第i个信道的比特宽度。C表示信道共享相同位宽的集合，|C|表示集合C中的信道数量。例如，对于3位量化，团队将一些异常显著信道分配给4位，并将相同数量的异常非客户端信道重新分配给2位。这样，层权重的平均总位宽仍然是3位，而不添加任何额外参数。
为了加速优化过程，团队将k个通道设置为一个搜索组，离群显著通道的搜索区域约束在 中。团队根据经验将k设置为 ，因此搜索时间是5倍，这对于层比特分配来说足够有效。
时间步平滑的关系蒸馏

对于扩散优化，现有的方法，如EfficientDM，通过对齐全精度（FP）模型和量化模型的输出来优化量化模型的训练参数:

其中θf和θq分别表示FP模型和量化模型，xt是通过用FP模型迭代T−t步对高斯噪声xT∼N(0, I)进行去噪而获得的。


△图4 从timesetp（左）T、（左中）T-1、（右中）T-2、（右）T-10的特征图可视化


△图5 （左）FP模型、（中）训练好的量化模型、（右）未训练量化模型的不同激活图的可视化
在极低位量化中，仅对齐模型的最终输出是不够的，因为量化模型的表达能力严重不足。团队通常将量化模型的最终投影层设置为8位，并将投影层之前的层设置为低位。这样，团队就不会直接感知到量化信息严重缺失的部分。因此，团队可以通过在最后一个项目层之前提取模型特征层，为极低位量化模型提供更细粒度的指导

其中Ff和Fq分别表示FP模型和量化模型的最后一个项目层之前的特征图。D是测量两个特征图距离的度量。
然而，极低位量化的特征呈现出高度的离散化（例如，4位激活量化只有16个值）。
这种离散化特征在数值上显示出高度的不真实性，直接提取它效果不佳，甚至会干扰量化模型的正常训练。此外，由于扩散模型的独特迭代去噪过程，特征的输出将继续与去噪相加，进一步放大这种不真实性。
通常，团队可以使用额外的投影头将它们映射到均匀空间或对其进行数值正则化。然而，投影头会带来额外的训练参数，进行正则化并不能解决迭代计算中的误差累积问题。团队发现扩散模型的特征与时间步长高度相关，并在某些时间步长中表现出相似性。
在图4中，团队发现连续时间步长上的特征高度相似，而相距甚远的时间步长则截然不同。特征的相似性表明，去噪轨迹中连续时间步长的差异很小，因此团队可以通过融合多个连续步骤的中间特征来缓解不同时间步长的特征不真实性。
而不是强迫高度离散化的量化模型严格学习FP模型每一步的去噪轨迹，而是学习连续多步的去噪声轨迹。
因此，团队使用N个连续步骤的中间特征作为平滑特征表示进行提取。对于T时间步长优化，团队将蒸馏公式改写为

其中FT−t,f 和FT−t,q 分别表示FP模型和量化模型在时间步长T-t中的最后一个特征图。
在图5a中，团队可视化了FP模型、训练好的量化模型和未训练的量化模型之间的特征图。
尽管时间平滑特征提高了量化特征的鲁棒性，但训练有素的模型和FP模型之间的数值表达式仍然存在不匹配。这被归咎于量化特征的离散潜空间和FP特征的连续潜空间之间的差异。因此，任何指标，如L2损耗数字对齐 和 无法避免这种不同空间之间的差异。
因此，建议使用关系对齐来代替严格的学习数值对齐 和 之间的特征相似关系。
在图5b，团队将离散潜空间和连续潜空间之间的数值关系转化为每个空间内的特征相似关系，将蒸馏目标统一到特征相似空间中。
这成功地解决了数值不匹配的问题。
形式上说，对于 ，为了简化写作，团队对其进行了重写 其中s=h×w。对于第i个特征表示 ，团队可以计算它的余弦相似性分布 因此，关系蒸馏度量为

其中 和 表示时间平滑特征相似度分布分别来自FP模型和量化模型。Dkl(·||·)代表两个分布之间的Kullback-Leibler（KL）散度。
团队在这里使用KL散度而不是L2损失。因为L2损失只能感知单个表示，而KL散度可以感知整个特征图的信息。总体优化目标为

实验

实验设置

团队在常用的数据集LSUN-Bedrooms 256×256、LSUN-Churchs 256×256和ImageNet 256×256上进行了LDM模型上无条件和有条件图像生成任务的实验。
还对Stable Diffusion进行了文本到图像的生成任务，使用IS、FID、sFID和Precision来评估LDM的性能。对于Stable Diffusion，团队使用CLIP评分进行评估。
为了应对2bit量化下的极端压缩性能下降，团队在OMQ的搜索过程中为2bit额外分配了10%的信道，称为MPQ-DM+。
与FP模型相比，模型尺寸仅增加了0.6%。团队将MPQ-DM与基线EfficientDM和分层混合精度HAWQ-v3以及其他具有相似耗时的基于PTQ的方法进行了比较。
实验结果

ImageNet 256×256上全量化LDM-4模型的性能比较：


△表1 最佳结果以粗体显示，次佳结果以下划线显示
LDM-4模型在LSUN卧室256×256上的无条件图像生成结果：


△表 2
LDM-8模型在LSUN Churches 256×256上的无条件图像生成结果：


△表 3
类条件生成
团队以LDM-4为重点，在ImageNet 256×256数据集上进行了条件生成实验。
表1中的结果显示，MPQ-DM在所有比特设置上都大大优于现有方法。MPQ-DM通常比分层方法HAWQ-v3表现更好，这证明了在层内进行混合精度量化的必要性。
MPQ-DM W3A4型号在FID上甚至超过了FP型号。在W2A4设置中，基于PTQ的方法无法生成图像，而EfficientDM的性能较差。MPQ-DM大大改善了基线，FID显著降低了27.86。MPQ-DM+仅使用4.8 MB的额外模型大小，就进一步导致FID减少9.48。
使用10k COCO2014验证集提示的Stable Diffusion v1.4的文本到图像生成结果（512×512）：


△表4
无条件的生成
团队在LDM-4的LSUN卧室数据集和LDM-8的LSUN教堂数据集上进行了无条件生成实验，分辨率为256×256。在表2和表3中，MPQ-DM在所有比特设置下仍然优于所有其他现有方法。
对于LSUN卧室数据集，与基线相比，W3A4的FID降低了5.59，W2A6降低了8.53，W2A4甚至降低了12.81。在W2A4设置下，团队是第一种将sFID降至20以下的方法，与基线相比，sFID降低了6.12。
文本到图像生成
团队在512×512分辨率的Stable Diffusion v1.4模型上，对随机选择的10k COCO2014验证集提示进行了文本到图像生成实验。在表4中，团队的方法比基线和SOTA PTQ方法具有更好的性能。在W3A4和W2A6设置中，团队实现了超过0.3的CLIP评分改善。MPQ-DM+在CLIP评分上进一步提高了1.79，只增加了10.2 MB的模型大小。
消融实验

组件研究
在表5中，团队对LDM-4 ImageNet 256×256模型进行了全面的消融研究，以评估每个拟议组件的有效性。
团队提出的OMQ从量化的角度解决了现有的逐层比特定位方法，在总平均比特宽度不变的情况下，将更多的比特宽度定位到层内具有显著异常现象的信道。
这种层内混合精度量化方法大大提高了基线的性能，增益IS增加了58.88。此外，TSD从模型优化的角度提高了蒸馏过程的鲁棒性，也取得了一定的改进。
通过量化和优化两个方面的并行改进，MPQ-DM实现了最先进的量化性能。


△表5 对提出方法的消融研究
异常值选择方法研究
在表6中，团队研究了LDM-4 ImageNet 256×256模型上混合精度量化中的不同异常显著信道选择方法。团队发现，即使随机选择一些通道到更高或更低的比特，也会带来一定的性能提升。
这表明，在极低比特量化中，增加某些信道的比特带来的增益远远大于减少某些信道带来的影响，这证明了混合量化的必要性。虽然在随机或从权重的头部和尾部选择通道方面有一些收获，但基于峰度的异常值选择方法实现了最显著的性能提升。
这表明峰度选择了性能提高最显著的异常显著信道，这证明了基于峰度的信道选择方法的有效性。
异常值选择函数的消融研究：


△表6 抽取10k个样本进行评估
蒸馏指标研究
在表7中，团队研究了LDM-4 ImageNet模型上方程9中使用的不同蒸馏度量。
团队与未蒸馏的模型进行对比。使用L2损耗来对齐 和 仅显示sFID几乎没有改善，但降低了FID和精度。
这表明离散特征和连续特征不能直接通过数值很好地对齐，从而导致负优化。然而，团队提出的关系蒸馏可以将所有特征转移到统一的相似性空间中。这打破了两个潜在空间之间的差异，提高了模型性能。
蒸馏度量的消融研究：


△表7 抽取10k个样本进行评估
论文链接：https://arxiv.org/pdf/2412.11549