PTQD: Accurate Post-Training Quantization for Diffusion Models阅读

Abstract

迭代去噪过程在推理时计算开销巨大，使得扩散模型在低延迟和可扩展的实际应用中不太实用。扩散模型的训练后量化可以显著减少模型大小并加速采样过程，而无需重新训练。然而，直接将现有的训练后量化方法应用于低比特的扩散模型，会显著损害生成样本的质量。具体来说，对于每一步去噪，量化噪声会导致估计均值的偏差，并与预定的方差计划不匹配。此外，随着采样过程的进行，量化噪声可能会积累，在后期去噪步骤中导致信噪比（SNR）较低。为了解决这些挑战，我们提出了一个统一的公式，用于描述量化噪声和量化去噪过程中的扩散扰动噪声。具体来说，我们首先将量化噪声与其全精度对应部分的相关部分和剩余不相关部分进行解耦。相关部分可以通过估计相关系数轻松校正。对于不相关部分，我们从量化结果中减去偏差以校正均值偏差，并校准去噪方差计划以吸收由量化引起的额外方差。此外，我们引入了一种混合精度方案，用于为每个去噪步骤选择最佳比特宽度，它优先选择较低比特宽度以加速早期去噪步骤，同时确保较高比特宽度在后期步骤中保持高信噪比（SNR）。广泛的实验表明，我们的方法在生成高质量样本方面优于之前的训练后量化扩散模型，与全精度的LDM-4在ImageNet 256×256上的FID得分相比，仅增加了0.06，同时节省了19.9倍的比特操作。代码可在https://github.com/ziplab/PTQD获取。

1 Introduction

扩散模型在多个领域展示了生成高质量样本的显著能力[10, 5, 58, 18, 38, 30, 13, 52, 8, 49]。与生成对抗网络（GANs） [15] 和变分自编码器（VAEs） [27] 相比，扩散模型不会面临模式崩溃和后验崩溃的问题，因此训练更加稳定。然而，扩散模型的应用受限于两个主要瓶颈。首先，扩散模型通常需要数百个去噪步骤来生成高质量样本，这使得过程显著慢于GANs。为了解决这个问题，许多研究[50, 36, 2, 28, 34] 提出了高级的无训练采样器，以减少去噪迭代的次数。其中，最近的一项研究DPM-solver [36] 通过解析计算扩散常微分方程（ODEs）的线性部分，将去噪过程缩短至十步。然而，即使使用这些快速采样器，扩散模型仍然不适合实时应用。例如，即使在高性能平台如RTX 3090上执行，使用DPM-Solver [36] 采样器的Stable Diffusion [45] 生成一个512×512的图像仍需超过一秒。其次，扩散模型在各种设备上的应用受到大量参数和计算复杂度的限制。举例来说，执行Stable Diffusion [45] 需要16GB的运行内存和超过10GB VRAM的GPU，这对大多数消费级PC来说是不可行的，更不用说资源受限的边缘设备了。

模型量化，通过使用较低的数值比特宽度来表示权重和激活，已被广泛研究以减少内存占用和计算复杂度。例如，使用8位模型相比浮点模型在ARM CPU上可以显著加速2.2倍[21]。采用4位量化相比8位量化可以进一步提高高达59%的吞吐量[3]。为了在无需重新训练的情况下促进量化过程，训练后量化（PTQ）已成为一种广泛使用的技术，具有高度的实用性和易于实施性。尽管PTQ在传统模型上已被广泛研究[40, 32, 20, 33, 57]，但其在扩散模型上的应用在基本层面上引发了两个新挑战。首先，随着噪声预测网络的量化，其量化噪声不可避免地在估计均值中引入偏差，并带来与每一步去噪步骤中预定方差计划冲突的额外方差。此外，随着迭代采样过程的进行，量化噪声会积累，导致噪声预测网络在后期去噪步骤中的信噪比（SNR）显著下降。这种信噪比的下降严重阻碍了去噪能力，导致生成图像质量明显下降。

