Temporal Dynamic Quantization for Diffusion Models阅读

Abstract

扩散模型因其出色的生成性能和多功能性在视觉应用中受到广泛关注。然而，由于模型大小和迭代生成带来的高存储和计算需求，限制了其在移动设备上的使用。现有的量化技术在8位精度下难以维持性能，因为扩散模型的激活具有时间变化的独特属性。我们引入了一种新颖的量化方法，该方法基于时间步信息动态调整量化区间，大幅提高了输出质量。与传统动态量化技术不同，我们的方法在推理过程中没有计算开销，并且兼容于训练后量化（PTQ）和量化感知训练（QAT）。我们的广泛实验表明，经过量化的扩散模型在各种数据集上输出质量显著提升。

什么是扩散模型？
扩散模型是一种生成模型，主要用于图像生成等任务。其基本思想是通过定义一个正向过程，将数据逐步加入噪声，直到接近纯噪声，然后学习逆过程，从噪声生成数据。扩散模型的结构一般包括以下部分：1. **正向扩散过程**：- 从原始数据开始，逐步添加噪声。- 每一步的小改动使数据分布逐渐变得模糊，直至接近高斯噪声。2. **逆向扩散过程**：- 学习去噪过程，通过逐步去除噪声回到数据分布。- 使用神经网络（例如UNet）预测和去除每一步的噪声。通常使用UNet等结构，因为它们擅长处理图像去噪任务。3. **训练**：- 通过最大化对数似然或最小化噪声预测误差来训练模型。扩散模型的关键在于其逐步变化的过程，使得生成的数据质量高且多样性好。它在图像生成、修复等任务中表现优异。

Introduction

生成建模在机器学习中对于图像、语音和文本合成等应用至关重要。扩散模型通过涉及数百次迭代推理的去噪过程逐步优化输入图像，因其优于GAN等替代方案的性能而受到关注。然而，扩散模型的高成本是其广泛应用的主要障碍。这些模型的尺寸较大，常达到数GB，并且单次图像生成需要大量迭代计算。因此，在资源有限的移动设备上执行扩散模型几乎是不可能的，大多数应用目前在昂贵的高性能服务器上实现。

为了充分利用扩散模型的潜力，已提出多种方法以减少其计算成本和内存需求，同时保持生成性能。例如，J. Song等人和L. Liu等人提出了更高效的采样调度方法，T. Salimans和J. Ho则提出使用知识蒸馏技术减少采样步骤。因此，可以用较少的采样步骤生成高质量图像，使得扩散模型的使用变得更具成本效益和可行性。尽管有这些进步，扩散模型的去噪过程仍需大量计算成本，因此需要进一步的性能增强和模型压缩。

虽然大多数先前的方法都集中在减少采样步骤以加速去噪过程，但减轻单个去噪步骤的负担也很重要。由于单个去噪步骤可以视为传统深度学习模型的推理，因此可以使用多种模型压缩技术。量化是一种广泛使用的压缩技术，它将权重和激活映射到低精度域。然而，由于扩散模型的独特性质，例如在迭代推理过程中激活分布的显著变化，随着激活位宽的减少，输出会严重失真。

为了解决扩散模型量化的独特挑战，我们引入了一种称为时间动态量化（TDQ）模块的新设计。该TDQ模块生成时间依赖的最佳量化配置，最小化激活量化误差。TDQ模块的显著优点在于可以无缝集成现有的QAT和PTQ算法，扩展这些算法以创建时间依赖的最佳量化配置，最小化激活量化误差。具体来说，该模块设计为在推理过程中不产生额外的计算开销，使其兼容于现有的加速框架而无需修改。TDQ模块通过生成优化的量化配置，显著提高了传统量化方案的质量。

Backgrounds and Related Works

2.1 扩散模型

扩散模型首次在2015年提出，革新了图像生成，通过将其描述为一个顺序去噪过程。如图1所示，正向扩散过程通过在每个时间步添加少量高斯噪声，将图像 (x_0) 逐渐转变为服从标准正态分布的随机数据 (x_T)。反向扩散过程则通过迭代去噪步骤，从随机数据 (x_T) 生成干净图像 (x_0)。因此，扩散模型学习的是反向过程，即在每个时间步 (x_t) 上估计给定噪声数据的噪声量。正向 (q) 和反向 (p_\theta) 过程可以描述为：

