基于扩散的图像生成模型近年来通过展示合成高质量内容的能力取得了巨大成功。然而，这些模型包含了大量的参数，导致模型体积显著增大。在多种应用中，尤其是那些在资源受限设备上运行的应用中，保存和传输它们成为一个主要瓶颈。
在这项工作中，作者开发了一种新颖的权重量化方法，将Stable Diffusion v1.5中的UNet量化到位，实现了比原始模型小的模型体积，同时展示出甚至比原始模型更优质的生成质量。
作者的方法包括了几项创新技术，例如为每个层分配最佳位数，初始化量化模型以获得更好的性能，以及改进训练策略以显著减少量化误差。此外，作者通过广泛的基准数据集评估和人工评估来证明作者量化模型的卓越生成质量。
1 Introduction
近期在开发基于扩散的图像生成模型方面的努力已经展示了在合成高保真和逼真图像方面的显著成果，这导致了各种应用，如内容创建和编辑，视频生成，以及3D资产合成等。
然而，扩散模型（DMs）的一个缺点是参数量巨大，例如百万甚至十亿 Level ，由于模型体积庞大，给转移和存储带来了重大负担，特别是在移动和可穿戴设备等资源受限的硬件上。
现有研究通过设计高效架构和网络剪枝探索减少大规模文本到图像扩散模型的大小。这些方法通常由于对预训练网络的修改而需要大量的训练。另一种有前景的减少模型存储量的方向是量化[12; 30]，其中浮点权重被转换为低比特固定点表示，从而节省计算内存和存储空间。
在通过量化压缩扩散模型（DMs）方面已有越来越多的努力。
然而，这些方法仍面临几个主要挑战，尤其是在对类似Stable Diffusion v1.5（SD-v1.5）[62]这样的大规模文本到图像扩散模型进行量化时。首先，这些方法大多是在相对小规模的DMs上开发的，这些DMs是在受限数据集上训练的。例如，在CIFAR-10上训练的模型大约需要 MB的存储空间[21; 38]。
相比之下，SD-v1.5在完全精度格式下需要 GB的存储空间。将这些方法适应到SD-v1.5仍是一个具有挑战性的问题。其次，当前的研究主要集中在将权重量化到位。如何将模型量化到极低比特尚未得到充分研究。第三，缺乏对量化方法在大规模DMs（即SD-v1.5）上表现的公平和广泛的评估。
为了解决上述挑战，本文提出了BitsFusion，这是一个量化感知训练框架，它采用一系列新颖技术将预训练的大规模DMs的权重压缩到极低比特（即1.99比特），实现了更好的性能（即更高的图像质量和更好的文本图像对齐）。因此，作者将SD-v1.5中的 GB UNet（FP16）1压缩到 MB的模型，实现了的压缩比。具体来说，作者的贡献可以总结为以下四个方面：针对DMs的混合精度量化。
作者提出了一种有效的混合精度量化DMs的方法。首先，作者彻底分析了适当的指标以理解在量化DMs中的量化误差（第3.2节）。其次，基于分析，作者根据各层的量化误差将不同层量化到不同比特（第3.3节）。
- 量化DMs的初始化。 作者介绍了几种技术来初始化量化模型以提高性能，包括时间嵌入的预计算和缓存、添加平衡整数以及用于缩放因子初始化的交替优化（第4.1节）。
- 改进量化DMs的训练流程。 作者通过提出的两阶段训练方法（第4.2节）改进了量化模型的训练流程。在第一阶段，作者使用全精度模型作为教师，通过蒸馏训练量化模型。作者的蒸馏损失强制量化模型从教师网络中学习预测噪声和中间特征。此外，作者在训练过程中调整时间步采样的分布，使得引起较大量化误差的时间步被更频繁地采样。在第二阶段，作者使用普通的噪声预测[21]对模型进行微调。
- 广泛的定量评估。 作者首次在文献中进行广泛的定量分析，比较量化模型与原始SD-v1.5的性能。作者包括在各种基准数据集上的结果，即TIFA[25]，GenEval[13]，CLIP分数[58]和FID[19]在MS-COCO 2014验证集[43]上的表现。此外，作者还对PartiPrompts[77]进行了人工评估。
- 作者的1.99位权重量化模型在所有评估中一致地超过了全精度模型，证明了作者方法的有效性。
2 Related Works
为了提高模型在存储和计算成本方面的效率，作者采用了量化方法对扩散模型进行优化，主要分为两种类型：
后训练量化（PTQ）和量化感知训练（QAT）。PTQ不需要完整的训练循环；
相反，它使用有限的校准数据集来调整量化参数。例如，PTQ4DM [65]校准量化参数以最小化DMs的量化误差。Q-Diffusion [37]通过块状重建[41]最小化量化误差。PTQD [18]将量化噪声与DMs采样步骤中固有的随机噪声结合在一起。TDQ [68]优化了不同时间步激活的缩放因子，适用于PTQ和QAT策略。TFMQ [27]专注于重构时间嵌入和投影层以防止过拟合。然而，与QAT相比，PTQ通常会导致性能下降，特别是在追求极低比特DMs时。相比之下，QAT涉及训练完整的权重以最小化量化误差，因此与PTQ相比，性能更高。例如，EfficientDM [17]，受LoRA [24]启发，引入了一个量化感知的低秩 Adapter 来更新LoRA权重，避免训练整个权重。Q-DM [38]通过提出的Q-注意力块对注意力特征进行归一化和平滑操作，以提高量化性能。然而，现有研究主要针对在受限数据集上训练的小规模DMs进行4比特及以上的量化。在本文中，作者专注于将大规模的Stable Diffusion量化到极低比特，并在不同的基准数据集上广泛评估性能。
3 Mixed Precision Quantization for Diffusion Models
在本节中，作者首先回顾权重量化和小波生成扩散模型的公式。然后确定混合精度策略，为不同层分配优化的位宽以减少整体量化误差。具体来说，作者首先分析扩散模型中每一层的量化误差并得出敏感属性，然后基于此分析，作者通过联合考虑参数效率（即，尺寸节省）为每一层分配适当的位数。
Preliminaries
Per-Layer Quantization Error Analysis
获取量化模型。 首先作者对扩散模型进行逐层敏感性分析。具体来说，给定一个预训练的全精度扩散模型，作者将每一层分别量化为1、2、3位，同时将其他层保持在全精度，并分别进行量化感知训练（QAT）。例如，对于具有256层（不包括时间嵌入、第一层和最后一层）的SD-v1.5 UNet，作者总共得到768个量化候选。作者在预定义的训练子数据集上对每个候选进行QAT，并通过与全精度模型比较来验证每个候选产生的量化误差（更多细节见附录B）。
测量量化误差。 为了找到适当解释量化误差的方法，作者分析了四种指标：均方误差（MSE），它量化了图像之间（在作者的案例中是浮点模型和量化模型生成的图像）的像素级差异，LPIPS[80]评估类似人类的感知相似性判断，PSNR[23]通过比较信号可能的最大功率与受损噪声的功率来衡量图像质量，以及CLIP分数[58]评估图像与其语言描述之间的相关性。在收集了分数（图1(b)和图1(c)中的示例，完整的指标列在附录F中）之后，作者进一步通过计算相同位宽下的不同指标（见表1）以及相同指标下的不同位宽（见表2）的一致性来衡量它们之间的皮尔逊相关性[8]。通过这些实证结果，作者得出以下两个主要观察。
观察1：MSE、PSNR和LPIPS显示出强烈的相关性，并且它们与图像质量的视觉感知密切相关。
表1显示，在相同的位宽下，MSE与PSNR和LPIPS高度相关。此外，作者观察到不同位宽下的逐层量化误差趋势相似，如表2所示。至于图1(a)和1(b)中的视觉质量，作者可以看到，更高的MSE误差导致图像质量严重下降，例如，突出显示的RB conv捷径。因此，MSE指标有效地反映了量化引起的质量退化，进一步包含PSNR和LPIPS是不必要的。
观察2：在低位量化后，CLIP分数的变化与MSE在不同层上的变化并不始终一致相关。尽管某些层的MSE较小，但它们可能会经历更大的语义退化，这在CLIP分数的变化中体现出来。
图2：对SD-v1.5的UNet的所有层进行1位量化误差分析。
作者注意到，在量化之后，所有层的CLIP分数变化与MSE只有很弱的关联，如表1所示。一些层显示出较小的MSE但较大的CLIP分数变化。例如，在图1(b)中，CA tok层（从左到右第五个突出层（绿色））的MSE小于RB conv层（从左到右第六个突出层（橙色）），但CLIP分数的变化却相反。正如图1(a)第一行所示，与RB conv层相比，量化这个CA tok层将图像内容从“一个泰迪熊”变为“一个人”，这与文本提示A teddy bear on a skateboard in Times Square, doing tricks on a cardboard box ramp相背离。这是因为MSE只测量了两幅图像之间的差异，这并没有捕捉到语义退化。相比之下，CLIP分数反映了文本和图像之间在语义信息方面的量化误差。因此，作者采用CLIP分数作为表示量化误差的补充指标。
Deciding the Optimal Precision
基于以上的观察，作者接着开发位宽分配的策略。作者选择MSE和CLIP作为作者的定量指标，同时使用每层参数的数量作为尺寸节省的指示器。
基于MSE分配位宽。 直观上，参数更多且量化误差较低的层更适合进行极低比特量化，因为模型的总体位宽可以显著减少。据此，作者提出了一个基于层大小感知的敏感度分数。对于第层，其位()的敏感度分数定义为，其中表示均方误差，是层的总参数数量，表示参数大小因子。为了确定每层的位宽（即），作者定义了一个敏感度阈值，第层被分配到位，其中。其余层为4位。
基于CLIP分数分配位宽。 对于量化后CLIP分数下降的层，作者不是像上面讨论的那样基于敏感度分数分配位宽，而是直接为这些层分配更高的位宽。因此，量化模型可以生成与提示的语义信息对齐的内容。作者在附录B的Alg. 1中提供了详细的混合精度算法。
4 Training Extreme Low-bit Diffusion Model
在确定了每一层的比特数之后，作者接着使用一系列技术来训练量化模型以提升性能。作者的方法的概述在图3中展示。
Initializing the Low-bit Diffusion Model
时间嵌入的预计算和缓存。 在扩散模型的推理过程中，时间步 通过投影层转换成嵌入以纳入扩散模型中。正如现有工作[27]所提到的，投影层的量化可能导致较大的量化误差。然而，每个时间步 的嵌入始终是相同的，这表明作者实际上可以在离线状态下预计算嵌入，并在推理过程中加载缓存值，而不是每次都计算嵌入。此外，时间嵌入的存储大小仅为投影层的 小。因此，作者预计算时间嵌入，并保存不包含投影层的模型。更多细节在附录C中提供。
添加平衡整数。 通常，深度神经网络中的权重分布被观察到是围绕零对称的[84]。为了在SD-v1.5上验证这一假设，作者分析了在完全精度下各层的权重分布，通过计算权重的偏度。值得注意的是，超过的层的偏度在范围内，表明几乎所有层的权重分布都是对称的。更多细节在附录D中提供。然而，现有的关于扩散模型量化的研究忽略了这一对称性[37; 65; 38]，因为它们执行相对较高的比特量化，例如比特或比特。这在极低比特水平下会损害模型性能。例如，在比特量化中，可能的最对称整数结果只能是或。同样，对于比特量化，最平衡的映射整数可以是或，这显著破坏了对称性。在低比特量化中，2个或4个数值中缺少一个单一值可能会产生重大影响。为了解决这一问题，作者利用位平衡策略[36; 51]来初始化模型。具体来说，作者引入一个额外的值来平衡原始量化值。也就是说，在1比特模型中，作者将候选整数集从调整为，以实现更平衡的分布。这样做，作者将平衡的比特权重视为比特。
通过交替优化初始化缩放因子。 初始化缩放因子是量化过程中的重要步骤。现有的QAT研究通常采用Min-Max初始化策略[17; 48]，以确保充分代表并保留异常值。然而，在比特这类极低比特量化设置中，这种方法面临挑战，因为它忽视了全精度权重的分布，导致较大的量化误差，并增加了收敛难度。因此，作者旨在最小化量化权重与全精度权重之间的误差，优化目标如下：
尽管如此，考虑到舍入操作，计算精确的闭式解并不是直接的[29]。受到Lloyd-Max算法[28; 49]的启发，作者使用一种优化方法来最小化缩放因子，以减少作者量化扩散模型的初始化误差，如下所示：
Two-Stage Training Pipeline
在初始化了混合精度模型后，作者引入了_两阶段_训练流程。在第一阶段，作者通过蒸馏损失用全精度模型作为教师训练量化模型。在第二阶段，作者用前一个阶段的模型通过噪声预测进行微调[21; 72]。
然后，作者通过添加平衡整数、预计算和缓存时间嵌入，以及交替优化缩放因子（第4.1节）来初始化量化UNet。中图：在第一阶段训练中，作者冻结教师模型（即SD-v1.5）并通过考虑量化误差的CFG感知量化蒸馏和特征蒸馏损失优化量化UNet。在第二阶段训练中，作者用噪声预测微调前一个模型。右图：对于推理阶段，使用预缓存的时间特征，作者的模型处理文本提示并生成高质量图像。
CFG感知量化蒸馏。与现有工作[11]类似，作者微调量化扩散模型以提高性能。这里权重和缩放因子都会被优化。此外，作者注意到，与直接用普通噪声预测训练相比，用全精度模型以蒸馏方式训练量化模型能获得更好的性能。此外，在蒸馏过程中，量化模型必须意识到CFG，即文本在蒸馏过程中被应用。具体来说，作者的训练目标如下：
其中控制训练过程中的文本丢弃概率，设置为。
使用噪声预测的微调。 在第一阶段用蒸馏损失训练模型之后，作者在第二阶段根据方程式（2）对其进行噪声预测的微调。作者以的概率应用文本丢弃，并根据上面介绍的量化误差修改时间步采样的分布。作者之所以采用两阶段微调，而不是将第一阶段和第二阶段结合起来，是因为作者观察到这样能获得更稳定的训练结果。
5 Experiments
实现细节。 作者使用diffusers库2开发代码，并使用AdamW优化器[33]和恒定学习率在作者内部数据集上训练模型。
对于第一阶段，作者使用8个NVIDIA A100 GPU，总批处理大小为来训练量化模型，迭代K次。对于第二阶段，作者使用个NVIDIA A100 GPU，总批处理大小为来训练量化模型，迭代K次。在推理过程中，作者采用PNDM调度器[45]和个采样步骤来生成用于比较的图像。
其他采样方法（例如，DDIM [70]和DPMSolver [50]）也得出相同的结论（附录K）。
评估指标。 作者在MS-COCO [44]，TIFA [26]，GenEval [14]上进行了CLIP分数和FID评估，并在PartiPrompts [77]上进行了人工评估。作者在CLIP分数中采用了ViT-B/32模型[10]，在GenEval中采用了Mask2Former(Swin-S-82) [5]。
参考
[1].BitsFusion: 1.99 bits Weight Quantization of Diffusion Model.
作者：小先知
来源：未来先知
推荐阅读
- GPTQ & SmoothQuant & AWQ 代码解析
- LLM PTQ量化经典研究解析
- PipeFusion：如何用PCIe互联GPU 低成本并行推理扩散模型
- I-LLM：首次实现了LLM全整形量化，精度超过Smooth/Omini/AffineQuant
欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式客栈专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
