[学习笔记]Post-training Quantization on Diffusion Models

Post-training Quantization on Diffusion Models

扩散模型的训练后量化(静态离线量化)

原文链接：[2211.15736] Post-training Quantization on Diffusion Models (arxiv.org)

“本作品仅供学习交流使用，严禁用于商业用途。如果保存，请在24小时内删除。作者不承担任何责任。”

首先明确什么是PQT?

训练后量化 (PTQ) 算法将训练过的 FP32 网络直接转换为定点计算的网络，过程中无需对原始模型进行任何训练。只对几个超参数调整就可完成量化过程，量化模型以一种更有效的计算方式进行模型推理。量化后模型中的参数使用低 bit 表示，在数据搬移时降低了带宽要求，在计算过程中一般硬件对于低 bit 整形数据具有更高的标称算力，因此模型量化的优化方案在多数情况下可获得较大的推理速度提升，此方法已被广泛应用于大量的端侧和云侧部署场景。

简单来说，就是将预训练的网络的部分参数进行float32到int8甚至更低的转化，从而对模型进行轻量化。但这种转化必然造成模型的误差，于是引入校准数据集进行校准。

在训练后量化（Post-Training Quantization，PTQ）中，权重和激活值通常是在训练过程之后进行量化。量化是指将高精度的浮点数表示转换为低精度的整数或浮点数表示，通常是在8位或16位精度之间进行。
反量化是指将量化后的权重和激活值转换回浮点数表示，以便在推理过程中使用。在PTQ中，反量化通常是在推理过程中进行的，以便在使用量化后的权重和激活值进行计算时能够获得与使用浮点数表示相同的精度。
在PTQ中，权重和激活值的量化和反量化通常是在训练过程之后和推理过程中进行的，它们之间的关系如下：

1. 在训练过程中，权重和激活值通常是在高精度的浮点数表示中进行计算和存储的。
2. 在训练结束后，权重和激活值通常会被量化为低精度的整数或浮点数表示。
3. 在推理过程中，量化后的权重和激活值会被反量化为浮点数表示，以便进行计算。
   总的来说，PTQ中的量化和反量化是为了在保持计算精度的同时减少模型的存储和计算需求，从而提高模型的效率。

摘要

去噪扩散（基于得分的）生成模型最近在生成逼真且多样的数据方面取得了显著的成就。不幸的是，当前的去噪扩散模型的生成过程由于依赖于繁琐的神经网络的长时间迭代噪声估计而著名地慢。这阻碍了扩散模型在边缘设备上的广泛部署。以前的工作通过找到更短但有效的采样轨迹来加速扩散模型的生成过程。但是，他们忽视了每次迭代中使用重型网络进行噪声估计的成本。在这项工作中，我们从压缩噪声估计网络的角度来加速生成。由于重新训练DMs的难度，我们排除了主流的训练意识压缩范式，并将后训练量化（PTQ）引入DM加速。然而，噪声估计网络的输出分布会随着时间步长而变化，这使得以前的PTQ方法在DMs中失败，因为它们是为单时间步长场景设计的。为了设计一个针对DM的PTQ方法，我们从三个方面探索DM的PTQ：量化操作、校准数据集和校准度量。我们总结并使用从全面调查中得出的几个观察结果来制定我们的方法，特别是针对DM的独特多时间步长结构。实验上，我们的方法可以直接将全精度DM量化为8位模型，同时在无训练的情况下保持甚至提高其性能。重要的是，我们的方法可以作为其他快速采样方法的即插即用模块。

研究思路

本研究的研究思路是针对扩散模型（Diffusion Models）进行后训练量化（Post-training Quantization）。传统的量化方法在扩散模型上不适用，因为扩散模型具有复杂的扩散和去噪过程。本研究提出了一种名为PTQ4DM的后训练量化方法，通过网络压缩来加速扩散模型。

具体而言，本研究首先分析了扩散模型的操作选择，确定了哪些操作可以进行量化。

然后，研究团队探索了如何收集适合量化的校准数据集。

使用适当的方法来选择权重张量和激活张量的量化参数。

基于这些观察，研究团队设计了一种有效的校准方法，称为PTQ4DM。 通过PTQ4DM方法，研究团队成功地将扩散模型进行了后训练量化。实验结果表明，经过量化的扩散模型在性能上与完整精度的模型相当。这项研究解决了在扩散模型上进行后训练量化的问题，并提供了一种有效的量化方法。

1.哪些操作可以进行量化

与以前的大多数PTQ方法一样，网络中计算密集的卷积层和全连接层应该被量化。批量归一化可以折叠到卷积层中。SiLU和softmax等特殊功能保持完全精确。

扩散模型还有两个问题：1.网络的输出µ和∑是否可以量化？2.采样图像$x_{t-1}$是否可以被量化？为了回答这两个问题，我们只量化生成µ、∑或$x_{t-1}$的运算。如表1所示，我们观察到它们对量化不敏感，并且我们表明它们可以被量化。可以理解为不确定是否可以量化的参数就先量化了试一试效果。

2.如何收集适合量化的校准数据集

一般的PTQ的校准数据集可以从测试集中选取，而扩散模型中的训练数据集是$x_0$，这不是网络的输入。实际输入是生成的样本$x_t$。我们应该在扩散过程中使用生成的样本，还是在去噪过程中使用产生的样本？

为了量化扩散模型，我们需要在这种特定的多时间步长场景中设计一种新颖有效的校准数据集收集方法。

观察0：激活的分布随着时间步长的变化而变化。

观察1：在去噪过程中生成的样本对校准更有建设性。

观察2：接近真实图像x0的样本xt更有利于校准。

观察3：校准样本应以不同的时间步长生成，而不是在同一时间步长生成一组样本。

于是，在这种方法中，校准集$x_{t_i}$是通过去噪过程生成的，其中时间步长$t_i$是从偏斜正态分布中采样的。如图所示：

参数标定的探讨

当收集校准样本时，第三步是为扩散模型中的张量选择量化参数。在本节中，我们将探讨校准张量的度量。如下表所示，MSE优于L1距离、余弦距离和KL散度。因此，我们将MSE作为量化扩散模型的度量。

结论

两个正交因素减缓了去噪过程：i）从噪声中采样图像的漫长迭代，以及ii）在每次迭代中估计噪声的繁琐网络。与主流DM加速作品侧重于前者不同，我们的作品挖掘了后者。在本文中，我们提出了用于扩散模型的后训练量化（PTQ4DM），其中预训练的扩散模型可以直接量化为8比特，而不会经历性能的显著下降。重要的是，我们的方法可以添加到其他快速采样方法中，如DDIM[24]

ps：既然量化网络是也是为了加速，为什么没有采样速度以及对硬件的需求的对比呢，还是说目标仅仅就是把float量化为int从而实现优化，具体速度的加快不考虑？？