[Paper Review] Temporal Dynamic Quantization for Diffusion Models

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

Stochastic Gradient Optimization Techniques4 citations

TL;DR

This paper proposes Temporal Dynamic Quantization (TDQ), a novel quantization method for diffusion models that dynamically adjusts quantization intervals based on time step information, significantly improving generation quality. Unlike static or conventional dynamic quantization, TDQ incurs no inference overhead and enhances both post-training quantization (PTQ) and quantization-aware training (QAT), achieving state-of-the-art performance at low bit-widths (e.g., W4A4 FID of 4.48 on CIFAR-10).

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.

Motivation & Objective

• Address the performance degradation in low-bit quantization of diffusion models due to time-varying activation distributions across iterative denoising steps.
• Overcome limitations of static and conventional dynamic quantization, which fail to adapt to temporal shifts in activation statistics.
• Develop a quantization method that generates time-dependent optimal quantization parameters without introducing computational overhead during inference.
• Ensure compatibility with both post-training quantization (PTQ) and quantization-aware training (QAT) for broad deployment on mobile and edge devices.
• Achieve high-fidelity image generation at low precision (e.g., 4-bit weights and activations) while maintaining performance close to full-precision baselines.

Proposed method

• Introduce the TDQ module, a learnable neural network that predicts quantization parameters (scale and zero-point) for each time step based on the current denoising step.
• Design the TDQ module to be seamlessly integrated into existing PTQ and QAT pipelines without modifying inference computation graphs.
• Train the TDQ module end-to-end using a differentiable quantization scheme (e.g., LSQ) to minimize activation quantization error across all time steps.
• Use a multi-layer perceptron (MLP) with 4 layers as the TDQ backbone, empirically selected for optimal trade-off between performance and complexity.
• Generate time-dependent quantization intervals that adapt to the evolving activation distribution during the reverse denoising process.
• Ensure inference efficiency by computing quantization parameters only once per time step during the forward pass, avoiding runtime computation.

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

Experimental results

Research questions

• RQ1 Can dynamic quantization that adapts to time-step-specific activation distributions significantly improve low-bit quantization performance in diffusion models?
• RQ2 Does a time-dependent quantization strategy outperform static and input-dependent dynamic quantization in maintaining generation quality at low bit-widths?
• RQ3 Can the proposed TDQ module be effectively integrated into both PTQ and QAT frameworks without introducing inference overhead?
• RQ4 How does the number of layers in the TDQ module affect the final model performance and convergence stability?
• RQ5 To what extent does TDQ maintain performance under varying sampling schedules (e.g., reduced number of steps)?
• RQ6 How do the quantization intervals evolve across time steps, and do they correlate with activation dynamics?

Key findings

• TDQ achieves a FID of 4.48 on CIFAR-10 using W4A4 quantization with DDIM, significantly outperforming LSQ (7.30) and other baselines.
• On the LDM-Churches dataset, TDQ achieves a W4A4 FID of 4.64, outperforming static quantization (5.06) and dynamic quantization (5.17).
• With W3A3 quantization, TDQ achieves a FID of 6.48 on CIFAR-10, substantially better than LSQ (7.63) and other dynamic schemes.
• The ablation study shows that 4-layer TDQ achieves the best balance of performance and complexity, with FID dropping to 4.48 (vs. 5.18 for 2-layer and 4.88 for 8-layer).
• TDQ maintains consistent performance even when the number of sampling steps is reduced from 100 to 10, unlike LSQ, which degrades sharply under such conditions.
• Visualization confirms that TDQ intervals adapt to activation variations across time steps, with strong temporal correlation in ~70% of layers, especially in attention and residual blocks.