PaperlayerView exact revision on arXiv

Source-linked AI summary

PTQD: Accurate Post-Training Quantization for Diffusion Models

Yefei He, Luping Liu, Jing Liu, Weijia Wu, Hong Zhou, Bohan Zhuang

arXiv:2305.10657v4cs.CV

TL;DR

Diffusion sampling is computationally expensive, and directly applying existing PTQ methods can degrade low-bit sample quality through accumulated quantization noise. PTQD unifies quantization and diffusion noise correction with step-aware mixed precision, achieving near-full-precision quality while substantially reducing bit operations.

  • Problem

    Low-bit post-training quantization can impair diffusion sample quality because quantization noise alters denoising means, variances, and later-step SNR.

  • Method

    PTQD separates quantization noise into correlated and residual uncorrelated parts, applies correlation, bias, and variance corrections, and adaptively selects denoising-step bitwidths.

  • Results

    PTQD reaches state-of-the-art post-training quantization performance, with a 0.06 increase in FID versus full-precision LDM-4 on ImageNet 256x256 while saving 19.9x bit operations.

  • Takeaways & Limitations

    PTQD provides a unified framework for accurate, efficient post-training quantization of diffusion models without retraining.

  • Takeaways & Limitations

    PTQD can still face low SNR in the quantized noise prediction network, and like other generative models may enable counterfeit media for malicious purposes.

Abstract

from arXiv · show

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.

1 Introduction

Diffusion models offer high-quality generation but remain costly because sampling requires many denoising steps and substantial memory and computation. PTQD addresses low-bit quantization errors with noise correction and step-aware mixed precision.

  • Motivation: Hundreds of denoising steps make diffusion sampling significantly slower than GAN-based generation.
  • Motivation: 16GB of running memory and over 10GB of VRAM are required to execute Stable Diffusion, limiting deployment across devices.
  • Post-training quantization: 8-bit models provide a 2.2× speed-up over floating-point models on ARM CPUs, while 4-bit quantization offers up to 59% higher throughput than 8-bit quantization.
  • PTQD: PTQD disentangles quantization noise into correlated and residual uncorrelated components, correcting them through correlation estimation, bias subtraction, and variance schedule calibration.
  • PTQD: Step-aware mixed precision dynamically selects bitwidths for synonymous denoising steps to preserve signal-to-noise ratio throughout sampling.
  • Results: PTQD reaches new state-of-the-art performance for post-training quantization of diffusion models in extensive experiments.

2 Related Work

Prior work improves diffusion efficiency through faster samplers and studies model quantization, but few methods specifically address post-training quantization of pretrained diffusion models.

  • Efficient diffusion models: Efficient diffusion research includes retraining methods such as distillation and scheme learning, alongside training-free advanced samplers.
  • Model quantization: Post-training quantization reduces memory and computation without fine-tuning, requiring only a small amount of unlabeled calibration data.
  • Model quantization: Mixed precision accelerates inference by representing part of a model with lower bitwidths, commonly selected using Hessian spectra or Pareto frontiers.
  • Diffusion-model quantization: Only a few studies quantize pretrained diffusion models without retraining; PTQ4DM targets 8-bit models, while Q-Diffusion evaluates broader datasets using advanced PTQ techniques.
  • Diffusion-model quantization: PTQD aims to systematically analyze quantization effects in diffusion models and establish a unified framework for accurate post-training quantization.

3 Preliminaries

Diffusion models add Gaussian noise in a forward process and learn a reverse denoising process. The paper studies pretrained models without retraining, using inherited schedules and uniform quantization to represent tensors at lower bitwidths.

  • Diffusion models: The forward process gradually applies Gaussian noise to real data, while the learned reverse process denoises and generates images.
  • Diffusion models: The reverse process approximates each intractable conditional distribution with a learned Gaussian pθ(x_t−1|x_t) parameterized by a predicted mean and variance.
  • Diffusion models: The reverse-process variance may be reparameterized or fixed to a constant schedule σ_t used during sampling.
  • Scope: PTQD uses pretrained diffusion models without training, inherits their hyperparameters and variance schedules, and can extend its DDPM derivations to DDIM.
  • Model quantization: Uniform quantization maps a floating-point vector to a target bitwidth b using a round operation and a range-derived quantization interval.
  • Model quantization: The notation defines X as a full-precision tensor, ˆX as its quantized counterpart, and ∆X as the difference between them.

4 Method