为了应对上述挑战，我们提出了PTQD，一个用于扩散模型的创新训练后量化框架。为了解决每一步去噪中的均值偏差和额外方差，我们通过将量化噪声与其全精度对应部分的相关部分和剩余不相关部分进行解耦，并为其设计了独立的校正方法。通过估计相关系数，可以轻松校正相关部分。对于剩余的不相关部分，我们从估计均值中减去偏差，并提出方差计划校准，将额外方差吸收到扩散扰动噪声中。为了克服后期去噪步骤中信噪比低的问题，我们引入了一个步骤感知的混合精度方案，自适应地为同义步骤分配不同的比特宽度，以在整个去噪过程中保持高信噪比。

总而言之，我们的主要贡献如下：

• 我们提出了PTQD，一个用于扩散模型的创新训练后量化框架，它为量化噪声和扩散扰动噪声提供了统一的公式。
• 我们将量化噪声与其全精度对应部分的相关和不相关部分进行解耦，然后通过估计相关系数校正相关部分，并提出方差计划校准来修正剩余的不相关部分。
• 我们引入了一个步骤感知的混合精度方案，动态选择适当的比特宽度，以在去噪过程中保持信噪比。
• 我们的广泛实验表明，我们的方法在扩散模型的训练后量化方面达到了新的最先进性能。

2 Related Work

3 Preliminaries

3.1 扩散模型

扩散模型[50, 17]在前向过程中逐步对真实数据x0应用高斯噪声，并学习一个反向过程来去噪和生成高质量的图像。对于DDPMs[17]，前向过程是一个马尔可夫链，可以表述为：

[ q(x1:T|x0) = \prod_{t=1}^{T} q(xt|xt−1), ]
[ q(xt|xt−1) = N(xt; \sqrt{\alpha_t} xt−1, \beta_t I), ]

其中，(\alpha_t)和(\beta_t)是超参数，并且(\beta_t = 1 - \alpha_t)。

在反向过程中，由于直接估计q(xt−1|xt)的真实分布是不可行的，扩散模型通过变分推理来近似它，学习一个高斯分布：

[ pθ(xt−1|xt) = N(xt−1; \muθ(xt,t), Σθ(xt,t)) ]

并通过噪声预测网络(\epsilonθ(xt,t))重新参数化其均值：

[ \muθ(xt,t) = \frac{1}{\sqrt{\alpha_t}} xt − \frac{\beta_t}{\sqrt{1 − \bar{\alpha_t}}} \epsilonθ(xt,t) ]

其中，(\bar{\alpha_t} = \prod_{s=1}^{t} \alpha_s)。方差(\Sigmaθ(xt,t))可以重新参数化，也可以固定为一个常数计划(\sigma_t)。当使用一个常数计划时，xt−1的采样可以表述为：

[ xt−1 = \frac{1}{\sqrt{\alpha_t}} xt − \frac{\beta_t}{\sqrt{1 − \bar{\alpha_t}}} \epsilonθ(xt,t) + \sigma_t z, 其中 z \sim N(0,I)。 ]

我们的方法专注于无需训练的扩散模型的训练后量化。相反，我们使用预训练的扩散模型，并继承它们的超参数和方差计划进行推理。尽管本文中展示的推导基于DDPM，但它们可以很容易地扩展到其他快速采样方法，如DDIM[50]。更多信息可以在补充材料中找到。

3.2 模型量化

在我们的研究和所有实验中，我们使用了均匀量化。对于均匀量化，给定一个浮点向量x，目标比特宽度b，量化过程可以定义为：

$$\hat{x}=\Delta \cdot \text{clip}(\lfloor \frac{x}{\Delta }\rceil +Z,0,2^b-1)-Z,$$

其中 (\lfloor·\rceil) 是四舍五入操作，(\Delta = \frac{\max(x)−\min(x)}{2^b−1}) 和 (Z = −\lfloor \frac{\min(x)}{\Delta} \rceil)。

为了确保清晰一致，我们引入符号来定义本文中使用的变量。设X是全精度模型中的一个张量（权重或激活），经过归一化层后的结果表示为X。量化模型的对应张量表示为(\hat{X})。量化噪声表示为(\Delta X)，它是(\hat{X})和X之间的差异。

