QuEST: Low-bit Diffusion Model Quantization via Efficient Selective Finetuning

公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）

目录

0. 摘要

2. 相关工作

3. 方法

3.1. 预备知识

3.2. 扩散模型的量化感知属性

3.3. 通过高效选择微调进行量化

3.3.1. 微调的充分性

3.3.2. 选择性微调进行量化

4. 实验

---

0. 摘要

扩散模型在图像生成任务中取得了显著成功，然而其实际应用受到高内存和时间消耗的限制。虽然量化为扩散模型的压缩和加速铺平了道路，但在模型被量化为低比特时，现有方法完全失败。在本文中，我们揭示了量化扩散模型中三个影响当前方法有效性的属性：不平衡的激活分布、不精确的时间信息和对特定模块扰动的脆弱性。为了缓解由分布不平衡引起的强化低比特量化困难，我们提出了对量化模型进行微调，以更好地适应激活分布。基于这个思路，我们确定了两种关键类型的量化层：持有重要时间信息的层和对减少位宽敏感的层，并对它们进行微调以提高效率。我们在实证中验证了我们的方法修改了激活分布并提供了有意义的时间信息，从而促进了更轻松、更准确的量化。我们的方法在三个高分辨率图像生成任务上进行了评估，在各种位宽设置下取得了最先进的性能，同时也是第一个在全 4 比特（即 W4A4，权重 W，激活 A）Stable Diffusion 上生成可读图像的方法。

神经网络量化为加速推断速度和同时减少内存消耗提供了可行的解决方案（Gholami等，2022年）。它旨在将高比特模型参数压缩为低比特近似，几乎没有性能降级。例如，使用 4 比特权重和4 比特激活量化理论上可以实现高达 8 倍的推断时间加速和 8 倍的内存减少（Liang等，2021年）。当位宽进一步降低时，这些优势将持续增加。因此，扩散模型的低比特量化成为提高其效率的可行方法。

2. 相关工作

模型量化是优化深度学习模型推断内存和速度的主导技术，通过降低计算中使用的张量的精度。对于扩散模型的量化，进行了一些研究，主要分为两个主流框架：量化感知训练（Quantization-Aware Training，QAT）（Jacob等，2017年；Li等，2022年；Xu等，2023年；Li等，2023b年）和训练后量化（Post-Training Quantization，PTQ）（Wang等，2020年；Nahshan等，2021年；Li等，2021年；Wei等，2023a年；Liu等，2023年）。

量化感知训练，如 Q-DM（Li等，2023b年），将数值精度视为附加的优化约束，从头开始训练所有参数以实现任务目标。虽然对于低比特量化是有效的，但比训练全精度模型更消耗资源。Efficient-DM（He等，2023a年）进一步使用 LoRA（Hu等，2021年）降低训练成本。然而，它引入了额外的权重参数，并仍需要大量的训练迭代。

训练后量化方法旨在根据小型校准集（calibration set）校准量化参数。开创性的作品，如PTQ4DM（Shang等，2023年）和 Q-Diffusion（Li等，2023a年），专注于在不同时间步采样全精度模型的输出，并以一定的概率构建校准集。ADP-DM（Wang等，2023年）将时间步骤分成不同组，并使用可微搜索算法优化每个组。PTQ-D（He等，2023b年）分解量化噪声，并根据全精度模型的统计数据单独校正它们。然而，这些方法通常在 4 比特或更低位宽下失败。为了实现与高效率的低比特兼容性，我们的方法采用了 PTQ 的一般框架，但进一步引入了高效的权重微调策略。表 1 显示了我们方法的独特能力，我们能够实现 QAT 能力与 PTQ 效率的结合。

3. 方法

3.1. 预备知识

对于向量中的单个值 x 的量化过程可以表示为：

其中 ^x 是量化的整数结果，round(·) 表示舍入算法，如 round-to-nearest 运算符（Li等，2021年）和 AdaRound（Nagel等，2020年），s 被称为缩放因子，Z 是零点（zero-point）。clamp 是将值钳制在 [q_min, q_max] 范围内的函数，由位宽确定。反之，将量化值转换回完整精度形式的过程表示为：