PTQD corrects quantization-induced mean and variance deviations by separating correlated and uncorrelated noise, then uses step-aware precision selection to maintain denoising quality across sampling steps.

  • Quantization Noise in Reverse Sampling: Quantized reverse sampling introduces quantization noise that alters the mean and variance of x_t−1, reducing SNR and sample quality.The method therefore targets mean and variance correction at each reverse-process step.
  • 4.1 Correlation Disentanglement: Normalization layers can create correlation between quantization noise and the full-precision noise-prediction output, confirmed empirically in Figure 2.The correlated component is modeled as kϵ_θ(x_t,t), while the residual is assumed uncorrelated with ϵ_θ(x_t,t).
  • 4.1 Correlation Disentanglement: PTQD estimates the correlation coefficient k and corrects the correlated component by dividing the quantized noise-prediction output by 1+k.The remaining uncorrelated noise has reduced mean and variance when k ≥ 0.
  • 4.2.2 Uncorrelated Noise Correction: Bias Correction subtracts the channel-wise mean of uncorrelated noise, while Variance Schedule Calibration absorbs its additional variance into the Gaussian diffusion noise.The uncorrelated noise is modeled as Gaussian, with statistics estimated from quantized and full-precision samples; the calibration fails analytically for deterministic sampling when σ_t is zero.
  • 4.3 Step-aware Mixed Precision: Even after mean and variance correction, low-bit diffusion models remain challenged by low SNR in the quantized noise-prediction network.This limitation is especially relevant for later denoising steps, where quantization noise can become comparable to the original signal.
  • 4.3 Step-aware Mixed Precision: SNR_Q decreases as denoising proceeds toward smaller t and improves with higher bitwidths and the proposed corrections, motivating step-aware mixed precision.The scheme varies activation bitwidth by step while keeping weight bitwidth fixed and selects the minimum bitwidth whose SNR_Q remains within a reasonable range.

5 Experiments

Experiments evaluate PTQD on class-conditional ImageNet, unconditional LSUN, and deployment-efficiency settings, showing improved quality and reduced computation across quantization configurations.

  • Implementation details: 1024 samples from quantized and full-precision models provide the correlation, mean, and variance statistics used for calibration.These statistics are collected at each denoising step before applying PTQD.
  • Ablation study: 0.48 lower FID and 6.55 lower sFID result from Correlated Noise Correction on ImageNet 256 × 256 LDM-4.The ablation also evaluates Bias Correction and Variance Schedule Calibration for residual uncorrelated noise.
  • Class-conditional generation: 0.06 FID gap to full precision is achieved with W4A8 at 250 ImageNet generation steps, alongside 6.83× model compression and 19.96× fewer bit operations.At 20 steps, PTQD reaches FID 7.75 despite variance schedule calibration being unavailable when eta is zero.
  • Unconditional generation: 0.1 FID increase is reported for 8-bit LSUN-Churches generation, while W4A8 reduces LSUN-Bedrooms FID and sFID by 0.78 and 3.61 versus Q-Diffusion.On LSUN-Churches mixed precision, FID decreases from 218.59 to 17.99 using Correlated Noise Correction and Bias Correction.
  • Deployment efficiency: 2.03× and 3.34× latency reductions are measured for W8A8 and W4A4 matrix multiplication and convolution over full precision on an RTX3090.The mixed-precision deployment experiments use W8A8 and W4A4 kernels because a fast W4A8 implementation is unavailable.

6 Conclusion and Future Work

PTQD unifies quantization and diffusion noise modeling, corrects correlated and uncorrelated components, and adaptively selects denoising-step precision; the paper also notes misuse risks and future extensions.

  • Conclusion: PTQD unifies the formulation of quantization noise and diffusion perturbed noise in post-training quantized diffusion models.The framework targets accurate quantization without retraining.
  • Conclusion: Correlated Noise Correction estimates a correlation coefficient, while Bias Correction and Variance Schedule Calibration address mean deviations and excess variance.Step-aware mixed precision additionally selects the bitwidth for each denoising step.
  • Conclusion: 0.06 FID increase versus full-precision LDM-4 on ImageNet 256 × 256 accompanies 19.9× fewer bit operations.The result is reported as the paper’s overall performance and efficiency outcome.
  • Future work: Future work includes quantizing the text encoder and image decoder and extending PTQD to a wider range of generative tasks.The latter is proposed to assess efficacy and generalizability.
  • Limitations and Broader Impacts: PTQD may be used to produce counterfeit images and videos for malicious purposes, despite its efficiency and energy-saving properties.This broader-impact concern is stated for PTQD alongside other deep generative models.

B Statistical analysis