其中 (\beta_t) 表示高斯噪声的大小。

文献[3]引入了一种重参数化技巧用于 (\mu_\theta) 和相应的损失函数，促进了扩散模型的训练。

虽然扩散模型能够生成高质量的图像，但迭代去噪过程使其难以在实际场景中使用。早期的扩散模型研究，如DDPM[3]，需要数百到数千次迭代推理来生成一张图像，导致采样速度极慢。因此，众多研究致力于探索算法增强和各种优化，以提高性能，使扩散模型更高效且适用于实际应用。

DDIM[2]引入了一种隐式概率模型，通过重新解释DDPM方法的马尔可夫过程，以十分之一的去噪步骤实现了具有竞争力的图像质量。基于蒸馏的方法[16, 17]提出使用知识蒸馏技术减少去噪步骤的数量。

另一方面，为了解决生成高分辨率图像的显著计算开销，[4]提出了一种潜在扩散模型（LDM），其中扩散模型处理的是潜在变量而非图像。尤其是大规模扩散模型（如Stable Diffusion [4]）利用了LDM并从大规模数据集（LAION-Dataset [27]）中学习，实现了基于文本输入生成高质量、高分辨率图像。

2.2 量化

量化是一种显著的神经网络优化技术，通过低精度表示减少存储需求，并在可用低精度加速时提升性能。使用(b)位精度时，仅有(2^b)个量化级别可用，相较于浮点表示极大地限制了量化模型的自由度。因此，最终质量会显著依赖于量化超参数的调整，如量化区间或零偏移。为了在低精度下保持输出质量，至关重要的是在更新模型参数时，结合目标任务的特定特征和要求调整量化参数。

2.3 量化感知训练和训练后量化

量化算法大致可分为两类：量化感知训练（QAT）[19, 28, 29, 20, 30]和训练后量化（PTQ）[23, 24, 25]。QAT在引入量化算子后进行额外训练，允许更新网络参数，以考虑量化算子的影响来最小化最终损失值。而PTQ在量化后不进行端到端的前向/后向传播，而是专注于减少量化引起的块级重构误差。QAT在低精度场景中通常优于PTQ，但由于数据集不足、训练管道限制和资源约束等原因，QAT可能并不总是适用。由于其实用性，PTQ近年来被积极研究。在扩散模型的文献中，[18]引入了一种专门的8位训练后量化方法，展示了高保真度的图像生成性能。

另一方面，尽管大多数QAT和PTQ算法专注于静态量化，近期研究强调了输入依赖的动态量化的优势[31, 32, 33, 34]。动态量化能够根据激活的输入依赖分布调整量化区间，具有减少量化误差的潜力。然而，实施这些动态方法通常会因提取激活统计信息而产生额外成本，使得实际中难以实现性能提升。例如，[35]指出输入依赖的激活函数可能导致显著的延迟增加，尽管其计算量相对较小。

尽管先前有多项工作提出加速扩散模型采样过程的不同方法，但很少有工作尝试利用扩散模型的动态特性。本文提出了一种新颖的扩散模型量化方案，通过基于时间步信息生成合适的量化区间，最大限度减少激活量化误差，同时不产生额外的推理成本。

TemporalDynamic Quantization

3.1 量化方法

在详细说明所提方法之前，我们定义本文中使用的量化函数。本研究专注于(b)位线性量化，其中(2^b)个可能的量化级别是均匀分布的。线性量化涉及两个关键超参数：量化区间(s)和零偏移(z)。给定全精度数据(x)，量化后的数据(\hat{x})可按如下计算：

其中(n§)是最小（最大）量化索引，(\text{clip}(·))是裁剪函数，(\lfloor·\rceil)是舍入函数。为了实际加速，我们对权重使用对称量化，其中(z=0)且(p=-n=2^{{(b-1)}-1)。另一方面，对于激活量化，我们假设使用不对称量化，其中(z=0)、(n=0)和(p=2}b-1)，这是一种常用方法[36]。