4 Method

模型量化

模型量化将权重和激活值离散化，这不可避免地会在结果中引入量化噪声。如公式 (3) 所示，在量化扩散模型的反向过程中，xt−1 的采样可以表示为：

$$xt-1=\frac{1}{\sqrt{{\alpha }_t}}xt-\frac{{\beta }_t}{\sqrt{1-\overline{{\alpha }_t}}}{\widehat{ϵ}}_{\theta }(xt,t)+{\sigma }_{t}z=\frac{1}{\sqrt{{\alpha }_t}}xt-\frac{{\beta }_t}{\sqrt{1-\overline{{\alpha }_t}}}({ϵ}_{\theta }(xt,t)+\Delta {ϵ}_{\theta }(xt,t))+{\sigma }_{t}z$$

这里，(\hat{\epsilon}\theta(xt, t)) 是量化噪声预测网络的输出，(\Delta \epsilon\theta(xt, t)) 是量化噪声。额外的量化噪声不可避免地会改变 xt−1 的均值和方差，降低信噪比（SNR），并对生成样本的质量产生不利影响。因此，为了减轻量化的影响，有必要在反向过程的每一步中校正均值和方差，以恢复 SNR。

4.1 相关性解耦

我们首先假设量化噪声与全精度噪声预测网络的结果之间存在相关性。虽然其他因素（如非线性操作）可能会导致这种相关性，但命题 1 证明归一化层是其中的一个原因。

命题 1：设 Y 和 (\hat{Y}) 分别是全精度模型和其量化版本中归一化层的输入，其中量化噪声 (\Delta Y = \hat{Y} - Y) 最初与 Y 无关，则在归一化层之后，量化噪声与全精度模型的输出之间会存在相关性。

证明基于以下事实：(\hat{Y}) 的均值和方差将与 Y 的不同（取决于具体的量化方案）。因此，归一化层之后的量化噪声可以表示为：

$$\Delta \hat{Y}=\frac{Y-{\hat{\mu }}_Y}{{\hat{\sigma }}_Y}-\frac{Y-{\mu }_Y}{{\sigma }_Y}=\frac{{\sigma }_Y\Delta Y-({\hat{\sigma }}_Y-{\sigma }_Y)Y+{\hat{\sigma }}_Y{\mu }_Y-{\sigma }_Y{\hat{\mu }}_Y}{{\hat{\sigma }}_Y{\sigma }_Y}$$

这里，为简化起见，我们省略了归一化层中的仿射变换参数。从公式 (6) 可以看出，分子中的第二项与 Y 相关，而其他三项则无关。因此，归一化层之后，量化噪声 (\Delta Y) 将与 Y 相关。

图 2 中的经验观察结果证实了量化噪声与全精度噪声预测网络输出之间存在很强的相关性，进一步验证了我们的假设。基于该假设和观察结果，量化噪声可以解耦为两部分：

[ \Delta \epsilon_\theta(xt, t) = k \epsilon_\theta(xt, t) + \Delta’ \epsilon_\theta(xt, t) ]

第一部分，用 (k \epsilon_\theta(xt, t)) 表示，与 (\epsilon_\theta(xt, t)) 线性相关。第二部分，用 (\Delta’ \epsilon_\theta(xt, t)) 表示，代表量化噪声的剩余部分，假设与 (\epsilon_\theta(xt, t)) 无关。这里，k 是相关系数，可以通过将量化噪声 (\Delta \epsilon_\theta(xt, t)) 和原始值 (\epsilon_\theta(xt, t)) 应用线性回归来估计。详情见第 5.1 节。

通过公式 (7) 中的解耦，xt−1 的采样可以进一步表示为：

$$xt-1=\frac{1}{\sqrt{{\alpha }_t}}xt-\frac{{\beta }_t}{\sqrt{1-\overline{{\alpha }_t}}}({ϵ}_{\theta }(xt,t)+\Delta {ϵ}_{\theta }(xt,t))+{\sigma }_{t}z=\frac{1}{\sqrt{{\alpha }_t}}xt-\frac{{\beta }_t}{\sqrt{1-\overline{{\alpha }_t}}}(1+k){ϵ}_{\theta }(xt,t)+{\Delta }^{\prime }{ϵ}_{\theta }(xt,t)+{\sigma }_{t}z$$