这被称为去量化过程。量化和去量化过程都在模型权重和层输出（也称为 '激活'）上执行。方程式2 表明量化误差由两个因素组成：由范围钳制引起的剪切误差和由舍入函数引起的舍入误差，它们呈现一种权衡关系（Li等，2021年）。虽然先前的方法努力在这两个误差之间找到最佳平衡，但它们忽视了量化扩散模型的固有机制。在下一节中，我们研究了与量化性能相关的扩散模型的三个属性，并展示了为什么典型的量化方法往往会失败。 

模型量化：以较低的推理精度损失将连续取值的浮点型模型权重进行裁剪和取舍，以更少比特的数据类型近似表示 32 位有限范围浮点型数据的过程，而模型的输入输出依然是浮点型，从而达到减少模型尺寸大小、减少模型内存消耗及加快模型推理速度等目标。

比如 int8 量化，就是让原来 32bit 存储的数字映射到 8bit 存储。int8 范围是 [-128,127]， uint8 范围是 [0,255]。

参考：https://errolyan.medium.com

3.2. 扩散模型的量化感知属性

属性 1：激活输出很可能具有大多数接近零的值，而其他值则具有数值较大且重要。先前的研究（Shang等，2023年；Li等，2023a年）关注了不同时间步上激活分布的变化，但没有一个关注激活分布本身。如图 2 所示，我们发现在一些层中，虽然大多数值接近 0，但存在相对较大且多样的值。以左侧的条形图为例，原始激活值（蓝线）范围为 [-10, 34]，但超过 99% 的值在 [-0.6, 1.7] 之间，这使得最小化剪切误差变得困难。这种现象不同于 LLMs（Lin等，2023年；Xiao等，2023年）中观察到的异常值分布（outlier distribution），其中异常值始终为正，并在特定位置一致出现（Wei等，2023b年）。此外，我们观察到较大的值对生成性能很重要。由于较大的值很重要，而较小的值经常出现，两者都不可忽视，需要同时进行仔细的量化。

然而，典型的量化方法在低比特设置下往往无法实现这种能力，其中在优化过程中舍入误差往往超过剪切误差，导致过度剪切的大值，生成损坏的图像。这种现象启发我们通过模型微调来重新塑造激活分布，以获得更适合量化的分布。

属性 2：准确的时间信息对于量化扩散模型的性能至关重要。扩散模型的预测与时间步信息高度相关。具体而言，时间步被转换并投影为时间嵌入，然后添加到图像特征中，引导模型正确去噪。表 2 显示了在 LSUN-Bedrooms 的 LDM-4 中量化时间嵌入或保留其全精度值的性能比较（Yu等，2015年）。在 W8A8 和 W4A8 下，量化时间嵌入导致 FID 分别增加 0.81 和 1.04（相对增加15%）。这表明不准确的时间嵌入量化可能导致生成性能的显著降级，在较低位宽下甚至可能导致更大的图像破坏。由于扩散模型在不同时间步对图像生成的关注不同，我们推断不准确的时间嵌入可能导致输入图像和模型功能不匹配，可能导致噪声去除序列中的振荡。因此，在量化下保持扩散模型性能需要准确的时间信息。

属性 3：不同的激活对减少位宽有不同的敏感性。低比特激活量化通常比权重量化更困难，仍然是扩散模型量化的主要挑战。与可以分类为线性层和卷积层的权重不同，激活类型更为多样且复杂。因此，我们研究了在量化不同激活时对模型性能的影响。图 3 显示了在 Stable Diffusion 中不同激活对量化的敏感性。我们将三种不同类型的激活量化到较低比特，同时保持其他激活的位宽（8 比特）。我们观察到 FeedForward 层（Feng等，2023年）激活在 6 比特时导致生成失败，而其他线性层的激活（包含 5 倍更多层）在 4 比特时几乎不失败，所有卷积层的激活（包含 3 倍更多层）仅在 4 比特时失败。这种现象表明一小部分激活对低比特量化特别敏感，因此需要特别处理。