3.2 扩散模型量化的挑战

量化扩散模型的最大挑战在于为激活找到最优量化参数（(s)和(z\）以最小化量化误差。如图2所示，扩散模型的激活分布由于迭代去噪过程具有独特特性，其分布高度依赖于时间步（(t)），而与层索引无关。

因此，使用静态值作为量化参数会在不同时间步引发显著的量化误差，如图3所示。先前研究[37,18]也报告了扩散模型中激活的动态特性，并尝试通过在所有时间框架内采样校准数据集来解决这一问题。然而，尽管有这些努力，这些研究仍依赖于静态参数，导致在最小化量化误差时收敛效果欠佳。要从根本上解决这个问题，至关重要的是能够根据输入激活分布的动态性更新量化参数。

3.3 TDQ模块的实现

为应对输入激活的快速变化，一种简单思路是输入依赖的动态量化。先前研究表明，结合生成量化参数的量化模块（基于输入特征如最小值、最大值和平均值）可以显著提高特定应用的准确性[38,32]。

根据我们的观察（详见附录），我们自己实现的输入依赖动态量化在质量上有显著提升。然而，收集给定激活的统计数据会导致复杂的实现和显著的开销[35]。尽管存在质量提升的潜力，这种方法在实际中可能不具吸引力。因此，我们提出了一种新颖的动态量化，利用时间信息而非输入激活。我们发现，当测量时间步与每个张量激活变化的皮尔逊相关系数时，62.1%的层表现出中等的时间依赖（|r| > 0.5），38.8%表现出强时间依赖（|r| > 0.7）。这表明扩散模型中的许多层确实具有强烈的时间依赖性。尽管激活分布可能根据输入数据变化，但在同一时间框架内整体趋势相似（见图2）。因此，我们可以基于时间步的差异确定最优区间，从而获得更可靠和稳健的量化结果。

直观上，我们可以为每个时间步采用一个标量量化区间。然而，这种策略难以捕捉时间间关系，且不能应用于训练和推理使用不同时间步的情况。因此，我们提出使用一个可学习的小型网络，通过时间步信息预测量化区间。该策略可以跨多个时间步整体评估激活变化，从而使学习过程更稳定并提高性能，如表3所示。

图4（b）展示了我们的想法概览，称为TDQ模块。在TDQ模块中，动态区间(\tilde{s})基于提供的时间步(t)生成，如下所示：

其中(\text{enc}(·))表示时间步的编码函数，将在3.4节中描述。这里，(I)是编码特征，(f(·))表示生成模块函数。

在我们的方法中，TDQ模块附着于每个量化算子，独立的动态量化区间基于给定时间步生成。如图4所示，生成器通过简单堆叠多个线性层和一个softplus函数实现，以将数据范围限制为非负值。请注意，生成器的所有组件都是可微的。因此，在PTQ或QAT过程中，可以通过梯度下降更新区间以最小化量化误差。

例如，当使用知名QAT算法如LSQ [29]时，静态区间可以用TDQ模块的输出替代。方程5中的量化函数修改如下：

当我们使用直通估计器[39]时，梯度可以通过生成器的可学习参数传播。这些参数然后通过迭代梯度下降更新，目标是最小化最终任务损失。同样的流程可以应用于利用局部梯度最小化重构误差的PTQ算法[23, 24]。因此，TDQ模块很容易适用于QAT和PTQ方案。

然而，与输入依赖动态量化不同，TDQ模块实现了无成本推理。在PTQ或QAT过程后，时间依赖的区间可以离线预计算，在推理期间，我们可以利用预计算的量化区间。这是TDQ模块在现有框架中无须修改即可实现无缝集成的一大优势。

3.4 工程细节

我们进行了系列实验以增强TDQ模块的稳定性和效果，发现一些工程改进在实现可靠结果方面起到了关键作用。本节详细描述了这些改进及其对TDQ模块性能的影响。

时间步的频率编码

直接将时间步输入生成器会导致收敛质量较差。这主要是由于神经网络的低频归纳偏差[40, 41]。如果直接将时间步输入生成器，它往往会生成一个几乎不随时间步变化的区间。为减轻这种低频偏差，我们对时间步使用几何傅里叶编码[42, 43]，如方程6所示：

其中(t)是当前时间步，(d)是编码向量的维度，(I)是编码后的向量。这种编码方法允许TDQ模块适应时间步的高频动态。在本文中，我们经验性地将(t_{\text{max}})设置为10000。

TDQ模块的初始化

量化区间的适当初始化至关重要，因为初始化不当会导致QAT或PTQ过程中的不稳定。现有的量化技术仅需初始化静态步长值，但我们需要将TDQ模块的输出初始化为期望值。为此，我们对TDQ模块（MLP）的最后一层线性层的权重使用He初始化[44]，并将偏差设置为期望值。由于MLP的输入（几何傅里叶编码）可视为均值为零的随机变量，经过He初始化的MLP的输出也将具有零均值。因此，我们可以通过偏差调整将MLP输出的均值控制为期望值。提取整个时间步的1000个样本后，我们初始化量化区间以最小化整体误差，然后进行更新以适应每个时间步。

Experimental Setup

为了展示TDQ的卓越性能，我们在两个不同的模型上进行了测试：DDIM[2]（像素空间扩散模型）和LDM[4]（潜在空间扩散模型）。在DDIM实验中，我们使用了CIFAR-10数据集[45]（32x32），而在LDM实验中，我们采用了LSUN Churches数据集[46]（256x256）。这使我们能够在低分辨率和高分辨率图像生成场景中展示所提方法的有效性。

我们将PTQ和QAT应用于这两个模型。然而，值得注意的是，虽然潜在扩散模型由VAE和扩散模型组成，但我们专注于量化扩散模型，并未对VAE部分进行量化。

由于缺乏扩散模型的先前QAT研究，我们尝试了知名的静态量化方法，如PACT[19]、LSQ[29]和NIPQ[30]作为基线。我们的构想是将TDQ模块的输出替换静态区间，集成于LSQ之上。我们在附录中提供了TDQ与其他方法集成的额外实验。对所有卷积层和线性层的激活和权重，包括注意力层的激活，进行逐层量化。模型在CIFAR-10和LSUN-churches上分别训练了200K次迭代，批量大小分别为128和32。学习率与全精度模型一致。

在PTQ中，我们使用了PTQ4DM[18]作为基线进行广泛分析。为公平比较，我们遵循PTQ4DM的实验设置，但修改了激活量化算子，使用TDQ模块生成动态量化区间。与PTQ4DM类似，我们对权重使用逐通道不对称量化，对激活使用逐张量不对称量化。权重量化范围通过逐通道最小/最大值确定，激活量化范围通过梯度下降训练以最小化分块重构误差，如BRECQ[24]等方法。量化应用于所有层，但在LDM的PTQ实验中，注意力矩阵的激活未量化。按照PTQ4DM的方法，我们使用由5120个样本组成的校准集进行PTQ，包括256张图像，每张图像随机选择20个时间步。

为了衡量模型性能，我们使用Fréchet Inception Distance（FID）[47]和Inception Score（IS）[48]评估CIFAR-10，并使用FID评估LSUN-churches。在评估中，我们使用QAT和DDIM的200步采样生成50,000张图像，使用PTQ和DDIM的100步生成50,000张图像。在QAT的情况下，我们选择了验证损失最低的5个检查点，并报告了表现最佳模型的分数。

所有实验在高性能服务器上进行，配备4个A100 GPU和8个RTX3090 GPU，使用PyTorch[49] 2.0框架。源代码可在https://github.com/ECoLab-POSTECH/TDQ_NeurIPS2023获取。此外，我们使用WxA表示x位权重和y位激活量化以简化表达。

Results

5.1 QAT和PTQ后的质量分析

表1比较了TDQ模块与现有静态量化方法。随着激活位数减少，所有静态量化方法的质量都下降，而我们的方法则保持一致的输出。即使在8位精度下，TDQ模块也显著提高了质量，在4位精度下甚至能将质量降至全精度输出的水平。TDQ通过有效分配有限的激活量化级别实现了这些优势。

此外，NIPQ通过基于人工噪声的伪量化来解决直通估计器的不稳定性问题。然而，NIPQ的噪声与输入噪声难以区分，阻碍了扩散模型的收敛。要利用基于PQN的QAT在扩散模型中的优势，需付出额外努力。

表2展示了TDQ模块与现有静态PTQ方案的比较。MinMax方法代表一种简单的线性量化方法，其范围由目标张量的最小值和最大值确定。实验表明，所有基线在激活位数高时保持良好的FID水平，但随着激活位数减少，性能显著下降。相比之下，TDQ仅表现出轻微的FID下降，表明TDQ模块在QAT和PTQ场景中均表现出色。

图5到图7展示了量化扩散模型生成的图像。在相同位宽配置下，我们的方法在生成高保真图像方面始终优于其他量化技术。图5和图6进一步说明了这一点，传统的QAT和PTQ生成模糊且不可识别的图像，而我们的方法生成逼真的图像。激活量化中融入时间信息在保持输出感知质量方面非常有效。

5.2 TDQ模块的泛化性能

本节展示了TDQ模块在快速推理中的性能实验结果。训练涵盖所有时间步（1到1000），但推理可以在较少时间步（50到100）下执行以提高性能。这改变了训练/测试时的时间步分布，因此TDQ模块的生成需要良好的泛化能力。

图8显示了在W8A4和W4A4配置下，通过LSQ算法量化的DDIM模型的FID测量。推理过程中时间步从100逐渐减少到10。如图所示，随着时间步减少，LSQ的性能显著下降，而我们的方法性能的下降与全精度基线相似。该实验表明，即使采样时间步变化，TDQ模块依然有效运作。

5.3 TDQ模块的消融研究

为研究TDQ模块输出的动态性，我们可视化了与时间步相关的动态区间更新（图9）。在DDIM上使用W4A4 LSQ训练的区间显示出与激活变化一致的趋势。然而，这种模式并非所有层都一致。此不一致可能表明TDQ模块尝试生成一个最小化最终损失值的区间，因为LSQ相应地调整了量化区间。

为提供对TDQ模块优点的全面分析，我们将其与输入依赖的动态量化方法进行了比较，并对TDQ模块进行了消融研究。如表4所示，即使不直接利用层的分布信息，TDQ仍表现出优越的性能。我们还研究了TDQ模块中层数对性能的影响。结果表明，当层数超过4时，输出质量显著稳定。基于这些发现，我们经验性地选择了4层MLP，以平衡性能和复杂性。

我们也进行了实验，将TDQ与每个时间步直接学习的量化区间配置进行比较。如表3所示，直接为每个时间步学习量化区间并在多个时间步中使用共享量化区间的概念相比基线（S1）可提高输出质量，但TDQ依然表现出卓越的输出质量和更大的灵活性。TDQ的成功在于其能够在考虑相邻时间步激活分布演变的同时，促进量化区间的连续稳定学习。

Discussion

在本节中，我们讨论了一些额外的直觉和局限性。虽然许多层表现出强烈的时间依赖性，但大约30%的层并未表现出这种相关性。这些层主要位于U-Net的中间块，并受到实例语义信息的强烈影响。此外，我们观察到LDM模型的时间依赖性较少，因此性能提升较小。然而，重要的是要强调，即使在这些情况下，TDQ模块也能确保收敛，在各个时间步上保持一致的输出。这种一致性保证了即便时间依赖性较弱，输出质量仍与现有PTQ/QAT算法相当。

Conclusion

在本文中，我们探讨了扩散模型中激活量化的挑战，特别是时间步中激活的动态变化，并表明现有的静态量化方法未能有效解决此问题。我们介绍了TDQ模块，将其无缝集成到QAT和PTQ中，并基于动态量化区间生成显著提升输出质量。由于我们的方法在推理过程中没有额外开销，我们期望这项研究能在移动和边缘设备的低比特扩散模型中提高性能，同时保持质量。

思考和心得

我看这个期刊最近关于扩散模型相关的量化还是很多的。因为目前我关注的都是量化这一块。个人感觉，量化要最求效果好就势必涉及底层的优化，要追求创新就必须得寻求新的场景（比如这个扩散模型）。唉还好今天得知我之后走的是工程向，我觉得论文还是挺难发的，可恶。