[论文解读] Temporal Dynamic Quantization for Diffusion Models

Junhyuk So, Jungwon Lee+3 more|arXiv (Cornell University)|Jun 4, 2023

Stochastic Gradient Optimization Techniques被引用 4

一句话总结

本文提出时间动态量化（Temporal Dynamic Quantization, TDQ），一种针对扩散模型的新量化方法，其根据时间步信息动态调整量化区间，显著提升生成质量。与静态或传统动态量化不同，TDQ在推理过程中无额外开销，并同时提升后训练量化（PTQ）与量化感知训练（QAT），在低比特宽度下实现最先进性能（例如，在CIFAR-10上W4A4 FID为4.48）。

ABSTRACT

The diffusion model has gained popularity in vision applications due to its remarkable generative performance and versatility. However, high storage and computation demands, resulting from the model size and iterative generation, hinder its use on mobile devices. Existing quantization techniques struggle to maintain performance even in 8-bit precision due to the diffusion model's unique property of temporal variation in activation. We introduce a novel quantization method that dynamically adjusts the quantization interval based on time step information, significantly improving output quality. Unlike conventional dynamic quantization techniques, our approach has no computational overhead during inference and is compatible with both post-training quantization (PTQ) and quantization-aware training (QAT). Our extensive experiments demonstrate substantial improvements in output quality with the quantized diffusion model across various datasets.

研究动机与目标

• 解决扩散模型在低比特量化中因迭代去噪过程中激活分布随时间变化而导致的性能下降问题。
• 克服静态量化与传统动态量化在适应激活统计量时间漂移方面的局限性。
• 开发一种无需推理阶段引入计算开销即可生成与时间相关最优量化参数的量化方法。
• 确保与后训练量化（PTQ）和量化感知训练（QAT）的兼容性，以支持在移动设备与边缘设备上的广泛部署。
• 在低精度（如4比特权重与激活）下实现高保真图像生成，同时保持与全精度基线相近的性能。

提出的方法

• 引入TDQ模块，一个可学习的神经网络，基于当前去噪步骤预测各时间步的量化参数（缩放因子与零点）。
• 设计TDQ模块以无缝集成至现有PTQ与QAT流程中，无需修改推理计算图。
• 采用可微量化方案（如LSQ）端到端训练TDQ模块，以最小化所有时间步的激活量化误差。
• 选用含4层的多层感知机（MLP）作为TDQ主干网络，经实证选择以实现性能与复杂度之间的最佳平衡。
• 生成随时间变化的量化区间，以适应反向去噪过程中激活分布的演化。
• 通过在前向传播中仅对每个时间步计算一次量化参数，确保推理效率，避免运行时计算。

Figure 1: The forward and reverse processes of diffusion models.

实验结果

研究问题

• RQ1 能否通过适配时间步特定激活分布的动态量化方法，显著提升扩散模型在低比特量化下的性能？
• RQ2 与静态量化和输入依赖型动态量化相比，时间依赖的量化策略是否能在低比特宽度下更好地维持生成质量？
• RQ3 所提出的TDQ模块是否能有效集成至PTQ与QAT框架中，且不引入推理开销？
• RQ4 TDQ模块中的层数如何影响最终模型性能与收敛稳定性？
• RQ5 在不同采样调度（如减少采样步数）下，TDQ的性能保持程度如何？
• RQ6 量化区间如何随时间步演化？其变化是否与激活动态存在相关性？

主要发现

• 在CIFAR-10上，使用DDIM的W4A4量化，TDQ实现FID为4.48，显著优于LSQ（7.30）及其他基线方法。
• 在LDM-Churches数据集上，TDQ实现W4A4 FID为4.64，优于静态量化（5.06）与动态量化（5.17）。
• 在W3A3量化下，TDQ在CIFAR-10上实现FID为6.48，显著优于LSQ（7.63）及其他动态方案。
• 消融实验表明，4层TDQ在性能与复杂度之间达到最佳平衡，FID降至4.48（2层为5.18，8层为4.88）。
• 当采样步数从100减少至10时，TDQ保持一致性能，而LSQ在此条件下性能急剧下降。
• 可视化结果证实，TDQ区间能随时间步的激活变化自适应调整，约70%的层中表现出强时间相关性，尤其在注意力与残差模块中。