因此，偏差和额外方差来自量化噪声的相关部分和不相关部分。在接下来的部分中，我们将详细解释如何分别校正这两部分的量化噪声。

4.2 量化噪声校正

4.2.1 相关噪声校正

基于公式 (8)，量化噪声的相关部分可以通过将量化噪声预测网络 (\hat{\epsilon}_\theta(xt, t)) 的输出除以 (1+k) 来矫正，结果为：

$$xt-1=\frac{1}{\sqrt{{\alpha }_t}}xt-\frac{{\beta }_t}{\sqrt{1-\overline{{\alpha }_t}}}{ϵ}_{\theta }(xt,t)+\sigma tz-\frac{{\beta }_t}{\sqrt{{\alpha }_t}}\frac{{\Delta }^{\prime }{ϵ}_{\theta }(xt,t)}{\sqrt{1-\overline{{\alpha }_t}}(1+k)}$$

因此，只剩下不相关的量化噪声。此外，对于 (k \geq 0) 的情况，可以推导出不相关量化噪声的均值和方差被 (1 / (1 + k)) 减小。实践中，我们强制 (k) 为非负，如果它是负值，则将其重置为零。接下来，我们将解释如何处理公式 (9) 中存在的不相关量化噪声。

4.2.2 不相关噪声校正

不相关量化噪声的存在在每一步引入了额外方差，导致总方差超过预定值 (\sigma_t^2)。为了解决这个问题，我们提出了校准量化扩散模型的方差计划，其表示为 (\sigma_t’^2)，比原始计划 (\sigma_t^2) 更小。为了估计 (\sigma_t’^2)，我们进一步假设并将不相关量化噪声建模为均值为 (\mu_q)、方差为 (\sigma_q^2) 的高斯分布：

$${\Delta }^{\prime }{ϵ}_{\theta }(xt,t)\sim N({\mu }_q,{\sigma }_q^2)$$

为了验证这一假设，我们进行了统计测试（详见补充材料）并在图 3 中展示了不相关量化噪声的分布。均值和方差的值可以通过使用量化和全精度扩散模型生成样本，并收集不相关量化噪声的统计数据来估计。按照先前的工作 [41]，均值偏差可以通过偏差校正（BC）来纠正，我们收集不相关量化噪声的通道均值，并从量化噪声预测网络的输出中减去它们。对于不相关量化噪声的方差，我们提出方差计划校准（VSC），假设不相关量化噪声可以被吸收到高斯扩散噪声中。通过将校准方差计划 (\sigma_t’^2) 代入公式 (9) 而保持每一步的方差不变，我们可以使用以下方法求解最优方差计划：

