[Paper Review] PTQD: Accurate Post-Training Quantization for Diffusion Models

Yefei He, Luping Liu+4 more|arXiv (Cornell University)|May 18, 2023

Machine Learning in Materials Science13 citations

TL;DR

PTQD introduces a unified post-training quantization framework for diffusion models that disentangles quantization noise into correlated and uncorrelated parts, corrects them, calibrates the variance schedule, and uses step-aware mixed precision to preserve SNR, achieving near full-precision quality with substantial bit-operations speedups.

ABSTRACT

Diffusion models have recently dominated image synthesis tasks. However, the iterative denoising process is expensive in computations at inference time, making diffusion models less practical for low-latency and scalable real-world applications. Post-training quantization (PTQ) of diffusion models can significantly reduce the model size and accelerate the sampling process without re-training. Nonetheless, applying existing PTQ methods directly to low-bit diffusion models can significantly impair the quality of generated samples. Specifically, for each denoising step, quantization noise leads to deviations in the estimated mean and mismatches with the predetermined variance schedule. As the sampling process proceeds, the quantization noise may accumulate, resulting in a low signal-to-noise ratio (SNR) during the later denoising steps. To address these challenges, we propose a unified formulation for the quantization noise and diffusion perturbed noise in the quantized denoising process. Specifically, we first disentangle the quantization noise into its correlated and residual uncorrelated parts regarding its full-precision counterpart. The correlated part can be easily corrected by estimating the correlation coefficient. For the uncorrelated part, we subtract the bias from the quantized results to correct the mean deviation and calibrate the denoising variance schedule to absorb the excess variance resulting from quantization. Moreover, we introduce a mixed-precision scheme for selecting the optimal bitwidth for each denoising step. Extensive experiments demonstrate that our method outperforms previous post-training quantized diffusion models, with only a 0.06 increase in FID score compared to full-precision LDM-4 on ImageNet 256x256, while saving 19.9x bit operations. Code is available at https://github.com/ziplab/PTQD.

Motivation & Objective

• Motivate the need for post-training quantization (PTQ) of diffusion models to reduce memory and compute without re-training.
• Develop a unified noise model that separates quantization noise from diffusion denoising perturbations.
• Provide corrective mechanisms for both correlated and uncorrelated quantization noise during sampling.
• Introduce a step-aware mixed-precision strategy to maintain high SNR across denoising steps.

Proposed method

• Model quantization using uniform quantization with explicit noise notation.
• Disentangle quantization noise into a correlated part k*epsilon_theta and an uncorrelated residual (Delta epsilon_theta') as in Eq. (7).
• Correlated Noise Correction by dividing the correlated component output by 1+k (Eq. (9)).
• Uncorrelated Noise Correction via Bias Correction (BC) and Variance Schedule Calibration (VSC) (Eq. (10)-(12)).
• Estimate correlation k and uncorrelated noise statistics from quantized vs. FP runs (Algorithm 1).
• Step-aware Mixed Precision selects per-step activation bitwidths from a set B to satisfy SNR^Q(t) > SNR^F(t) (Eq. (13)-(15)).

Experimental results

Research questions

• RQ1 How does quantization noise affect mean and variance in diffusion denoising steps?
• RQ2 Can a unified decomposition of quantization noise into correlated and uncorrelated parts improve PTQ for diffusion models?
• RQ3 Can bias and variance be corrected without re-training to recover sampling quality?
• RQ4 Does step-aware mixed precision preserve SNR across denoising steps while maximizing speedups?

Key findings

• Disentangling quantization noise into correlated and uncorrelated parts enables targeted corrections.
• Correlated Noise Correction (CNC) reduces FID by 0.48 and sFID by 6.55 in ablations.
• Bias Correction (BC) and Variance Schedule Calibration (VSC) further reduce FID by 0.2 and 0.11 in sFID in ablations.
• PTQD achieves FID 6.44 and sFID 8.43 with W4A4/W4A8 mixed precision, only 1.33 sFID worse than FP, while saving 19.9x bit-operations.
• On ImageNet 256x256, PTQD reduces FID by about 0.06 compared to full-precision LDM-4 with 250 steps, while maintaining much smaller model size and significant BOP reductions.
• Step-aware Mixed Precision maintains higher SNR across steps, enabling effectively low-bit diffusion without large quality loss.