本文目录

1 Q-DM：低比特扩散模型
(来自北航，上海 AI Lab)
1 Q-DM 论文解读
1.1 把量化感知训练引入扩散模型
1.2 扩散模型背景简介
1.3 量化过程简介
1.4 扩散模型量化的挑战
1.5 实验结果

太长不看版

本文探索了扩散模型的低比特量化方案，是对扩散模型进行量化的先驱工作之一。量化 (Quantization) 是一种通过低比特精度参数和运算来显著减小内存消耗和计算消耗的技术。本文研究的是如何将量化技术应用到扩散模型。

本文观察到在扩散模型中应用低比特量化技术得到的结果很差，原因是作者发现低比特扩散模型性能存在瓶颈的原因是：多步去噪过程引起的累计量化误差，以及激活值分布的振荡剧烈。本文提出 Timestep-aware Quantization (TaQ) 和 Noise-estimating Mimicking (NeM) 来针对消除振荡，以及累积误差，最终达到帮助训练低比特的扩散模型的目的。

通过这种方式，本文提出一种有效的扩散模型量化方法 Q-DM，在 DDPM 和 DDIM 两个框架上做了验证。

本文做了什么工作

1. 提出了一种针对扩散模型的 QAT 量化方法 Q-DM。
2. 提出 Timestep-aware Quantization (TaQ) 方法减少输出分布的振荡，提出 Noise-estimating Mimicking (NeM) 策略来减小累积误差。
3. 作者在 CIFAR-10 和 ImageNet 上做了验证，Q-DM 相比于 8 位 PTQ 方法，性能有大幅提升，并以相当大的加速率实现了与全精度模型相当的性能。

1 Q-DM：低比特扩散模型

论文名称：Q-DM: An Efficient Low-bit Quantized Diffusion Model (NeurIPS 2023)

论文地址：

https://openreview.net/pdf?id=sFGkL5BsPi

1.1 把量化感知训练引入扩散模型

扩散模型 (Diffusion Models, DM) 在各种生成式任务中取得了显著的成功。扩散模型的过程涉及前向过程，即逐渐将真实数据转换为高斯噪声；以及反向过程，即逐步去噪以生成真实的数据。扩散模型已被证明在生成质量和多样性方面表现很出色。但是，DM 的去噪过程是非常耗时的，并且要在数千个 time steps 上不断迭代通过模型来进行噪声的估计。

因此，研究人员希望能够加速这个生成过程，要么减少很长且耗时的迭代过程，要么减少样本生成的高推理成本。这就对应了两种扩散模型压缩的技术路线：开发更快的采样策略，以及压缩和加速噪声估计网络。

那么量化技术就是后者的一种很好的选择，因为 low-bit 的参数和计算很适合 AI 芯片。量化技术可以分为训练后量化 (Post-Training Quantization, PTQ) 和量化感知训练 (Quantization-Aware Training, QAT)。PTQ 直接基于预训练的全精度模型计算量化参数，这种方法操作简单，无需微调，但是量化结果的性能往往不令人满意。而且，这种方法在超低比特 (4比特或更低) 的情况是无效的，且会伴随着性能的显著下降。PTQ 在 DM 上的应用已有先例[1]。QAT 在训练过程中进行量化，且通常以比 PTQ 可以实现更高的压缩率，和更少的性能下降。但是这种方法需要从头训练模型，对计算资源的要求较高，不适合很大的模型，例如 LLMs。QAT 的代表方法是 LSQ[2]。QAT 以被证明对 CNN，ViT 和 BERT 都好用，但对 DM 的探索还很少。因此，本文致力于构建一个使用 QAT 的低比特的量化 DM 基线。

1.2 扩散模型背景简介

1.2.1 扩散模型的前向过程

假设 是从数据分布 当中采样的一个样本, 正向扩散过程将高斯噪声添加到样本 次, 得到一系列噪声样本 为:

式中 是在每一步控制高斯噪声的强度的超参数。前向扩散过程满足马尔可夫性质, 因为每一步都仅依赖于前面一步。此外, 随着步数增加到无穷大 , 最终状态 收玫到各向同性高斯分布。前向过程的一个显着特性是它允许以封闭形式在任意时间步 采样 :

1.2.2 扩散模型的反向过程

反向扩散过程是从高斯噪声输入 中生成样本。但是, 由于实际的反向条件分布 是未知的, 扩散模型使用可学习的条件分布 来逼近真实的反向条件概率分布, 并且可以表示为:

通过使用 re-parameterization trick, 和 可以表示为:

其中, 是一个函数逼近器, 旨在从 中预测 。

1.2.3 扩散模型的训练目标

训练时, 优化的目标是最小化负对数似然, 即 。通过变分推理, 可以找到它的下界, 表示为 :

使用简化的损失函数来替代 通常可以获得更好的性能:

1.2.4 采样过程

在采样过程中, 扩散模型对高斯噪声张量 进行采样, 并通过重复采样反向分布 去噪。将 作为最终生成结果, 最后一步去噪中不添加噪声。

1.3 量化过程简介

假设 和 分别是一个 CNN 网络第 层的卷积权重和输入特征。

量化神经网络旨在以低位格式表示 和 , 使得浮点卷积输出可以近似为：

其中, ○代表 channel-wise的相乘, 代表 GEMM 乘法。

代表了量化超参数 scaling factor, 是 channel-wise 的。本文对输入激活使用逐层量化, 每层的 是一个标量。

1.4 扩散模型量化的挑战

作者列举出了2点扩散模型量化的挑战，尤其针对于扩散模型的多步推理过程和随机采样步骤训练过程，它们分别是：

1) 训练过程中激活值的分布振荡： 为了了解扩散模型的分布变化，作者研究了训练过程中的激活分布与 time step 的关系。理论上，如果这个分布会随着 time step 发生变化，那就会很难实现以前的 QAT 方法。作者分析了噪声估计网络的整体激活分布，如图1和图2所示。可以看到，不同的 time step 对应激活值的分布有很大的差异。为之前的 QAT 量化方法 (如[3]等) 在扩散模型的应用带来了不小的难度。

2) 推理时多步去噪过程产生的累积量化误差

噪声估计网络的量化会给权重和激活带来干扰，导致每一层的输出产生错误。在扩散模型中，每个 time step 模型的输入 是从模型在前一个 time step 的输出 中获得的，如图3所示，作者通过 MSE 距离来表示低比特扩散模型的量化误差。可以看到采样的进行，即 time step 的不断下降，这个 MSE 误差出现显著的增长。这意味着随着推理时多步去噪过程的进行，量化误差的累积将更加突出。

1.4 Q-DM 方法介绍

1.4.1 Timestep-aware Quantization

为了解决训练过程中的分布振荡问题，作者设计了考虑 time step 的量化注意力机制，这种结构允许对激活值的范围进行数值分析，以减轻量化扩散模型的分布振荡。

式中, 是注意力分数。

Timestep-aware Quantization 可以写成：

其中，

式中, 是第 个 Batch Size, 是 Batch 的数量。 和 分别是 query 激活值的统计均值和方差。

实际上, 可以看作是先做了下面的平滑计算之后, 再进行量化:

这些平滑之后的激活值对于训练过程中随时间采样不敏感，因此 Timestep-aware Quantization 在一定的程度上可以削减分布振荡。

1.4.2 Noise-estimating Mimicking

为了减轻量化误差的累积对训练量化扩散模型 带来的负面影响, 作者引入了全精度模型 来进行蒸馏。

低比特量化 DM 模型: 。

全精度 DM 模型: 。

Noise-estimating Mimicking 的目标函数可以写成:

其中, 是与 和