Statistical analyses support modeling residual quantization noise as Gaussian, show bitwidth-dependent variance and channel-wise bias, and find strong correlation in low-bit settings.

  • Distribution of residual quantization noise: At significance level 0.01, normality tests cannot reject a Gaussian distribution for residual quantization noise at any step.The assumption is based on D’Agostino–Pearson testing.
  • Distribution of residual quantization noise: Lower bitwidth increases residual quantization-noise variance, but its relatively small coefficient permits absorption into the calibrated diffusion variance schedule.This supports Variance Schedule Calibration for the residual uncorrelated component.
  • Distribution of residual quantization noise: Channel-wise bias varies substantially, motivating distinct correction procedures for different channels.The bias affects the estimated mean introduced by residual quantization noise.
  • Correlation analysis: High Pearson R values for W4A4 indicate that quantization noise primarily consists of a component correlated with full-precision noise predictions.Correlation-based correction also addresses a substantial portion of W4A8 and W8A8 quantization noise.

C.1 Implementation details of step-aware mixed precision

The step-aware mixed-precision analysis allocates activation bitwidths by denoising step using SNR statistics, while component ablations and PTQ4DM comparisons show stronger gains at lower precision.

  • Implementation details of step-aware mixed precision: Activation bitwidths are selected from a predefined set using step-wise SNRQ evaluation, while weight bitwidth remains fixed across denoising steps.SNRF from the forward process provides a benchmark for data noise at each step.
  • Implementation details of step-aware mixed precision: SNRQ correlation coefficients are analyzed for LDM-8 on LSUN-Churches and LDM-4 on LSUN-Bedrooms.The supplied figure references identify per-step coefficient analyses for these model-dataset pairs.
  • Implementation details of step-aware mixed precision: Lower-bit W3A8 settings produce more significant improvements because they contain greater quantization noise.Constant-precision ablations show consistent gains from each proposed component.
  • Comparisons with PTQ4DM: PTQD outperforms PTQ4DM on LSUN-Bedrooms under both W4A8 and W3A8 bitwidth settings.The comparison is reported in Table G over the LDM-4 model.

C.4 Evaluation with advanced sampler

PTQD remains effective with advanced samplers, reducing distributional distances on CelebA-HQ and maintaining comparable quality at lower precision on LSUN-Churches. Additional ablations indicate that its advantage grows as bitwidth decreases.

  • CelebA-HQ with PLMS: 3.23 and 4.73 lower FID and sFID, respectively, than Q-Diffusion on CelebA-HQ with PLMS.Table H evaluates recent DDPM variants with the PLMS sampler.
  • LSUN-Churches with DPM++: PTQD with W3A8 achieves sFID comparable to W4A8 Q-Diffusion on LSUN-Churches with DPM++.The comparison uses the latest DPM++ solver.
  • Bitwidth allocation: Table E reports the bitwidth allocation used for each dataset.
  • Constant-precision ablation: As bitwidth decreases, the efficacy of PTQD becomes increasingly pronounced in the constant-precision LSUN-Bedrooms ablation.

C.5 Evaluation with different variance schedule

The paper evaluates sample quality across sampling configurations and variance schedules, including deterministic and stochastic sampling and comparisons with PTQ4DM and other sampler settings.

  • Deterministic and stochastic sampling: Stochastic perturbations lower FID and sFID for deterministic DDIM sampling on FFHQ, while greatly improving ImageNet IS with little FID and sFID increase.These results are reported for LDM-4 experiments.
  • Comparison with PTQ4DM: Table G compares PTQD with PTQ4DM on LSUN-Bedrooms using the LDM-4 model.
  • PLMS sampling: Table H reports experimental results on CelebA-HQ with the PLMS sampler.
  • DPM++ sampling: Table I reports experimental results on LSUN-Churches with the DPM++ sampler.
  • Variance schedules: Table J compares generated sample quality under different variance schedules denoted by eta.

D Additional visualization results

Additional visualizations show PTQD producing high-fidelity or less quantization-affected samples across step counts, datasets, and low-bit mixed-precision settings. The comparisons emphasize closer agreement with full-precision generation and robustness where Q-Diffusion fails.

  • Sampling steps: The visualizations include samples generated with 20 steps and 250 steps.
  • ImageNet visualization: 20-step PTQD generation with LDM-4 at W4A8 produces high-fidelity class-conditional ImageNet 256 × 256 images.
  • LSUN-Bedrooms visualization: On LSUN-Bedrooms 256 × 256, PTQD samples are less affected by quantization noise and more closely resemble full-precision LDM-4 than Q-Diffusion samples.
  • Mixed-precision visualization: PTQD generates high-quality LSUN-Churches images with W4A4 in mixed precision, whereas Q-Diffusion fails to denoise under the same setting.The comparison uses full-precision LDM-8 as an additional reference.
Loading 2305.10657v4…