$${\sigma }_t^{\prime 2}=\{\begin{array}{ll}{\sigma }_t^2-\frac{{\beta }_t^2}{{\alpha }_t(1-\overline{{\alpha }_t})(1+k{)}^2}{\sigma }_q^2,& \text{if }{\sigma }_t^2\ge \frac{{\beta }_t^2}{{\alpha }_t(1-\overline{{\alpha }_t})(1+k{)}^2}{\sigma }_q^2\\ 0,& \text{otherwise}\end{array}$$

可以看出，如果量化噪声的额外方差小于噪声超参数 (\sigma_t^2)，则可以消除量化引起的方差增加。根据公式 (12)，量化噪声方差的系数可以计算为：

$$\frac{{\beta }_t^2}{{\alpha }_t(1-\overline{{\alpha }_t})(1+k{)}^2}$$

通常该值足够小，确保量化噪声可以完全被吸收，除了确定性采样的情况，其中 (\sigma_t) 为零。在这种情况下，没有解析解 (\sigma_t’^2)，我们使用最优解 (\sigma_t’^2 = 0)。总体而言，过程量化噪声校正总结在算法 1 中。

尽管所提出的方法可以校正每一步的均值偏差和方差值，但由于量化噪声预测网络的低信噪比（SNR），使用低比特扩散模型生成令人满意的样本仍然具有挑战性。在下一节中，我们将详细分析这个问题。

4.3 步骤感知混合精度

给定全精度噪声预测网络 (\epsilon_\theta(xt, t)) 和相应的量化噪声 (\Delta \epsilon_\theta(xt, t)) 的输出，我们定义量化噪声预测网络的 SNRQ 为：

$$\text{SNRQ}(t)=\frac{\Vert {ϵ}_{\theta }(xt,t){\Vert }^2}{\Vert \Delta {ϵ}_{\theta }(xt,t){\Vert }^2}$$

图 4 展示了不同比特宽度和校正方法下的 SNRQ。图中揭示了几个见解：1）随着步骤 t 的减少，SNRQ 急剧下降；2）更高比特宽度的模型表现出更大的 SNRQ；3）所提出的校正方法显著提高了 SNRQ，尤其是对于大步数。第一个观察结果突出表明，使用低比特扩散模型生成高质量样本的挑战。特别是，当 t 接近零时，W4A4 扩散模型的 SNRQ 降低并接近 1，意味着量化噪声的幅度与噪声预测网络的原始结果相当。为了在保持良好生成性能的同时启用低比特扩散模型，我们提出了一种新方法，称为步骤感知混合精度，通过为同义步骤设置不同的比特宽度来保持所有步骤中的 SNRQ 在合理范围内。

具体来说，权重的比特宽度是固定的并在不同去噪步骤之间共享，这消除了在采样过程中存储和重新加载多个模型状态文件的需要。因此，我们仅调整激活值的比特宽度。正式地，我们为激活值预定义了一组比特宽度 B = {b1, b2, …, bn} 并在每个比特宽度下评估 SNRQ。为了建立 SNRQ 的基准，我们遵循先前的研究 [36, 26] 并引入基于前向过程的 SNRF，表示每一步的数据噪声程度：

$$\text{SNRF}(t)=\frac{{\alpha }_t^2}{{\sigma }_t^2}$$

图 4 展示了 SNRF(t)，其随步骤 t 的增加严格递减。为了确定每一步 t 的最优比特宽度，我们将每个比特宽度的 SNRQ 与 SNRF 进行比较，并选择满足条件的最小比特宽度 bmin：

$${\text{SNRQ}}_{bmin}(t)>\text{SNRF}(t)$$

如果没有比特宽度满足此条件，我们使用 B 中的最大比特宽度以实现更高的 SNR。实际上，不同比特宽度的模型是分别校准的，校准集从相应的步骤中收集。

5 Experiments

6 Conclusion and FutureWork

在本文中，我们提出了PTQD，一个用于扩散模型的创新训练后量化框架，它统一了量化噪声和扩散扰动噪声的表述。首先，我们将量化噪声相对于其全精度对应部分解耦为相关部分和剩余不相关部分。为了减少每一步中的均值偏差和额外方差，可以通过估计相关系数轻松校正相关部分。对于不相关部分，我们提出了方差计划校准来吸收其额外方差，并通过偏差校正来修正均值偏差。此外，我们引入了步骤感知的混合精度方案，自适应地为每个去噪步骤选择最佳比特宽度。通过结合这些技术，我们的PTQD在现有最先进的训练后量化扩散模型上取得了显著的性能提升，与全精度的LDM-4在ImageNet 256×256上的FID增加仅为0.06，同时节省了19.9倍的比特操作。在未来，我们可以进一步量化扩散模型中的其他组件，如文本编码器和图像解码器，以实现更高的压缩比和加速性能。我们还可能将PTQD扩展到更广泛的生成任务中，以评估其有效性和通用性。

限制和广泛影响。所提出的PTQD框架以其高效率和节能特性脱颖而出，这在减少因广泛部署扩散模型而产生的碳排放方面具有重要意义。然而，和其他深度生成模型类似，PTQD也有潜在的风险可能被用于生成伪造图像和视频，从而被恶意利用。