3.3. 通过高效选择微调进行量化

在本节中，我们首先在理论上讨论先前方法失败的潜在原因以及权重微调的基本原理。然后，我们介绍了一种名为 QuEST （(Quantization via Efficient Selective Fine-Tuning).）的扩散模型微调方法，它可以显著提升低比特性能，并具有较少的时间和内存使用。图 1 展示了提出的框架，这是一个基于蒸馏的微调策略，包括选择性权重优化和网络范围的缩放因子优化。为了进一步解决由于量化而导致的时间信息丢失，我们提出了一种高效的时域感知激活量化器。

3.3.1. 微调的充分性

首先，我们回顾先前方法的潜在理论基础，这些方法通常采用基于重构的训练后量化方法。将在时间 t 的扩散模型的激活表示为 z_t = [z_(1,t), z_(2,t), ..., z_(n,t)]，任务损失为L(z_(n,t), -z_(n,t))，其中 n 是层数，z_(n,t) 是地面实况。L 可以是任何损失函数，这里我们使用均方误差（MSE）。假设模型权重 w 被冻结，我们将量化视为一种扰动，并使用泰勒展开来公式化激活量化的影响： 

其中， Δ 是激活扰动，-g^(z) 是梯度，而 -H^(zn,t) 是 Hessian 矩阵。根据 (Li 等人，2021; Yuan 等人，2022)，对于一个训练良好的模型

趋近于 0。因此上述方程可以简化为：

这导致了一个非常直接的结论：通过重构模型输出可以减少量化误差。然而，在低比特设置下，从方程 4 到方程 5 的推理是不准确的，其中激活扰动 Δ 对于有意义的泰勒展开来说太大。因此，我们有以下命题：

命题 3.1. 基于重构的后训练量化方法可能会失去它们的理论保证，因为在低比特量化下，存在较大的值扰动。

为了解决上述问题，我们可以使用以下定理将 Δ 转化为一个较小的扰动 ϵ：

定理 3.2. 给定时间 t 的 n 层扩散模型，其量化激活为 ˜zt = [˜z_(1,t), ˜z_(2,t), ..., ˜z_(n,t)]，且 ˜z_(n,t) = z_(n,t) + Δ，其中 z_(n,t) 是地面真值，Δ 是由低比特量化引起的大扰动。将目标任务 MSE 损失表示为

则量化误差可以转化为： 

其中，w_n 是第 n 层的权重，ϵ 是一个足够小以进行泰勒展开的扰动，K 是一个常数，Δ = Kϵ。

定理 3.2 表明：为了最小化量化误差，理想情况下我们可以微调 w_n，使得对于任意 i，能够拟合数据/激活

这是全精度对应物不能拟合的，从而使模型对不同输入更加鲁棒。换句话说，我们可以微调模型权重以提高对大输入激活扰动的鲁棒性，从而更容易进行量化。我们还发现方程 6 中的第二项可以与第一项同时最小化，从而通过一个共同的目标减小总量化误差：通过微调 w_n 实现激活和输出对齐。

（注：对于公式 6，通过微调 w_n，使量化激活拟合全精度激活，从而最小化量化误差。）

3.3.2. 选择性微调进行量化

受量化扩散模型的揭示属性以及现有量化方法的不足启发，我们提出了一种有效的低比特扩散模型量化微调策略。我们方法的组成部分如下所述。 

网络逐部无数据训练（Data-Free Partial-network-wise Training）。为了减轻对大量训练数据的需求，我们建议采用训练后量化的一般设置，并以无数据的方式构建校准集。通过将采样的高斯噪声 xT 输入到全精度模型，并在不同的时间步骤上进行迭代采样，我们可以获得微调量化模型所需的校准数据。在实践中，每个时间步生成的样本数量为 256，这意味着我们只需对全精度模型进行少数次推断，即可获得所需数量的校准样本。

此外，受到我们可以将激活量化参数视为额外的模型参数，这些参数都对最终输出有贡献的事实的启发，我们提出了一种网络逐部训练策略。与使用将量化参数与其相应的层或块绑定的逐层或逐块重构的量化方法（Shang 等人，2023; Li 等人，2023a）不同，我们为每个时间步优化所有激活缩放因子，仅训练部分权重参数（例如，在 LDM-4 中的参数的 7%）。如图 1 所示，大多数模型权重在时间步 t 保持不变，只有与该时间步相关的一小部分参数（包括层权重和缩放因子）使用全精度模型的输出作为地面真值进行更新。对于要微调的权重的选择在以下部分讨论。此外，逐层/块的优化方法只能顺序重构，而我们的方法可以同时更新所需的参数。通过这种方式，我们显著节省了进行量化所需的时间和内存。请注意，在初始化后，权重的量化参数不会进行更新，以保持在部署期间的内存效率。

高效的时间感知激活量化（Time-aware Activation Quantization，TAQuant）。时间信息是扩散模型中的一个独特而重要的因素。正如在（Shang等人，2023; Li等人，2023a）中强调的那样，激活分布在不同时间步骤上的变化对于传统的 PTQ 方法是致命的。此外，正如在第 3.2 节中讨论的，准确的时间嵌入量化是关键的。为了区分不同的时间步骤，我们设计了时间感知激活量化器（TAQuant），它采用不同的激活量化参数集合用于不同的时间步骤。虽然这些参数消耗的内存微乎其微，不影响推断速度，但为每个时间步学习独立的参数会导致巨大的时间成本。受到相邻时间步产生扩散模型中类似功能的事实的启发（Choi等人，2022），我们使用简单的策略对相邻的时间步进行聚类，使它们共享相同的校准数据和量化参数。然后，激活量化参数集合被制定为：

其中用于优化的时间步骤是均匀采样的，每个选择的时间步骤的统计数据用于代表相邻的 T/t 个时间步骤。尽管大多数时间步骤的量化参数是近似的，这种策略节省了 t 倍优化时间，而几乎不影响性能。第 4.3 节进一步分析了时间效率与模型性能之间的权衡。 

选择性 & 渐进式层对齐（Selective & Progressive Layer Alignment）。正如在第 3.2 节中所研究的，我们确定了两种对图像生成性能至关重要的激活类型，并选择它们进行权重微调：时间嵌入和与注意力相关的激活（例如 Stable Diffusion 中的 FeedForward 层）。由于它们具有不同且不重叠的功能，我们逐步更新这些组件。

在单个前向过程中，相同的时间嵌入被注入到模型的不同部分，通过投影层传递，并与潜在图像表示合并。这意味着时间信息独立于主要的网络流动。因此，我们首先通过以下方式微调时间嵌入层 l 的权重 w_l 以及其激活量化参数 s^t_l：

其中，C_TE 表示时间嵌入层的集合，e(t) 是时间 t 的固定时间嵌入，-O(e(t);wl) 是全精度模型的激活输出，代表地面真值，而 ˜O(e(t);wl, stl) 则是量化输出。 这个损失函数表明，选择的权重参数在不同的时间步骤中被一致地更新，以确保对不同的输入具有鲁棒性，而仅调整与每个特定时间步骤相关的缩放因子子集 s^t_l，以进行时间信息的区分。 

除了时间嵌入层之外，注意力相关层（FeedForward 层）也非常重要。正如图 3 所示，减少这些层的激活位宽最大程度地降低了模型性能。将 C_A 表示为包含所有这些层的集合，并给定图像校准样本 xt，我们优化 wl 和 st，但不包括在方程 8 中已经优化的量化参数：

其中 w_(:l)​ 是前 l 层的权重参数。

这个策略类似于线性探测，只微调与激活输出直接相关的层。此外，我们的目标不仅是减小每个独立层的量化误差，还要提高最终生成的图像的质量。因此，引入了目标任务损失 ∑t​ L^t_d​，其中地面真值由全精度对应物的输出近似。通过整合方程 8 和方程 9，最终目标被制定为：

4. 实验